Use MemorySegmentReaderSupplier in ReaderSupplierFactory

Author: mdogan

jvector-examples/pom.xml

@@ -214,7 +214,7 @@
             </build>
         </profile>
         <profile>
-            <id>jdk21</id>
+            <id>jdk22</id>
             <activation>
                 <activeByDefault>true</activeByDefault>
             </activation>

jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/ReaderSupplierFactory.java

@@ -19,15 +19,26 @@
 import io.github.jbellis.jvector.disk.SimpleMappedReaderSupplier;
 
 import java.io.IOException;
+import java.lang.reflect.Constructor;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
 public class ReaderSupplierFactory {
     private static final Logger LOG = Logger.getLogger(ReaderSupplierFactory.class.getName());
+    private static final String MEMORY_SEGMENT_READER_CLASSNAME = "io.github.jbellis.jvector.disk.MemorySegmentReaderSupplier";
 
     public static ReaderSupplier open(Path path) throws IOException {
+        try {
+            var supplierClass = Class.forName(MEMORY_SEGMENT_READER_CLASSNAME);
+            Constructor<?> ctor = supplierClass.getConstructor(Path.class);
+            return (ReaderSupplier) ctor.newInstance(path);
+        } catch (Exception e) {
+            LOG.log(Level.WARNING, "MemorySegmentReaderSupplier not available, falling back to MMapReaderSupplier. Reason: {0}: {1}",
+                    new Object[]{e.getClass().getName(), e.getMessage()});
+        }
+
         try {
             return new MMapReaderSupplier(path);
         } catch (UnsatisfiedLinkError|NoClassDefFoundError e) {

jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java

@@ -40,8 +40,8 @@ public class MemorySegmentReader implements RandomAccessReader {
     private static final OfFloat floatLayout = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN);
     private static final OfLong longLayout = ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN);
 
-    private final Arena arena;
-    private final MemorySegment memory;
+    final Arena arena;
+    final MemorySegment memory;
     private long position = 0;
 
     public MemorySegmentReader(Path path) throws IOException {
@@ -54,7 +54,7 @@ public MemorySegmentReader(Path path) throws IOException {
         }
     }
 
-    private MemorySegmentReader(Arena arena, MemorySegment memory) {
+    MemorySegmentReader(Arena arena, MemorySegment memory) {
         this.arena = arena;
         this.memory = memory;
     }
@@ -134,6 +134,11 @@ public void close() {
         arena.close();
     }
 
+    /**
+     * Creates a shallow copy of the MemorySegmentReader.
+     * Underlying memory-mapped region will be shared between all copies.
+     * When MemorySegmentReader is closed, all copies will become invalid.
+     */
     public MemorySegmentReader duplicate() {
         return new MemorySegmentReader(arena, memory);
     }

jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReaderSupplier.java

@@ -16,13 +16,15 @@
 package io.github.jbellis.jvector.disk;
 
 import java.io.IOException;
+import java.lang.foreign.Arena;
+import java.lang.foreign.MemorySegment;
 import java.nio.file.Path;
 
 public class MemorySegmentReaderSupplier implements ReaderSupplier {
-    private final MemorySegmentReader reader;
+    private final InternalMemorySegmentReader reader;
 
     public MemorySegmentReaderSupplier(Path path) throws IOException {
-        reader = new MemorySegmentReader(path);
+        reader = new InternalMemorySegmentReader(path);
     }
 
     @Override
@@ -34,4 +36,31 @@ public RandomAccessReader get() {
     public void close() {
         reader.close();
     }
+
+    private static class InternalMemorySegmentReader extends MemorySegmentReader {
+
+        private final boolean shouldClose;
+
+        private InternalMemorySegmentReader(Path path) throws IOException {
+            super(path);
+            shouldClose = true;
+        }
+
+        private InternalMemorySegmentReader(Arena arena, MemorySegment memory) {
+            super(arena, memory);
+            shouldClose = false;
+        }
+
+        @Override
+        public void close() {
+            if (shouldClose) {
+                super.close();
+            }
+        }
+
+        @Override
+        public InternalMemorySegmentReader duplicate() {
+            return new InternalMemorySegmentReader(arena, memory);
+        }
+    }
 }

jvector-native/src/test/java/io/github/jbellis/jvector/disk/MemorySegmentReaderTest.java

@@ -100,6 +100,32 @@ public void testReaderClose() throws Exception {
         }
     }
 
+    @Test
+    public void testSupplierClose() throws Exception {
+        var s = new MemorySegmentReaderSupplier(tempFile);
+        var r1 = s.get();
+        var r2 = s.get();
+
+        // Close on supplied readers are nop.
+        r1.close();
+        r1.readInt();
+        r2.close();
+        r2.readInt();
+
+        // Backing memory-map will be closed when supplier is closed.
+        s.close();
+        try {
+            r1.readInt();
+            fail("Should have thrown an exception");
+        } catch (IllegalStateException _) {
+        }
+        try {
+            r2.readInt();
+            fail("Should have thrown an exception");
+        } catch (IllegalStateException _) {
+        }
+    }
+
     private void verifyReader(MemorySegmentReader r) {
         r.seek(0);
         var bytes = new byte[7];