From 264979e677aeee0ba8ead7164672099cd3108da7 Mon Sep 17 00:00:00 2001 From: ikaitist Date: Tue, 25 Aug 2026 00:58:37 +0000 Subject: [PATCH] Make threadLocalSupplier() return a closeable supplier and close it at build call sites --- .../jvector/graph/GraphIndexBuilder.java | 40 ++++-- .../graph/RandomAccessVectorValues.java | 31 ++++- .../jvector/graph/TestThreadLocalCopies.java | 131 ++++++++++++++++++ 3 files changed, 189 insertions(+), 13 deletions(-) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestThreadLocalCopies.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java index 4139a14b6..222191512 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java +++ b/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 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(); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java index eb8f6df24..532692b6a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/RandomAccessVectorValues.java @@ -95,6 +95,12 @@ default void getVectorInto(int node, VectorFloat destinationVector, int offse /** * Returns a supplier of thread-local copies of the RAVV. + *

+ * 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 threadLocalSupplier() { if (!isValueShared()) { @@ -102,10 +108,31 @@ default Supplier threadLocalSupplier() { } 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, AutoCloseable { + private final ExplicitThreadLocal tl; + + ThreadLocalCopies(ExplicitThreadLocal tl) { + this.tl = tl; + } + + @Override + public RandomAccessVectorValues get() { + return tl.get(); + } + + @Override + public void close() throws Exception { + tl.close(); + } } /** diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestThreadLocalCopies.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestThreadLocalCopies.java new file mode 100644 index 000000000..f8082dc0d --- /dev/null +++ b/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()); + } +}