prefetch may select the wrong memory segment for multi-segment slices

apache · Jan 7, 2025 · 7ba2afe · 7ba2afe
1 parent d92d045
commit 7ba2afe
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 4 deletions.
diff --git a/lucene/core/src/java21/org/apache/lucene/store/MemorySegmentIndexInput.java b/lucene/core/src/java21/org/apache/lucene/store/MemorySegmentIndexInput.java
@@ -337,8 +337,6 @@ public void prefetch(long offset, long length) throws IOException {
 
     ensureOpen();
 
-    Objects.checkFromIndexSize(offset, length, length());
-
     if (BitUtil.isZeroOrPowerOfTwo(consecutivePrefetchHitCount++) == false) {
       // We've had enough consecutive hits on the page cache that this number is neither zero nor a
       // power of two. There is a good chance that a good chunk of this index input is cached in
@@ -381,8 +379,6 @@ void advise(long offset, long length, IOConsumer<MemorySegment> advice) throws I
 
     ensureOpen();
 
-    Objects.checkFromIndexSize(offset, length, length());
-
     final NativeAccess nativeAccess = NATIVE_ACCESS.get();
 
     try {
@@ -818,6 +814,12 @@ public MemorySegment segmentSliceOrNull(long pos, long len) throws IOException {
         throw handlePositionalIOOBE(e, "segmentSliceOrNull", pos);
       }
     }
+
+    @Override
+    public void prefetch(long offset, long length) throws IOException {
+      Objects.checkFromIndexSize(offset, length, this.length);
+      super.prefetch(offset, length);
+    }
   }
 
   /** This class adds offset support to MemorySegmentIndexInput, which is needed for slices. */
@@ -903,5 +905,11 @@ public MemorySegment segmentSliceOrNull(long pos, long len) throws IOException {
     MemorySegmentIndexInput buildSlice(String sliceDescription, long ofs, long length) {
       return super.buildSlice(sliceDescription, this.offset + ofs, length);
     }
+
+    @Override
+    public void prefetch(long offset, long length) throws IOException {
+      Objects.checkFromIndexSize(offset, length, this.length);
+      super.prefetch(this.offset + offset, length);
+    }
   }
 }
diff --git a/lucene/core/src/test/org/apache/lucene/store/TestMMapDirectory.java b/lucene/core/src/test/org/apache/lucene/store/TestMMapDirectory.java
@@ -329,4 +329,42 @@ public void testNoGroupingFunc() {
     assertFalse(func.apply("segment.si").isPresent());
     assertFalse(func.apply("_51a.si").isPresent());
   }
+
+  public void testPrefetchWithSingleSegment() throws IOException {
+    testPrefetchWithSegments(64 * 1024);
+  }
+
+  public void testPrefetchWithMultiSegment() throws IOException {
+    testPrefetchWithSegments(16 * 1024);
+  }
+
+  static final Class<IndexOutOfBoundsException> IOOBE = IndexOutOfBoundsException.class;
+
+  // does not verify that the actual segment is prefetched, but rather exercises the code and bounds
+  void testPrefetchWithSegments(int maxChunkSize) throws IOException {
+    byte[] bytes = new byte[(maxChunkSize * 2) + 1];
+    try (Directory dir =
+        new MMapDirectory(createTempDir("testPrefetchWithSegments"), maxChunkSize)) {
+      try (IndexOutput out = dir.createOutput("test", IOContext.DEFAULT)) {
+        out.writeBytes(bytes, 0, bytes.length);
+      }
+
+      try (var in = dir.openInput("test", IOContext.READONCE)) {
+        in.prefetch(0, in.length());
+        expectThrows(IOOBE, () -> in.prefetch(1, in.length()));
+        expectThrows(IOOBE, () -> in.prefetch(in.length(), 1));
+
+        var slice1 = in.slice("slice-1", 1, in.length() - 1);
+        slice1.prefetch(0, slice1.length());
+        expectThrows(IOOBE, () -> slice1.prefetch(1, slice1.length()));
+        expectThrows(IOOBE, () -> slice1.prefetch(slice1.length(), 1));
+
+        // we sliced off all but one byte from the first complete memory segment
+        var slice2 = in.slice("slice-2", maxChunkSize - 1, in.length() - maxChunkSize + 1);
+        slice2.prefetch(0, slice2.length());
+        expectThrows(IOOBE, () -> slice2.prefetch(1, slice2.length()));
+        expectThrows(IOOBE, () -> slice2.prefetch(slice2.length(), 1));
+      }
+    }
+  }
 }