Merge remote-tracking branch 'upstream/main' into success-exception-handling

Author: thecoop

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java

@@ -45,6 +45,7 @@
 import org.apache.lucene.store.FileTypeHint;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.store.PreloadHint;
 import org.apache.lucene.store.RandomAccessInput;
 import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
@@ -109,7 +110,10 @@ public Lucene99HnswVectorsReader(SegmentReadState state, FlatVectorsReader flatV
               Lucene99HnswVectorsFormat.VECTOR_INDEX_EXTENSION,
               Lucene99HnswVectorsFormat.VECTOR_INDEX_CODEC_NAME,
               state.context.withHints(
-                  FileTypeHint.DATA, FileDataHint.KNN_VECTORS, DataAccessHint.RANDOM));
+                  FileTypeHint.DATA,
+                  FileDataHint.KNN_VECTORS,
+                  DataAccessHint.RANDOM,
+                  PreloadHint.INSTANCE));
     } catch (Throwable t) {
       IOUtils.closeWhileSuppressingExceptions(t, this);
       throw t;

lucene/core/src/java/org/apache/lucene/store/MMapDirectory.java

@@ -99,6 +99,13 @@ public class MMapDirectory extends FSDirectory {
    */
   public static final BiPredicate<String, IOContext> NO_FILES = (_, _) -> false;
 
+  /**
+   * Argument for {@link #setPreload(BiPredicate)} that configures files to be preloaded when they
+   * are hinted to do so.
+   */
+  public static final BiPredicate<String, IOContext> PRELOAD_HINT =
+      (_, c) -> c.hints().contains(PreloadHint.INSTANCE);
+
   /**
    * This sysprop allows to control the total maximum number of mmapped files that can be associated
    * with a single shared {@link java.lang.foreign.Arena foreign Arena}. For example, to set the max
@@ -233,6 +240,7 @@ public MMapDirectory(Path path, LockFactory lockFactory, long maxChunkSize) thro
    * @param preload a {@link BiPredicate} whose first argument is the file name, and second argument
    *     is the {@link IOContext} used to open the file
    * @see #ALL_FILES
+   * @see #PRELOAD_HINT
    * @see #NO_FILES
    */
   public void setPreload(BiPredicate<String, IOContext> preload) {

lucene/core/src/java/org/apache/lucene/store/PreloadHint.java

@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.store;
+
+/** A hint that the file should be preloaded into memory */
+public enum PreloadHint implements IOContext.FileOpenHint {
+  INSTANCE
+}

lucene/core/src/test/org/apache/lucene/store/TestMMapDirectory.java

@@ -131,6 +131,31 @@ public void testWithNormal() throws Exception {
     }
   }
 
+  public void testPreload() throws Exception {
+    assumeTrue("madvise for preloading only works on linux/mac", MMapDirectory.supportsMadvise());
+
+    final int size = 8 * 1024;
+    byte[] bytes = new byte[size];
+    random().nextBytes(bytes);
+
+    try (MMapDirectory dir = new MMapDirectory(createTempDir("testPreload"))) {
+      dir.setPreload(MMapDirectory.PRELOAD_HINT);
+      try (IndexOutput out =
+          dir.createOutput("test", IOContext.DEFAULT.withHints(PreloadHint.INSTANCE))) {
+        out.writeBytes(bytes, 0, bytes.length);
+      }
+
+      try (final IndexInput in =
+          dir.openInput("test", IOContext.DEFAULT.withHints(PreloadHint.INSTANCE))) {
+        // the data should be loaded in memory
+        assertTrue(in.isLoaded().orElse(false));
+        final byte[] readBytes = new byte[size];
+        in.readBytes(readBytes, 0, readBytes.length);
+        assertArrayEquals(bytes, readBytes);
+      }
+    }
+  }
+
   // Opens the input with ReadAdvice.READONCE to ensure slice and clone are appropriately confined
   public void testConfined() throws Exception {
     final int size = 16;