Make threadLocalSupplier() return a closeable supplier and close it at build call sites

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java

@@ -35,6 +35,7 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.util.*;
+import java.util.function.Supplier;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.ForkJoinPool;
@@ -435,17 +436,31 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi
 
     public ImmutableGraphIndex build(RandomAccessVectorValues ravv) {
         var vv = ravv.threadLocalSupplier();
-        int size = ravv.size();
+        try {
+            int size = ravv.size();
 
-        simdExecutor.submit(() -> {
-            IntStream.range(0, size).parallel().forEach(node -> {
-                addGraphNode(node, vv.get().getVector(node));
-            });
-        }).join();
+            simdExecutor.submit(() -> {
+                IntStream.range(0, size).parallel().forEach(node -> {
+                    addGraphNode(node, vv.get().getVector(node));
+                });
+            }).join();
+        } finally {
+            closeThreadLocalSupplier(vv);
+        }
 
         cleanup();
         return graph;
     }
+
+    private static void closeThreadLocalSupplier(Supplier<RandomAccessVectorValues> supplier) {
+        if (supplier instanceof AutoCloseable) {
+            try {
+                ((AutoCloseable) supplier).close();
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        }
+    }
     /**
      * Validates that the current entry node has been completely added.
      */
@@ -1069,11 +1084,14 @@ public static ImmutableGraphIndex buildAndMergeNewNodes(RandomAccessReader in,
             );
 
             var vv = newVectors.threadLocalSupplier();
-
-            // parallel graph construction from the merge documents Ids
-            simdExecutor.submit(() -> IntStream.range(startingNodeOffset, newVectors.size()).parallel().forEach(ord -> {
-                builder.addGraphNode(ord, vv.get().getVector(ord));
-            })).join();
+            try {
+                // parallel graph construction from the merge documents Ids
+                simdExecutor.submit(() -> IntStream.range(startingNodeOffset, newVectors.size()).parallel().forEach(ord -> {
+                    builder.addGraphNode(ord, vv.get().getVector(ord));
+                })).join();
+            } finally {
+                closeThreadLocalSupplier(vv);
+            }
 
             builder.cleanup();
             return builder.getGraph();

jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java

@@ -95,17 +95,44 @@ default void getVectorInto(int node, VectorFloat<?> destinationVector, int offse
 
     /**
      * Returns a supplier of thread-local copies of the RAVV.
+     * <p>
+     * For shared RAVVs the returned supplier is {@link AutoCloseable}: closing it invokes
+     * close() on every AutoCloseable copy created so far and drops the per-thread cache.
+     * Callers that hold the supplier across a bounded operation should close it when done;
+     * heap-only copies are additionally collected once the supplier itself becomes
+     * unreachable.
      */
     default Supplier<RandomAccessVectorValues> threadLocalSupplier() {
         if (!isValueShared()) {
             return () -> this;
         }
 
         if (this instanceof AutoCloseable) {
-            LOG.warning("RAVV is shared and implements AutoCloseable; threadLocalSupplier() may lead to leaks");
+            LOG.warning("RAVV is shared and implements AutoCloseable; close the supplier returned by threadLocalSupplier() to release per-thread copies");
         }
         var tl = ExplicitThreadLocal.withInitial(this::copy);
-        return tl::get;
+        return new ThreadLocalCopies(tl);
+    }
+
+    /**
+     * Thread-local RAVV supplier whose close() releases the per-thread copies.
+     */
+    final class ThreadLocalCopies implements Supplier<RandomAccessVectorValues>, AutoCloseable {
+        private final ExplicitThreadLocal<RandomAccessVectorValues> tl;
+
+        ThreadLocalCopies(ExplicitThreadLocal<RandomAccessVectorValues> tl) {
+            this.tl = tl;
+        }
+
+        @Override
+        public RandomAccessVectorValues get() {
+            return tl.get();
+        }
+
+        @Override
+        public void close() throws Exception {
+            tl.close();
+        }
     }
 
     /**

jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestThreadLocalCopies.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 io.github.jbellis.jvector.graph;
+
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+import org.junit.Test;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
+
+public class TestThreadLocalCopies {
+    private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+    /** Shared RAVV whose copies count their close() calls; the source itself never counts. */
+    static class CloseTrackingRavv implements RandomAccessVectorValues, AutoCloseable {
+        private final AtomicInteger closedCount;
+        private final boolean source;
+
+        CloseTrackingRavv(AtomicInteger closedCount) {
+            this(closedCount, true);
+        }
+
+        CloseTrackingRavv(AtomicInteger closedCount, boolean source) {
+            this.closedCount = closedCount;
+            this.source = source;
+        }
+
+        @Override
+        public int size() {
+            return 1;
+        }
+
+        @Override
+        public int dimension() {
+            return 1;
+        }
+
+        @Override
+        public VectorFloat<?> getVector(int nodeId) {
+            return vts.createFloatVector(1);
+        }
+
+        @Override
+        public boolean isValueShared() {
+            return true;
+        }
+
+        @Override
+        public RandomAccessVectorValues copy() {
+            return new CloseTrackingRavv(closedCount, false);
+        }
+
+        @Override
+        public void close() {
+            if (!source) {
+                closedCount.incrementAndGet();
+            }
+        }
+    }
+
+    @Test
+    public void closingSupplierClosesAllThreadLocalCopies() throws Exception {
+        AtomicInteger closed = new AtomicInteger();
+        CloseTrackingRavv source = new CloseTrackingRavv(closed);
+        var supplier = source.threadLocalSupplier();
+
+        ExecutorService pool = Executors.newFixedThreadPool(2);
+        try {
+            for (int i = 0; i < 2; i++) {
+                pool.submit(() -> supplier.get()).get();
+            }
+        }
+        finally {
+            pool.shutdown();
+            pool.awaitTermination(10, TimeUnit.SECONDS);
+        }
+
+        ((AutoCloseable) supplier).close();
+        assertEquals(2, closed.get());
+    }
+
+    @Test
+    public void supplierRemainsUsableAfterClose() throws Exception {
+        AtomicInteger closed = new AtomicInteger();
+        CloseTrackingRavv source = new CloseTrackingRavv(closed);
+        var supplier = source.threadLocalSupplier();
+
+        supplier.get();
+        ((AutoCloseable) supplier).close();
+        assertEquals(1, closed.get());
+
+        // The per-thread cache is dropped on close: a fresh copy is created and closed again.
+        supplier.get();
+        ((AutoCloseable) supplier).close();
+        assertEquals(2, closed.get());
+    }
+
+    @Test
+    public void unsharedRavvSupplierReturnsTheSource() {
+        AtomicInteger closed = new AtomicInteger();
+        RandomAccessVectorValues unshared = new CloseTrackingRavv(closed) {
+            @Override
+            public boolean isValueShared() {
+                return false;
+            }
+        };
+        var supplier = unshared.threadLocalSupplier();
+        assertSame(unshared, supplier.get());
+    }
+}