From 3c814fac81cc1c09328c3bef8eed7848fd6d2e35 Mon Sep 17 00:00:00 2001 From: "Dian-Lun (Aaron) Lin" Date: Thu, 2 Jul 2026 22:14:58 +0000 Subject: [PATCH 01/13] Compaction: fix cold compaction time by eliminating disk scan in computeLayerInfoFromSources MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit computeLayerInfoFromSources called getNodes(0) on each source graph to count live nodes at level 0. getNodes(0) sequentially seeks through every node record on disk to filter out deleted entries. On a cold page cache this touches large amounts of source data before compaction even begins, significantly delaying the start of actual graph merging. Since every live node is present at level 0 by the HNSW invariant, the count is simply liveNodes.get(s).cardinality() — an in-memory popcount requiring no I/O. Also switch PQ retraining from ProductQuantization.compute() (full k-means++ init) to basePQ.refine() (Lloyd's iterations only, warm-started from the existing codebook). The source codebooks are already trained on the same distribution, so warm-starting converges in far fewer passes with no recall loss. --- .../graph/disk/OnDiskGraphIndexCompactor.java | 17 +++++++--- .../jvector/graph/disk/PQRetrainer.java | 33 +++++++++++-------- .../quantization/ProductQuantization.java | 2 +- 3 files changed, 33 insertions(+), 19 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index deffb719f..ddf5d3cc5 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -1354,11 +1354,18 @@ private List computeLayerInfoFromSources() { int count = 0; for (int s = 0; s < sources.size(); s++) { if (level > sources.get(s).getMaxLevel()) continue; - NodesIterator it = sources.get(s).getNodes(level); - FixedBitSet alive = liveNodes.get(s); - while (it.hasNext()) { - int node = it.next(); - if (alive.get(node)) count++; + if (level == 0) { + // Every live node is present at level 0 (HNSW base layer invariant), + // so count directly from the in-memory bitset instead of scanning node + // records on disk (which touches gigabytes of source data on a cold cache). + count += liveNodes.get(s).cardinality(); + } else { + NodesIterator it = sources.get(s).getNodes(level); + FixedBitSet alive = liveNodes.get(s); + while (it.hasNext()) { + int node = it.next(); + if (alive.get(node)) count++; + } } } layerInfo.add(new CommonHeader.LayerInfo(count, maxDegrees.get(level))); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index 280d75c9c..c45120f70 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -22,6 +22,7 @@ import io.github.jbellis.jvector.quantization.ProductQuantization; import io.github.jbellis.jvector.util.DocIdSetIterator; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -32,6 +33,7 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.List; +import java.util.concurrent.ForkJoinPool; import java.util.concurrent.ThreadLocalRandom; /** @@ -96,22 +98,27 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, log.info("Collected {} training samples", samples.size()); - // Extract vectors sequentially in sorted (source, node) order so disk reads are - // purely sequential and the OS read-ahead can cover them efficiently. We do this - // here rather than letting ProductQuantization.compute() drive the reads via its - // parallel stream, which would scatter page faults across a potentially very large - // file and cause I/O that scales with dataset size rather than sample count. List> trainingVectors = extractVectorsSequential(samples); - var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - boolean center = similarityFunction == VectorSimilarityFunction.EUCLIDEAN; + long t0 = System.nanoTime(); + log.info("Extracted {} vectors in {}ms; starting PQ refinement", + trainingVectors.size(), (System.nanoTime() - t0) / 1_000_000L); + + var ravv = new ListRandomAccessVectorValues(trainingVectors, dimension); - return ProductQuantization.compute( - ravv, - basePQ.getSubspaceCount(), - basePQ.getClusterCount(), - center - ); + // Warm-start from the existing codebook via Lloyd's-only refinement rather than + // re-running k-means++ from scratch. k-means++ initialization visits every point + // once per centroid (256 passes for k=256), which dominates training time. + // Since the source codebooks are already trained on data from the same underlying + // distribution, this warm-start converges in far fewer passes with no recall loss. + long t1 = System.nanoTime(); + ProductQuantization result = basePQ.refine(ravv, + ProductQuantization.K_MEANS_ITERATIONS, + -1.0f, // UNWEIGHTED / isotropic + PhysicalCoreExecutor.pool(), + ForkJoinPool.commonPool()); + log.info("PQ refinement complete in {}ms", (System.nanoTime() - t1) / 1_000_000L); + return result; } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index c84b7b955..6d9d23879 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -60,7 +60,7 @@ public class ProductQuantization implements VectorCompressor>, A private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); static final int DEFAULT_CLUSTERS = 256; // number of clusters per subspace = one byte's worth - static final int K_MEANS_ITERATIONS = 6; + public static final int K_MEANS_ITERATIONS = 6; public static final int MAX_PQ_TRAINING_SET_SIZE = 128000; final VectorFloat[] codebooks; // array of codebooks, where each codebook is a VectorFloat consisting of k contiguous subvectors each of length M From 785cea4fa0137c0bfa855658cdb7f0eef82acb61 Mon Sep 17 00:00:00 2001 From: "Dian-Lun (Aaron) Lin" Date: Thu, 16 Jul 2026 14:39:39 +0000 Subject: [PATCH 02/13] Compaction: windowed streaming prefetch for bulk-phase reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Source graphs are mapped with MADV_RANDOM (correct for search-time access), which disables kernel readahead and makes compaction's bulk phases fault one page at a time on a cold cache (measured: retrain ~37s, pre-encode ~42s on a disk-cold 10M-node compaction that takes ~3s/~8s warm). Adds ReaderSupplier.prefetch(offset, length) — streams a byte range into the page cache through a separate readahead-enabled descriptor — and OnDiskGraphIndex.prefetchL0Records(minNode, maxNode) on top of it. Each bulk phase warms exactly the records it is about to read, from the worker that will read them: - PQ retrain prefetches its (source, node)-sorted sample ranges before extraction (bounded by training-set size, not file size), - code pre-encode prefetches each chunk's records at task start, - L0 batch processing prefetches each batch's own records at task start (cross-source search reads are data-dependent and stay demand-faulted). Transient cache demand is proportional to the in-flight windows, so there is no up-front whole-file streaming pass and no memory-availability gate to mistune: on a box that cannot hold the sources, pages are simply evicted and reads degrade per-page to the old fault-on-demand behavior. --- .../jbellis/jvector/disk/ReaderSupplier.java | 11 +++ .../graph/disk/FusedCompactionStrategy.java | 3 + .../jvector/graph/disk/OnDiskGraphIndex.java | 16 +++++ .../graph/disk/OnDiskGraphIndexCompactor.java | 6 ++ .../jvector/graph/disk/PQRetrainer.java | 29 ++++++++ .../jvector/disk/MemorySegmentReader.java | 70 +++++++++++++++---- 6 files changed, 121 insertions(+), 14 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java index 30431d6ca..7ebb3f9b0 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java @@ -30,6 +30,17 @@ public interface ReaderSupplier extends AutoCloseable { */ RandomAccessReader get() throws IOException; + /** + * Streams the byte range {@code [offset, offset + length)} of the underlying storage into + * the page cache, if the implementation supports it. Synchronous and best-effort; a no-op + * by default. Sized for repeated windowed calls — a caller that knows which records a bulk + * phase is about to read (e.g. graph compaction, whose readers otherwise advise + * {@code MADV_RANDOM} and forgo kernel readahead) can warm exactly those, keeping transient + * cache demand proportional to the window rather than the file. + */ + default void prefetch(long offset, long length) { + } + default void close() throws IOException { } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java index f759bc702..b6a44ece3 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java @@ -230,6 +230,9 @@ private void precomputeCodes(CompactWriter writer) throws IOException { final int cStart = chunkStart; final int cEnd = Math.min(chunkStart + chunkSize, upper); tasks.add(() -> { + // Stream this chunk's records into the page cache before the encode loop; + // the mapping's MADV_RANDOM otherwise faults them one page at a time. + source.prefetchL0Records(cStart, cEnd - 1); ByteSequence code = vectorTypeSupport.createByteSequence(cs); VectorFloat vec = vectorTypeSupport.createFloatVector(ctx.dimension); long count = 0; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index 3fb69d967..61d8157fb 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java @@ -443,6 +443,22 @@ public int maxDegree() { return layerInfo.stream().mapToInt(li -> li.degree).max().orElseThrow(); } + /** + * Streams the L0 records of ordinals {@code [minNode, maxNode]} (inclusive) into the page + * cache; see {@link ReaderSupplier#prefetch(long, long)}. An L0 record holds the node's id, + * inline features (vector, fused codes), and adjacency, so warming it covers every read a + * bulk scan makes for that node. Best-effort no-op when unsupported. + */ + public void prefetchL0Records(int minNode, int maxNode) { + if (maxNode < minNode) { + return; + } + long blockBytes = Integer.BYTES + inlineBlockSize + + (long) Integer.BYTES * (layerInfo.get(0).degree + 1); + long start = neighborsOffset + blockBytes * minNode; + readerSupplier.prefetch(start, blockBytes * (maxNode - minNode + 1)); + } + // re-declared to specify type @Override public View getView() { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index ddf5d3cc5..613c664e3 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -997,6 +997,12 @@ private List computeBaseBatch(CompactWriter writer, CompactionParams params) throws IOException { List out = new ArrayList<>(bs.end - bs.start); + if (bs.end > bs.start) { + // Stream this batch's own records into the page cache before processing. Search + // reads into other sources are data-dependent and stay demand-faulted, but each + // node's own record read (adjacency + vector) is fully predictable. + sources.get(bs.sourceIdx).prefetchL0Records(bs.nodes[bs.start], bs.nodes[bs.end - 1]); + } for (int i = bs.start; i < bs.end; i++) { int node = bs.nodes[i]; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index c45120f70..616263951 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -218,6 +218,8 @@ private List sampleBalanced(int totalSamples) { * cover them efficiently. Each source's view is opened once and reused for all its samples. */ private List> extractVectorsSequential(List samples) { + prefetchSampleRanges(samples); + OnDiskGraphIndex.View[] views = new OnDiskGraphIndex.View[sources.size()]; for (int s = 0; s < sources.size(); s++) { views[s] = (OnDiskGraphIndex.View) sources.get(s).getView(); @@ -232,6 +234,33 @@ private List> extractVectorsSequential(List samples) { return vectors; } + // Merging samples closer than this many ordinals into one range keeps the prefetch mostly + // sequential without dragging in long runs of unsampled records. + private static final int PREFETCH_MERGE_GAP = 256; + + /** + * Streams the records of the (source, node)-sorted sample list into the page cache before + * extraction. The mappings advise {@code MADV_RANDOM}, so without this the extraction loop + * faults one page at a time on a cold cache; total demand is bounded by the training-set + * size, not the file size. + */ + private void prefetchSampleRanges(List samples) { + int i = 0; + while (i < samples.size()) { + SampleRef first = samples.get(i); + int last = first.node; + int j = i + 1; + while (j < samples.size() + && samples.get(j).source == first.source + && samples.get(j).node - last <= PREFETCH_MERGE_GAP) { + last = samples.get(j).node; + j++; + } + sources.get(first.source).prefetchL0Records(first.node, last); + i = j; + } + } + /** * Reference to a sampled vector from a specific source index. */ diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java b/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java index c9ac39c21..7e3c40499 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/disk/MemorySegmentReader.java @@ -148,25 +148,15 @@ public void close() { public static class Supplier implements ReaderSupplier { private final Arena arena; private final MemorySegment memory; + private final Path path; public Supplier(Path path) throws IOException { + this.path = path; this.arena = Arena.ofShared(); try (var ch = FileChannel.open(path, StandardOpenOption.READ)) { this.memory = ch.map(MapMode.READ_ONLY, 0L, ch.size(), arena); - - // Apply MADV_RANDOM advice - var linker = Linker.nativeLinker(); - var maybeMadvise = linker.defaultLookup().find("posix_madvise"); - if (maybeMadvise.isPresent()) { - var madvise = linker.downcallHandle(maybeMadvise.get(), - FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT)); - int result = (int) madvise.invokeExact(memory, memory.byteSize(), MADV_RANDOM); - if (result != 0) { - throw new IOException("posix_madvise failed with error code: " + result); - } - } else { - logger.warn("posix_madvise not found, MADV_RANDOM advice not applied"); - } + // Search-time access is random; disable kernel readahead by default. + madvise(memory, MADV_RANDOM, true); } catch (Throwable e) { arena.close(); if (e instanceof IOException) { @@ -176,6 +166,58 @@ public Supplier(Path path) throws IOException { } } + private static void madvise(MemorySegment memory, int advice, boolean strict) throws IOException { + var linker = Linker.nativeLinker(); + var maybeMadvise = linker.defaultLookup().find("posix_madvise"); + if (maybeMadvise.isPresent()) { + var madvise = linker.downcallHandle(maybeMadvise.get(), + FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT)); + int result; + try { + result = (int) madvise.invokeExact(memory, memory.byteSize(), advice); + } catch (Throwable t) { + throw new RuntimeException(t); + } + if (result != 0 && strict) { + throw new IOException("posix_madvise failed with error code: " + result); + } + } else { + logger.warn("posix_madvise not found, advice {} not applied", advice); + } + } + + // Windowed prefetch streams ranges through a separate read-ahead-enabled file descriptor: + // the mapping itself advises MADV_RANDOM (disables kernel readahead), and MADV_WILLNEED + // proved to be a no-op hint for large ranges. The page cache is shared per file, so + // subsequent random access through the mapping hits the populated pages. Called + // concurrently from many worker threads; positional reads on a shared channel are + // thread-safe, and the per-thread buffer avoids allocation churn. + private static final ThreadLocal PREFETCH_BUF = + ThreadLocal.withInitial(() -> java.nio.ByteBuffer.wrap(new byte[4 << 20])); + + @Override + public void prefetch(long offset, long length) { + if (length <= 0) { + return; + } + var buf = PREFETCH_BUF.get(); + long end = Math.min(offset + length, memory.byteSize()); + try (var ch = FileChannel.open(path, StandardOpenOption.READ)) { + long pos = Math.max(0, offset); + while (pos < end) { + buf.clear().limit((int) Math.min(buf.capacity(), end - pos)); + int n = ch.read(buf, pos); + if (n < 0) { + break; + } + pos += n; + } + } catch (IOException e) { + logger.warn("ranged prefetch of {} [{}, {}) failed; continuing without warm cache", + path, offset, end, e); + } + } + @Override public MemorySegmentReader get() { return new MemorySegmentReader(memory); From a6ff3c61f70982d4a0871b85f713d88e182f6425 Mon Sep 17 00:00:00 2001 From: "Dian-Lun (Aaron) Lin" Date: Thu, 16 Jul 2026 18:57:14 +0000 Subject: [PATCH 03/13] Compaction: use searchTopK as rerankK for L0 cross-source searches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The L0 cross-source candidate search passed beamWidth (= 2x searchTopK) as rerankK, doubling the approximate-phase beam over what the candidate budget needs. A seeding-vs-beam decomposition study (7 paired disk-cold arms, cohere-10M, median-of-3) showed the narrower beam is where the time goes: S=2: L0 165.2s -> 133.4s (-19%), recall 0.5718 -> 0.5660 S=4: L0 260.3s -> 224.6s (-14%), recall 0.5733 -> 0.5659 Query latency on the merged index is unaffected (0.55ms avg both ways). The wider beam's extra candidates were largely pruned by diversity selection, which keeps at most degree edges per node. The same study found warm-start seeding of these searches (from finished neighbors' merged adjacency, reverse candidates, or upper-layer descent) is net-negative: at matched beam width, seeded searches run 8-17% slower with equal recall — the per-search seeding overhead exceeds the few cheap descent hops it saves. Beam width is the whole lever. --- .../jvector/graph/disk/OnDiskGraphIndexCompactor.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index 613c664e3..bb1a8362d 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -84,7 +84,6 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final ForkJoinPool executor; private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; - /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -1223,8 +1222,11 @@ private int gatherFromOtherSource(int node, int level, int sourceIdx, ); if (level == 0) { + // rerankK = searchTopK, not beamWidth: the wider beam's extra candidates are largely + // pruned by diversity selection, so the doubled approximate-phase cost buys almost + // no recall. SearchResult results = scratch.gs[sourceIdx].search( - ssp, params.searchTopK, params.beamWidth, 0f, 0f, indexAlive + ssp, params.searchTopK, params.searchTopK, 0f, 0f, indexAlive ); for (var r : results.getNodes()) { From d00a7c5b8f170e9f94df456bc40d9f4720737441 Mon Sep 17 00:00:00 2001 From: "Dian-Lun (Aaron) Lin" Date: Thu, 16 Jul 2026 19:36:58 +0000 Subject: [PATCH 04/13] Compaction: make post-compaction refinement optional (default off) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refinement ablations (three datasets, disk-cold, paired same-window runs) show its recall contribution on the merged index is ~0: same-beam arms with and without refinement land within 0.001. Skipping it saves 45-58s, ~20-25% of total compaction time at 10M nodes. What it buys is navigability — query latency on the merged index rises from ~0.55ms to ~0.95ms avg (p99 2.0ms to 3.9ms, cohere-10M) without it. That is a workload tradeoff, not a correctness call: default to compaction throughput, and let latency-sensitive pipelines opt back in with setRefineAfterCompaction(true). --- .../graph/disk/OnDiskGraphIndexCompactor.java | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index bb1a8362d..be8ae8c7e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -84,6 +84,19 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final ForkJoinPool executor; private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; + private boolean refineAfterCompaction = false; + + /** + * Whether to run the second-pass neighbor refinement after the merged graph is written + * (default false). Refinement is a navigability pass: it has no measurable effect on + * recall, but it improves query latency on the merged index at the cost of a significant + * fraction of total compaction time. Enable it when search latency matters more than + * compaction throughput. + */ + @Experimental + public void setRefineAfterCompaction(boolean refineAfterCompaction) { + this.refineAfterCompaction = refineAfterCompaction; + } /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -297,7 +310,9 @@ public void compact(Path outputPath) throws FileNotFoundException { try { compactGraphImpl(outputPath, strategy); releaseSourcesBeforeRefine(strategy); - refineCompactedGraph(outputPath, strategy); + if (refineAfterCompaction) { + refineCompactedGraph(outputPath, strategy); + } } finally { // Delayed until after refinement so refineCompactedGraph can read from the pre-encoded // code cache appended past the projected EOF; onAfterClose unmaps it and truncates. @@ -329,7 +344,9 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce try { sidecarStrategy.retrain(similarityFunction); compactGraphImpl(graphPath, inlineStrategy); - refineCompactedGraph(graphPath, inlineStrategy); + if (refineAfterCompaction) { + refineCompactedGraph(graphPath, inlineStrategy); + } sidecarStrategy.writeSidecar(compressedPath); } catch (IOException e) { throw new RuntimeException("Sidecar compaction failed", e); From a4958dbcd5b7b9a2e51428831f1f890ba9550144 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 1 Jul 2026 08:02:17 +0000 Subject: [PATCH 05/13] support cooperative resource shareing when embedded --- .../graph/disk/OnDiskGraphIndexCompactor.java | 200 +++++++++--- .../jvector/util/work/LeakyBucketLimiter.java | 67 ++++ .../jvector/util/work/ProgressLimiter.java | 97 ++++++ .../jvector/util/work/ProgressTracker.java | 43 +++ .../jvector/util/work/WorkLimiter.java | 59 ++++ .../jbellis/jvector/util/work/WorkStage.java | 33 ++ .../disk/TestOnDiskGraphIndexCompactor.java | 199 ++++++++++++ .../util/work/TestProgressLimiter.java | 289 ++++++++++++++++++ 8 files changed, 951 insertions(+), 36 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index be8ae8c7e..9f1c89d68 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -35,6 +35,9 @@ import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; import io.github.jbellis.jvector.util.*; +import io.github.jbellis.jvector.util.work.ProgressLimiter; +import io.github.jbellis.jvector.util.work.WorkLimiter; +import io.github.jbellis.jvector.util.work.WorkStage; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.graph.similarity.ScoreFunction; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -86,6 +89,51 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final VectorSimilarityFunction similarityFunction; private boolean refineAfterCompaction = false; + // Embedder progress + work-admission control surface (see io.github.jbellis.jvector.util.work). + // Default UNLIMITED = no observation and no throttling → byte-identical output and equivalent + // timing to no SPI installed. + private volatile ProgressLimiter limiter = ProgressLimiter.UNLIMITED; + + /** + * The stages a compaction reports progress and acquires work admission against. Names are + * stable so an embedder can distinguish them. The write side of {@code MERGE_LEVELS} carries the + * graph-body IO — there is no separate flush phase. + */ + @Experimental + public enum Phase implements WorkStage { + /** Merging source graphs level-by-level and writing the compacted graph body. */ + MERGE_LEVELS, + /** Second pass refining neighborhoods in the compacted graph. */ + REFINE + } + + /** + * Installs an embedder control surface for progress observation and work admission. Both facets + * are optional (see {@link ProgressLimiter}); passing {@code null} restores + * {@link ProgressLimiter#UNLIMITED}. Returns {@code this} for chaining. + * + *

Admission ({@code acquire}) is invoked only on the orchestrating thread — the caller of + * {@code compact} — never on a pool worker, so a blocking limiter back-pressures batch dispatch + * without a {@code ForkJoinPool.ManagedBlocker}. The unit of {@code acquire} for this consumer + * is bytes about to be written. + */ + @Experimental + public OnDiskGraphIndexCompactor setProgressLimiter(ProgressLimiter limiter) { + this.limiter = (limiter == null) ? ProgressLimiter.UNLIMITED : limiter; + return this; + } + + /** + * The compactor's task window: the number of batches kept in flight, equal to the injected + * pool's {@code getParallelism()} (or the {@link PhysicalCoreExecutor#pool()} default's). It + * bounds both compaction concurrency and the in-flight memory window; embedders can log or + * assert it to confirm the parallelism they injected. + */ + @Experimental + public int getTaskWindowSize() { + return taskWindowSize; + } + /** * Whether to run the second-pass neighbor refinement after the merged graph is written * (default false). Refinement is a navigability pass: it has no measurable effect on @@ -120,6 +168,13 @@ public OnDiskGraphIndexCompactor( * that ship alongside each graph. Pass {@code null} when sources carry * quantization inline (FUSED_PQ) or have none. Must not be combined * with sources that carry the FUSED_PQ feature. + * @param executor the pool that runs compaction batches. Its {@code getParallelism()} sets + * both CPU concurrency and the in-flight task/memory window + * ({@link #getTaskWindowSize()}), so one knob bounds two dimensions. Compaction + * is compute- and memory-bandwidth-bound; size this from physical cores — + * logical-core sizing oversubscribes hyperthreaded hosts and costs throughput. + * The compactor never owns or shuts down the pool. Pass {@code null} to use the + * shared {@link PhysicalCoreExecutor#pool()} default. */ @Experimental public OnDiskGraphIndexCompactor( @@ -306,12 +361,33 @@ private void validateFeatures(List sources) { */ @Experimental public void compact(Path outputPath) throws FileNotFoundException { + compact(outputPath, 0L); + } + + /** + * No-copy compaction entry point: writes the compacted graph into {@code outputPath} at + * {@code startOffset}, leaving any bytes in {@code [0, startOffset)} untouched. An embedder + * that wraps the graph in its own container reserves its header by passing the header size as + * {@code startOffset}, so jvector's body lands directly inside the container — removing the + * temp-file-and-copy. jvector writes only {@code [startOffset, projectedSize)} and never reads + * or clobbers the reserved prefix. The file is opened read/write and not truncated, so a + * prefix the embedder pre-wrote survives. {@code compact(path, 0)} equals {@link #compact(Path)}. + * + * @param outputPath the file to write into + * @param startOffset the byte offset at which jvector's output begins; {@code 0} for a + * standalone file + */ + @Experimental + public void compact(Path outputPath, long startOffset) throws FileNotFoundException { + if (startOffset < 0) { + throw new IllegalArgumentException("startOffset must be >= 0, got " + startOffset); + } QuantizationCompactionStrategy strategy = detectInlineStrategy(); try { - compactGraphImpl(outputPath, strategy); + compactGraphImpl(outputPath, startOffset, strategy); releaseSourcesBeforeRefine(strategy); if (refineAfterCompaction) { - refineCompactedGraph(outputPath, strategy); + refineCompactedGraph(outputPath, startOffset, strategy); } } finally { // Delayed until after refinement so refineCompactedGraph can read from the pre-encoded @@ -343,9 +419,9 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce QuantizationCompactionStrategy sidecarStrategy = detectSidecarStrategy(); try { sidecarStrategy.retrain(similarityFunction); - compactGraphImpl(graphPath, inlineStrategy); + compactGraphImpl(graphPath, 0L, inlineStrategy); if (refineAfterCompaction) { - refineCompactedGraph(graphPath, inlineStrategy); + refineCompactedGraph(graphPath, 0L, inlineStrategy); } sidecarStrategy.writeSidecar(compressedPath); } catch (IOException e) { @@ -417,7 +493,7 @@ private CompactionContext buildContext() { * mmap cleanup) are delegated to {@code strategy}. For sources with no inline quantization, * pass {@link QuantizationCompactionStrategy#NONE} for a fully no-op strategy hook set. */ - private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy strategy) throws FileNotFoundException { + private void compactGraphImpl(Path outputPath, long startOffset, QuantizationCompactionStrategy strategy) throws FileNotFoundException { strategy.retrain(similarityFunction); boolean fusedPQEnabled = strategy.writesCodesInline(); @@ -433,7 +509,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st log.info("Writing compacted graph : {} total nodes, maxOrdinal={}, dimension={}, degree={}", numTotalNodes, maxOrdinal, dimension, maxDegrees.get(0)); - try (CompactWriter writer = new CompactWriter(outputPath, maxOrdinal, numTotalNodes, 0, layerInfo, entryNode, dimension, maxDegrees, outputFusedFeature)) { + try (CompactWriter writer = new CompactWriter(outputPath, maxOrdinal, numTotalNodes, startOffset, layerInfo, entryNode, dimension, maxDegrees, outputFusedFeature)) { // Header has to be written first so the writer's position is past the header // before any strategy that mmaps past the projected end of the output runs. writer.writeHeader(); @@ -473,7 +549,7 @@ private void compactGraphImpl(Path outputPath, QuantizationCompactionStrategy st * Only L0 records are written. Upper-layer neighbor lists live in an in-memory map after * load and have no addressable file offset, so they're left as written by compactLevels. */ - private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrategy strategy) { + private void refineCompactedGraph(Path outputPath, long startOffset, QuantizationCompactionStrategy strategy) { log.info("Refining compacted graph: {}", outputPath); long t0 = System.nanoTime(); @@ -500,7 +576,7 @@ private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrateg // useFooter=false because the file's logical EOF (where the v6 footer trailer sits) is // before the still-attached pre-encode cache section. loadFromFooter() would seek to // the actual file length and read garbage as the magic. - OnDiskGraphIndex mergedGraph = OnDiskGraphIndex.load(supplier, 0, false); + OnDiskGraphIndex mergedGraph = OnDiskGraphIndex.load(supplier, startOffset, false); // Pick the iteration set: when there's a hierarchy, refine only L1 nodes (each also // lives in L0, so their L0 record is what we rewrite). Mirrors GraphIndexBuilder's @@ -536,33 +612,51 @@ private void refineCompactedGraph(Path outputPath, QuantizationCompactionStrateg log.info("Refining {} live nodes at level {} (hierarchy maxLevel={}, fusedPQ={}, codeCache={})", total, iterationLevel, mergedGraph.getMaxLevel(), fpq, cache != null); - int submitted = 0; - for (int start = 0; start < total; start += batchSize) { - final int s = start; - final int e = Math.min(start + batchSize, total); - ecs.submit(() -> { - RefineScratch scratch = tls.get(); - for (int i = s; i < e; i++) { - int node = ords[i]; - refineOneNode(node, scratch, fc, baseDegree, fpq, codeSize, cmp, bw, - graphRef, cache, cacheSz); - } - return e - s; - }); - submitted++; - } - - int completed = 0; - int nodesDone = 0; + // Windowed submit/drain (mirrors runBatchesWithBackpressure) so a blocking WorkLimiter — + // a rate limiter, or a semaphore that releases permits on Grant.close() — is + // deadlock-free: admission is on the orchestrator before each submit, and exactly one + // grant is released per completed batch. This also bounds in-flight refine batches (and + // thus peak memory) to taskWindowSize, which the prior submit-all loop did not. Amount is + // an estimate of the per-node record bytes rewritten. UNLIMITED makes it all no-ops. + final long refinedRecordSize = (long) dimension * Float.BYTES + + (long) baseDegree * Integer.BYTES + codeSize; + final int nBatches = (total + batchSize - 1) / batchSize; + java.util.ArrayDeque grants = new java.util.ArrayDeque<>(); + + int nextStart = 0, inFlight = 0, completed = 0, nodesDone = 0; int progressStep = Math.max(1, total / 10); int nextProgress = progressStep; - while (completed < submitted) { - nodesDone += ecs.take().get(); - completed++; - if (nodesDone >= nextProgress) { - log.info("Refinement progress: {}/{} nodes", nodesDone, total); - nextProgress += progressStep; + try { + while (completed < nBatches) { + while (inFlight < taskWindowSize && nextStart < total) { + final int s = nextStart; + final int e = Math.min(nextStart + batchSize, total); + nextStart = e; + grants.add(limiter.acquire((long) (e - s) * refinedRecordSize)); // may park the orchestrator + ecs.submit(() -> { + RefineScratch scratch = tls.get(); + for (int i = s; i < e; i++) { + int node = ords[i]; + refineOneNode(node, scratch, fc, baseDegree, fpq, codeSize, cmp, bw, + graphRef, cache, cacheSz); + } + return e - s; + }); + inFlight++; + } + nodesDone += ecs.take().get(); + completed++; + inFlight--; + WorkLimiter.Grant g = grants.poll(); + if (g != null) g.close(); + limiter.onProgress(Phase.REFINE, nodesDone, total); + if (nodesDone >= nextProgress) { + log.info("Refinement progress: {}/{} nodes", nodesDone, total); + nextProgress += progressStep; + } } + } finally { + for (WorkLimiter.Grant g : grants) g.close(); } // Per-thread scratches live in worker-thread ThreadLocals; closing the supplier in @@ -884,6 +978,11 @@ private void compactLevels(CompactWriter writer, new Scratch(maxCandidateSize, scratchDegree, dimension, sources, pq) ); + // MERGE_LEVELS progress denominator: base-layer live nodes. Level 0 (the bulk) runs first, + // so progress[0] climbs 0 -> numTotalNodes across it, then holds while the small upper tail + // runs (completed is clamped to the total). progress persists across all level calls. + long[] mergeProgress = { 0L, numTotalNodes }; + for (int level = 0; level < maxDegrees.size(); level++) { List batches = buildBatches(level); int searchTopK = Math.max(MIN_SEARCH_TOP_K, ((maxDegrees.get(level) + sources.size() - 1) / sources.size()) * SEARCH_TOP_K_MULTIPLIER); @@ -924,7 +1023,14 @@ private void compactLevels(CompactWriter writer, } catch (IOException e) { throw new RuntimeException(e); } - } + }, + Phase.MERGE_LEVELS, + (results) -> { // exact bytes: read before the write consumes the buffers + long s = 0; + for (WriteResult r : results) s += r.data.remaining(); + return s; + }, + mergeProgress ); } @@ -961,7 +1067,15 @@ private void compactLevels(CompactWriter writer, } catch (IOException e) { throw new RuntimeException(e); } - } + }, + Phase.MERGE_LEVELS, + (results) -> { // estimate: neighbor ids + optional PQ code per node + long s = 0; + for (UpperLayerWriteResult r : results) + s += (long) r.neighbors.length * Integer.BYTES + (r.pqCode == null ? 0 : r.pqCode.length()); + return s; + }, + mergeProgress ); } } @@ -1316,7 +1430,10 @@ private void runBatchesWithBackpressure( List batches, ExecutorCompletionService> ecs, java.util.function.Consumer submitOne, - java.util.function.Consumer> onComplete + java.util.function.Consumer> onComplete, + WorkStage stage, + java.util.function.ToLongFunction> batchBytes, + long[] progress ) throws InterruptedException, ExecutionException { final int total = batches.size(); @@ -1332,11 +1449,22 @@ private void runBatchesWithBackpressure( int completed = 0; while (completed < total) { List results = ecs.take().get(); - onComplete.accept(results); + + // Admission runs on this (orchestrating) thread, before the write. A blocking limiter + // back-pressures dispatch/consume while in-flight workers keep computing — no + // ManagedBlocker. The amount is the bytes this batch is about to write, read here before + // onComplete consumes the buffers. For the default UNLIMITED limiter both calls are no-ops. + long amount = batchBytes.applyAsLong(results); + try (WorkLimiter.Grant g = limiter.acquire(amount)) { + onComplete.accept(results); + } completed++; inFlight--; + progress[0] = Math.min(progress[0] + results.size(), progress[1]); + limiter.onProgress(stage, progress[0], progress[1]); + if (nextToSubmit < total) { submitOne.accept(batches.get(nextToSubmit++)); inFlight++; diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java new file mode 100644 index 000000000..0ebba8f3d --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java @@ -0,0 +1,67 @@ +/* + * 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.util.work; + +import java.util.concurrent.TimeUnit; + +/** + * A leaky-bucket rate meter realizing the {@link WorkLimiter} facet: {@link #acquire} paces the + * aggregate admitted amount to a fixed {@code unitsPerSecond}, blocking the caller when the rate + * would be exceeded. The bucket drains during idle gaps (a burst after a quiet period is not + * charged for the idle time), and the first request after an idle period is admitted without + * delay — the cost of each request is paid by the next one, which is the standard smooth + * shaping behaviour. {@link #onProgress} is inherited as a no-op: this limiter only throttles. + * + *

Thread-safe and reentrant: the emission clock is advanced under a short lock, then the caller + * sleeps outside the lock, so concurrent callers serialize their reservations but wait + * independently. The returned grant is a no-op — the cost is paid entirely at {@code acquire}. + * + *

Obtain instances via {@link ProgressLimiter#rateLimited(double)}. + */ +final class LeakyBucketLimiter implements ProgressLimiter { + private final double nanosPerUnit; + private final Object lock = new Object(); + // Earliest nanoTime at which the next reservation may start. Long.MIN_VALUE until the first + // acquire, so Math.max(now, nextFreeNanos) == now (a fully drained bucket) on the first call. + private long nextFreeNanos = Long.MIN_VALUE; + + LeakyBucketLimiter(double unitsPerSecond) { + if (!(unitsPerSecond > 0) || Double.isInfinite(unitsPerSecond)) { + throw new IllegalArgumentException("unitsPerSecond must be finite and > 0, got " + unitsPerSecond); + } + this.nanosPerUnit = 1_000_000_000.0 / unitsPerSecond; + } + + @Override + public Grant acquire(long amount) throws InterruptedException { + if (amount <= 0) { + return Grant.NOOP; + } + long startAt; + synchronized (lock) { + long now = System.nanoTime(); + startAt = Math.max(now, nextFreeNanos); // drain if idle, else queue behind backlog + long cost = (long) Math.min((double) Long.MAX_VALUE, amount * nanosPerUnit); + nextFreeNanos = startAt + cost; + } + // Sleep (interruptibly, so cancellation aborts) until this request's slot opens. + for (long remaining = startAt - System.nanoTime(); remaining > 0; remaining = startAt - System.nanoTime()) { + TimeUnit.NANOSECONDS.sleep(remaining); + } + return Grant.NOOP; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java new file mode 100644 index 000000000..34a23d9dd --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java @@ -0,0 +1,97 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +import java.util.Objects; +import java.util.function.Consumer; + +/** + * The {@link ProgressTracker tracker} and the {@link WorkLimiter throttle} melded into one control + * surface. A long-running jvector operation accepts a single {@code ProgressLimiter} and uses both + * facets; an embedder may override only the facet it needs — the other defaults to a no-op, so + * {@link #UNLIMITED} behaves exactly as if no SPI were installed. + * + *

Both methods default to no-ops here. A consumer that wants only one facet can still accept a + * lambda via the single-method parents ({@link ProgressTracker}, {@link WorkLimiter}); a consumer + * that wants both accepts a {@code ProgressLimiter}. + */ +@Experimental +public interface ProgressLimiter extends ProgressTracker, WorkLimiter { + + @Override + default void onProgress(WorkStage stage, long completed, long total) { } + + @Override + default Grant acquire(long amount) throws InterruptedException { return Grant.NOOP; } + + /** Observes nothing and limits nothing — behaviour identical to no SPI installed. */ + ProgressLimiter UNLIMITED = new ProgressLimiter() { }; + + /** + * A leaky-bucket rate meter realizing the throttle facet: {@link #acquire} paces the aggregate + * admitted amount to {@code unitsPerSecond} (bytes/sec for the compaction consumer), blocking + * the caller when the rate would be exceeded and draining during idle gaps. {@link #onProgress} + * is a no-op and the returned grant is a no-op (cost is paid at {@code acquire}). Compose with + * {@link #logging(ProgressLimiter, Consumer)} to also log. + * + * @param unitsPerSecond the sustained admission rate; must be finite and {@code > 0} + * @throws IllegalArgumentException if {@code unitsPerSecond} is not finite and positive + */ + static ProgressLimiter rateLimited(double unitsPerSecond) { + return new LeakyBucketLimiter(unitsPerSecond); + } + + /** + * Wraps {@code delegate}, emitting a one-line message to {@code sink} on each + * {@link #onProgress} and on each {@link #acquire} that actually blocked, then delegating both + * facets to {@code delegate}. Composes over any limiter — e.g. + * {@code logging(rateLimited(bytesPerSecond), log::info)} logs a rate-limited operation. The + * delegate's grant is returned unchanged, so a semaphore delegate still releases on close. + * + * @param delegate the limiter to observe and delegate to; {@code null} means {@link #UNLIMITED} + * @param sink receives formatted log lines (e.g. {@code msg -> logger.info(msg)}) + */ + static ProgressLimiter logging(ProgressLimiter delegate, Consumer sink) { + Objects.requireNonNull(sink, "sink"); + final ProgressLimiter d = (delegate == null) ? UNLIMITED : delegate; + return new ProgressLimiter() { + @Override + public void onProgress(WorkStage stage, long completed, long total) { + sink.accept("progress[" + stage.name() + "] " + completed + "/" + (total < 0 ? "?" : Long.toString(total))); + d.onProgress(stage, completed, total); + } + + @Override + public Grant acquire(long amount) throws InterruptedException { + long startNanos = System.nanoTime(); + Grant g = d.acquire(amount); + long waitedMs = (System.nanoTime() - startNanos) / 1_000_000L; + if (waitedMs > 0) { + sink.accept("acquire " + amount + " units - throttled " + waitedMs + "ms"); + } + return g; + } + }; + } + + /** Logging over no throttle: equivalent to {@code logging(UNLIMITED, sink)}. */ + static ProgressLimiter logging(Consumer sink) { + return logging(UNLIMITED, sink); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java new file mode 100644 index 000000000..d7d117269 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java @@ -0,0 +1,43 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Observation contract: receives progress updates for a stage of a long-running operation. + * + *

Best-effort and cheap: implementations must not throw (the caller invokes this on its + * orchestrating thread and treats it as fire-and-forget). See {@link ProgressLimiter} for the + * melded progress + throttle surface that most consumers accept. + */ +@Experimental +@FunctionalInterface +public interface ProgressTracker { + /** + * Reports progress for {@code stage}. + * + * @param stage the stage reporting progress + * @param completed work done so far in this stage, in stage-defined units; monotonically + * non-decreasing within a stage + * @param total total work for this stage, or {@code -1} if not yet known + */ + void onProgress(WorkStage stage, long completed, long total); + + /** A tracker that discards every update. */ + ProgressTracker NOOP = (stage, completed, total) -> { }; +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java new file mode 100644 index 000000000..5355f0757 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkLimiter.java @@ -0,0 +1,59 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Admission contract: blocks until an amount of work may proceed, returning a {@link Grant} that + * the caller closes once the admitted work has completed. + * + *

The unit of {@code amount} is defined by the consumer (e.g. bytes for IO, or rows, + * nodes, items); jvector fixes only the blocking-grant mechanism, never the meaning of the + * quantity. Implementations must be thread-safe and reentrant. {@code acquire} may block but must + * not throw for ordinary back-pressure. + */ +@Experimental +@FunctionalInterface +public interface WorkLimiter { + /** + * Blocks until {@code amount} units of work may proceed. + * + * @param amount the amount of work about to be performed, in consumer-defined units + * @return a non-null grant to {@link Grant#close() close} once that work has completed + * @throws InterruptedException if the calling thread is interrupted while blocked, which + * aborts the operation + */ + Grant acquire(long amount) throws InterruptedException; + + /** + * A handle released by the consumer once the admitted work has completed. For a rate-limiter + * realization (cost paid at {@link WorkLimiter#acquire}) {@link #close()} is a no-op; for a + * semaphore-style in-flight-amount realization it releases the permits taken by {@code acquire}. + */ + interface Grant extends AutoCloseable { + /** Releases the grant. Never throws. */ + @Override + void close(); + + /** A grant that holds nothing and releases nothing. */ + Grant NOOP = () -> { }; + } + + /** A limiter that admits everything immediately. */ + WorkLimiter UNLIMITED = amount -> Grant.NOOP; +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java new file mode 100644 index 000000000..001197a00 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/WorkStage.java @@ -0,0 +1,33 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.annotations.Experimental; + +/** + * Identifies a stage of a long-running operation. The consumer defines its own stages; an + * {@code enum} satisfies this for free via {@link Enum#name()}. + * + *

Part of the generic progress + work-admission SPI ({@link ProgressTracker}, + * {@link WorkLimiter}, {@link ProgressLimiter}). Neither the stage identity nor the unit of work + * is fixed by jvector; both are supplied by the consumer. + */ +@Experimental +public interface WorkStage { + /** The stage's name, stable within a single operation. */ + String name(); +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index 3517c7c40..409859d92 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -36,6 +36,9 @@ import io.github.jbellis.jvector.util.Bits; import io.github.jbellis.jvector.util.BoundedLongHeap; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.util.work.ProgressLimiter; +import io.github.jbellis.jvector.util.work.WorkLimiter; +import io.github.jbellis.jvector.util.work.WorkStage; import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -48,11 +51,16 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.*; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.Semaphore; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.IntFunction; import static io.github.jbellis.jvector.TestUtil.createRandomVectors; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.assertFalse; @@ -536,6 +544,197 @@ public void testCompact() throws Exception { searcher.close(); } + // ---- ProgressLimiter SPI (io.github.jbellis.jvector.util.work) ---- + + /** A ProgressLimiter that records observations and hands out no-op grants (rate-limiter style). */ + private static final class RecordingLimiter implements ProgressLimiter { + final Set stages = ConcurrentHashMap.newKeySet(); + final Map lastCompleted = new ConcurrentHashMap<>(); + final Map lastTotal = new ConcurrentHashMap<>(); + final Map monotonic = new ConcurrentHashMap<>(); + final AtomicInteger acquires = new AtomicInteger(); + final AtomicInteger closes = new AtomicInteger(); + final AtomicLong bytesAcquired = new AtomicLong(); + + @Override + public void onProgress(WorkStage stage, long completed, long total) { + String s = stage.name(); + stages.add(s); + lastTotal.put(s, total); + Long prev = lastCompleted.put(s, completed); + if (prev != null && completed < prev) monotonic.put(s, false); + else monotonic.putIfAbsent(s, true); + } + + @Override + public WorkLimiter.Grant acquire(long amount) { + acquires.incrementAndGet(); + bytesAcquired.addAndGet(amount); + return () -> { closes.incrementAndGet(); }; + } + } + + /** A ProgressLimiter whose grants hold real semaphore permits, released on close. */ + private static final class SemaphoreBytesLimiter implements ProgressLimiter { + final Semaphore permits; + final int cap; + final AtomicInteger open = new AtomicInteger(); + + SemaphoreBytesLimiter(int totalPermits) { + this.permits = new Semaphore(totalPermits); + this.cap = totalPermits; + } + + @Override + public WorkLimiter.Grant acquire(long amount) throws InterruptedException { + final int n = (int) Math.max(1, Math.min(amount, cap)); // never exceed total -> no single-acquire deadlock + permits.acquire(n); + open.incrementAndGet(); + return () -> { permits.release(n); open.decrementAndGet(); }; + } + } + + /** Loads the fixture's source graphs with every node live and identity remapping. */ + private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss) throws IOException { + return newAllLiveCompactor(rss, null); + } + + private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss, ForkJoinPool executor) throws IOException { + List graphs = new ArrayList<>(); + List liveNodes = new ArrayList<>(); + List remappers = new ArrayList<>(); + for (int i = 0; i < numSources; ++i) { + var rs = ReaderSupplierFactory.open(testDirectory.resolve("test_graph_" + i).toAbsolutePath()); + rss.add(rs); + graphs.add(OnDiskGraphIndex.load(rs)); + } + int globalOrdinal = 0; + for (int n = 0; n < numSources; n++) { + Map map = new HashMap<>(numVectorsPerGraph); + for (int i = 0; i < numVectorsPerGraph; i++) map.put(i, globalOrdinal++); + remappers.add(new OrdinalMapper.MapMapper(map)); + var lives = new FixedBitSet(numVectorsPerGraph); + lives.set(0, numVectorsPerGraph); + liveNodes.add(lives); + } + return new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, executor); + } + + @Test + public void testGetTaskWindowSizeReflectsInjectedParallelism() throws Exception { + List rss = new ArrayList<>(); + ForkJoinPool pool = new ForkJoinPool(3); + try { + var compactor = newAllLiveCompactor(rss, pool); + assertEquals("taskWindowSize should equal the injected pool's parallelism", + 3, compactor.getTaskWindowSize()); + } finally { + pool.shutdown(); + for (var rs : rss) rs.close(); + } + } + + @Test + public void testProgressLimiterObservesAndPairsGrants() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var rec = new RecordingLimiter(); + compactor.setProgressLimiter(rec); + // Refinement is opt-in (default off); this test asserts progress across both the + // MERGE_LEVELS and REFINE phases, so it must enable the refinement pass explicitly. + compactor.setRefineAfterCompaction(true); + + var outputPath = testDirectory.resolve("test_compact_progress"); + compactor.compact(outputPath); + + // Both stages observed, progress monotonic and finishing at 100% of a known total. + assertTrue("MERGE_LEVELS progress not reported", rec.stages.contains("MERGE_LEVELS")); + assertTrue("REFINE progress not reported", rec.stages.contains("REFINE")); + assertTrue("MERGE_LEVELS progress not monotonic", rec.monotonic.getOrDefault("MERGE_LEVELS", true)); + assertTrue("REFINE progress not monotonic", rec.monotonic.getOrDefault("REFINE", true)); + assertEquals("MERGE_LEVELS did not reach total", + rec.lastTotal.get("MERGE_LEVELS"), rec.lastCompleted.get("MERGE_LEVELS")); + assertEquals("REFINE did not reach total", + rec.lastTotal.get("REFINE"), rec.lastCompleted.get("REFINE")); + assertEquals("MERGE_LEVELS total should be the live-node count", + (long) numSources * numVectorsPerGraph, (long) rec.lastTotal.get("MERGE_LEVELS")); + + // Throttle invoked with real byte amounts, and every acquire paired with exactly one close. + assertTrue("acquire never called", rec.acquires.get() > 0); + assertTrue("acquire amounts were all zero", rec.bytesAcquired.get() > 0); + assertEquals("every grant must be closed exactly once", rec.acquires.get(), rec.closes.get()); + + // Output is a valid, complete graph. + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testSemaphoreLimiterReleasesAllGrantsWithoutDeadlock() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var sem = new SemaphoreBytesLimiter(64 * 1024 * 1024); // ample: exercises release-on-close, not undersized + compactor.setProgressLimiter(sem); + + var outputPath = testDirectory.resolve("test_compact_semaphore"); + final Throwable[] thrown = new Throwable[1]; + Thread t = new Thread(() -> { + try { compactor.compact(outputPath); } + catch (Throwable e) { thrown[0] = e; } + }, "compact-semaphore"); + t.start(); + t.join(90_000); + + assertFalse("compaction did not finish within 90s — possible throttle deadlock", t.isAlive()); + if (thrown[0] != null) throw new AssertionError("compaction failed under semaphore limiter", thrown[0]); + assertEquals("all semaphore grants must be released (acquire/close paired)", 0, sem.open.get()); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactAtNonZeroStartOffset() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + + // Reserve a prefix (as an embedder would for its container header) and record its bytes. + var outputPath = testDirectory.resolve("test_compact_offset"); + byte[] prefix = new byte[64]; + for (int i = 0; i < prefix.length; i++) prefix[i] = (byte) (0xA0 + (i % 16)); + Files.write(outputPath, prefix); + long startOffset = prefix.length; + + compactor.compact(outputPath, startOffset); + + // The reserved prefix must be untouched by the no-copy write. + byte[] afterPrefix = Arrays.copyOf(Files.readAllBytes(outputPath), prefix.length); + assertArrayEquals("compaction clobbered the reserved prefix", prefix, afterPrefix); + + // The graph loads from startOffset and is complete + searchable (proves refine wrote valid + // records at the base-shifted offsets). + try (ReaderSupplier rs = ReaderSupplierFactory.open(outputPath)) { + var compactGraph = OnDiskGraphIndex.load(rs, startOffset); + assertEquals(numSources * numVectorsPerGraph, compactGraph.size(0)); + + GraphSearcher searcher = new GraphSearcher(compactGraph); + for (int i = 0; i < 5; i++) { + VectorFloat q = allVecs.get(randomIntBetween(0, allVecs.size() - 1)); + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv); + SearchResult r = searcher.search(ssp, 10, Bits.ALL); + assertTrue("search from offset-loaded graph returned nothing", r.getNodes().length > 0); + } + searcher.close(); + } + for (var rs : rss) rs.close(); + } + /** * Tests compaction with deleted nodes. * Verifies that deleted nodes are properly excluded from the compacted graph. diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java new file mode 100644 index 000000000..16d02f31c --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java @@ -0,0 +1,289 @@ +/* + * 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.util.work; + +import io.github.jbellis.jvector.util.work.WorkLimiter.Grant; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestProgressLimiter { + + private static final WorkStage STAGE = () -> "TEST"; + + private static long millisFor(ThrowingRunnable r) throws Exception { + long t0 = System.nanoTime(); + r.run(); + return (System.nanoTime() - t0) / 1_000_000L; + } + + private interface ThrowingRunnable { void run() throws Exception; } + + // ---- rateLimited (leaky bucket) ---- + + @Test + public void rateLimitedRejectsNonPositiveOrNonFiniteRate() { + for (double bad : new double[]{0.0, -1.0, -0.0, Double.NaN, Double.POSITIVE_INFINITY}) { + try { + ProgressLimiter.rateLimited(bad); + fail("expected IllegalArgumentException for rate " + bad); + } catch (IllegalArgumentException expected) { + // ok + } + } + } + + @Test + public void rateLimitedAdmitsZeroOrNegativeAmountImmediately() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1.0); // 1 unit/sec: any real wait would be seconds + long ms = millisFor(() -> { + try (Grant g = limiter.acquire(0)) { assertNotNull(g); } + try (Grant g = limiter.acquire(-100)) { assertNotNull(g); } + }); + assertTrue("zero/negative amount must not block, waited " + ms + "ms", ms < 500); + } + + @Test + public void rateLimitedPacesSubsequentAcquire() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1000.0); // 1 unit/ms + limiter.acquire(200).close(); // warmup: drained bucket admits the first request immediately + + long ms = millisFor(() -> limiter.acquire(200).close()); // must wait ~200ms behind the warmup reservation + assertTrue("expected pacing >= ~100ms at 1000 units/s after a 200-unit warmup, got " + ms + "ms", ms >= 100); + } + + @Test + public void rateLimitedFirstAcquireIsNotDelayed() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(10.0); // slow: a delayed first call would be seconds + long ms = millisFor(() -> limiter.acquire(1000).close()); + assertTrue("first acquire on a drained bucket must not block, waited " + ms + "ms", ms < 500); + } + + @Test + public void rateLimitedIsInterruptible() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(100.0); // 100 units/sec + limiter.acquire(100).close(); // warmup reserves ~1s of future emission time + + AtomicReference caught = new AtomicReference<>(); + AtomicInteger returnedNormally = new AtomicInteger(); + Thread t = new Thread(() -> { + try { + limiter.acquire(1).close(); // blocks ~1s behind the warmup reservation + returnedNormally.incrementAndGet(); + } catch (Throwable e) { + caught.set(e); + } + }, "rate-limited-blocked"); + t.start(); + Thread.sleep(150); // let it reach the interruptible sleep + t.interrupt(); + t.join(5_000); + + assertFalse("interrupted acquire should not hang", t.isAlive()); + assertEquals("acquire should not have returned normally", 0, returnedNormally.get()); + assertTrue("expected InterruptedException, got " + caught.get(), + caught.get() instanceof InterruptedException); + } + + @Test + public void rateLimitedGrantIsNoopAndProgressIsNoop() throws Exception { + ProgressLimiter limiter = ProgressLimiter.rateLimited(1_000_000.0); + Grant g = limiter.acquire(10); + assertNotNull(g); + g.close(); + g.close(); // idempotent no-op + limiter.onProgress(STAGE, 1, 2); // rate limiter does not track progress; must not throw + } + + // ---- logging wrapper ---- + + @Test(expected = NullPointerException.class) + public void loggingRejectsNullSinkWithDelegate() { + ProgressLimiter.logging(ProgressLimiter.UNLIMITED, null); + } + + @Test(expected = NullPointerException.class) + public void loggingRejectsNullSink() { + ProgressLimiter.logging((java.util.function.Consumer) null); + } + + @Test + public void loggingNullDelegateBehavesAsUnlimited() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(null, log::add); + long ms = millisFor(() -> limiter.acquire(Long.MAX_VALUE).close()); // UNLIMITED: instant + assertTrue("null delegate should not throttle, waited " + ms + "ms", ms < 500); + } + + @Test + public void loggingDelegatesBothFacets() { + RecordingLimiter delegate = new RecordingLimiter(); + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(delegate, log::add); + + limiter.onProgress(STAGE, 3, 10); + assertEquals("onProgress must be delegated", 1, delegate.progressCalls.get()); + assertEquals(3, delegate.lastCompleted); + assertEquals(10, delegate.lastTotal); + assertTrue("onProgress should have been logged", + log.stream().anyMatch(s -> s.contains("TEST") && s.contains("3/10"))); + } + + @Test + public void loggingPreservesDelegateGrant() throws Exception { + RecordingLimiter delegate = new RecordingLimiter(); + ProgressLimiter limiter = ProgressLimiter.logging(delegate, s -> { }); + + Grant g = limiter.acquire(1234); + assertEquals("acquire must be delegated", 1, delegate.acquireCalls.get()); + assertEquals(1234, delegate.lastAmount); + assertEquals("grant must not be closed yet", 0, delegate.grantCloses.get()); + g.close(); + assertEquals("closing the wrapper grant must close the delegate's grant", 1, delegate.grantCloses.get()); + } + + @Test + public void loggingLogsAcquireOnlyWhenItBlocks() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + + // Instant delegate (UNLIMITED): no throttled line expected. + ProgressLimiter fast = ProgressLimiter.logging(ProgressLimiter.UNLIMITED, log::add); + fast.acquire(500).close(); + assertTrue("unblocked acquire should not log a throttle line", + log.stream().noneMatch(s -> s.contains("throttled"))); + + // Blocking delegate: a throttled line is expected. + log.clear(); + ProgressLimiter slow = ProgressLimiter.logging(new SleepingLimiter(60), log::add); + slow.acquire(500).close(); + assertTrue("blocked acquire should log a throttle line", + log.stream().anyMatch(s -> s.contains("throttled") && s.contains("500"))); + } + + // ---- composition ---- + + @Test + public void loggingComposesWithRateLimited() throws Exception { + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(ProgressLimiter.rateLimited(1000.0), log::add); + + limiter.acquire(200).close(); // warmup + long ms = millisFor(() -> limiter.acquire(200).close()); + + assertTrue("composed limiter should still pace, got " + ms + "ms", ms >= 100); + assertTrue("composed limiter should log the throttled acquire", + log.stream().anyMatch(s -> s.contains("throttled"))); + limiter.onProgress(STAGE, 5, 5); + assertTrue("composed limiter should log progress", + log.stream().anyMatch(s -> s.contains("TEST") && s.contains("5/5"))); + } + + // ---- melded SPI defaults ---- + + @Test + public void unlimitedIsFullyNoop() throws Exception { + long ms = millisFor(() -> { + try (Grant g = ProgressLimiter.UNLIMITED.acquire(Long.MAX_VALUE)) { + assertNotNull(g); + } + }); + assertTrue("UNLIMITED.acquire must not block, waited " + ms + "ms", ms < 500); + ProgressLimiter.UNLIMITED.onProgress(STAGE, 7, -1); // no-op, must not throw + + // Facet no-op constants exist and are safe. + WorkLimiter.Grant.NOOP.close(); + ProgressTracker.NOOP.onProgress(STAGE, 1, 1); + try (Grant g = WorkLimiter.UNLIMITED.acquire(99)) { + assertNotNull(g); + } + } + + @Test + public void facetsAreIndependentlyOverridable() throws Exception { + // Tracker-only: overrides onProgress, inherits no-op acquire. + AtomicInteger progressSeen = new AtomicInteger(); + ProgressLimiter trackerOnly = new ProgressLimiter() { + @Override public void onProgress(WorkStage stage, long completed, long total) { + progressSeen.incrementAndGet(); + } + }; + try (Grant g = trackerOnly.acquire(1_000_000)) { // inherited no-op: must not block + assertNotNull(g); + } + trackerOnly.onProgress(STAGE, 1, 1); + assertEquals(1, progressSeen.get()); + + // Throttle-only: overrides acquire, inherits no-op onProgress. + AtomicInteger acquireSeen = new AtomicInteger(); + ProgressLimiter throttleOnly = new ProgressLimiter() { + @Override public Grant acquire(long amount) { + acquireSeen.incrementAndGet(); + return Grant.NOOP; + } + }; + throttleOnly.onProgress(STAGE, 1, 1); // inherited no-op: must not throw + throttleOnly.acquire(5).close(); + assertEquals(1, acquireSeen.get()); + } + + // ---- test doubles ---- + + /** Records both facets and hands out a grant whose close is counted. */ + private static final class RecordingLimiter implements ProgressLimiter { + final AtomicInteger progressCalls = new AtomicInteger(); + final AtomicInteger acquireCalls = new AtomicInteger(); + final AtomicInteger grantCloses = new AtomicInteger(); + volatile long lastCompleted, lastTotal, lastAmount; + + @Override + public void onProgress(WorkStage stage, long completed, long total) { + progressCalls.incrementAndGet(); + lastCompleted = completed; + lastTotal = total; + } + + @Override + public Grant acquire(long amount) { + acquireCalls.incrementAndGet(); + lastAmount = amount; + return grantCloses::incrementAndGet; + } + } + + /** A throttle that always blocks for a fixed number of milliseconds. */ + private static final class SleepingLimiter implements ProgressLimiter { + private final long sleepMillis; + + SleepingLimiter(long sleepMillis) { this.sleepMillis = sleepMillis; } + + @Override + public Grant acquire(long amount) throws InterruptedException { + Thread.sleep(sleepMillis); + return Grant.NOOP; + } + } +} From a96fd1321848a8b7efb4de274af6a9b62f1d9ca2 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 1 Jul 2026 17:41:57 +0000 Subject: [PATCH 06/13] additional refinements for type safety, thread pools, and embedding flexibility --- .../jvector/bench/CompactorBenchmark.java | 2 +- docs/compaction.md | 98 +++++++++- .../jvector/disk/FileChannelSeekableSink.java | 70 +++++++ .../jbellis/jvector/disk/SeekableSink.java | 65 +++++++ .../jvector/graph/disk/CompactionContext.java | 6 +- .../graph/disk/CompactionDestination.java | 83 +++++++++ .../graph/disk/FileCompactionDestination.java | 66 +++++++ .../graph/disk/OnDiskGraphIndexCompactor.java | 133 +++++++++---- .../jvector/example/CompactionBench.java | 2 +- .../jvector/disk/TestSeekableSink.java | 85 +++++++++ .../disk/TestOnDiskGraphIndexCompactor.java | 176 +++++++++++++++--- 11 files changed, 720 insertions(+), 66 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java diff --git a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java index 3a3b9c8bf..bf3acb695 100644 --- a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java +++ b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java @@ -704,7 +704,7 @@ private long compactPartitions() throws Exception { liveNodes.add(randomLiveNodes(size, liveNodesRate, n)); globalOrdinal += size; } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); long startNanos = System.nanoTime(); compactor.compact(compactOutputPath); diff --git a/docs/compaction.md b/docs/compaction.md index af0c12317..a24c4485a 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -33,7 +33,8 @@ for (var src : sources) { var compactor = new OnDiskGraphIndexCompactor( sources, liveNodes, remappers, VectorSimilarityFunction.COSINE, - /* executor= */ null // null = use by default shared threadpool in compactor + /* executor= */ null, // null = jvector's shared physical-core pool + /* taskWindowSize= */ -1 // <= 0 derives the window from the pool's parallelism ); compactor.compact(Path.of("compacted.index")); @@ -57,6 +58,99 @@ for (int i = 0; i < source.size(); i++) { remappers.add(new OrdinalMapper.MapMapper(oldToNew)); ``` +## Embedding API + +For embedding the compactor into a host system (for example, a database's compaction pipeline), this branch adds `@Experimental` extension points that let the host supply its own threads, observe and throttle the merge, and place the output inside its own container file. All are additive: with none of them used, behavior and output are unchanged (a jvector-owned pool, no throttling, a standalone output file). + +### Supplying an executor + +The compactor dispatches its batch work through an internal `ExecutorCompletionService`, so it runs on any `Executor` the host supplies, with an explicit in-flight window: + +```java +new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, sim, executor, taskWindowSize); +``` + +- Pass a **`ForkJoinPool`** (work-stealing makes submit-and-block safe), or a **caller-runs** executor (`Runnable::run`) to run the entire merge on the *calling* thread — no jvector-owned pool. A host then gets parallelism by running multiple compactions concurrently rather than from a per-compaction pool. +- **Do not** pass a bounded `ThreadPoolExecutor` that is *also* running the calling thread: the caller blocks in `take()` while its sub-tasks queue behind it, which can thread-starvation-deadlock. +- `taskWindowSize` bounds in-flight batches (concurrency and peak write-side memory); `<= 0` derives it from a `ForkJoinPool`'s parallelism, else defaults to `1`. Read it back with `getTaskWindowSize()`. + +A `ForkJoinPool` is passed as the `executor` (with `taskWindowSize <= 0` to derive its window) — there is no separate `ForkJoinPool` constructor. + +### Progress and throttling + +`setProgressLimiter(ProgressLimiter)` installs a control surface (package `io.github.jbellis.jvector.util.work`) that melds two independently-optional facets: + +- **`onProgress(WorkStage stage, long completed, long total)`** — progress observation. Stages are `OnDiskGraphIndexCompactor.Phase.{MERGE_LEVELS, REFINE}`; `total` may be `-1` until known. +- **`acquire(long amount) → Grant`** — blocking work admission. For the compactor the unit is **bytes about to be written**. `acquire` is called on the orchestrating thread (never a pool worker), so a blocking limiter back-pressures dispatch without a `ForkJoinPool.ManagedBlocker` — exactly like ordinary code blocking on a rate limiter. + +The default is `ProgressLimiter.UNLIMITED` (both facets no-op), so output and timing are unchanged when unset. + +Two ready-made implementations compose as decorators: + +```java +// Cap merge write bandwidth to 100 MB/s, logging throttled writes and progress. +compactor.setProgressLimiter( + ProgressLimiter.logging( + ProgressLimiter.rateLimited(100.0 * 1024 * 1024), // leaky-bucket meter, bytes/sec + msg -> log.info("compaction: {}", msg))); +``` + +- `rateLimited(unitsPerSecond)` — a leaky-bucket rate meter (drains when idle, interruptible; grant is a no-op). +- `logging(delegate, sink)` / `logging(sink)` — logs each `onProgress` and each *blocked* `acquire`, delegating both facets; a `sink` is any `Consumer`. + +A host that draws from its own shared budget implements `acquire` directly: + +```java +compactor.setProgressLimiter(new ProgressLimiter() { + @Override public void onProgress(WorkStage stage, long completed, long total) { + metrics.report(stage, completed, total); + } + @Override public Grant acquire(long bytes) throws InterruptedException { + hostIoBudget.acquire(bytes); // block against a host-wide IO limiter + return Grant.NOOP; // rate-limiter model: nothing to release + } +}); +``` + +The returned `Grant` is closed when the admitted work completes. A **rate-limiter** realization pays its cost at `acquire` and returns `Grant.NOOP`; a **semaphore** realization (permits = in-flight bytes) releases on `Grant.close()`. + +> **Caveat (semaphore realization).** During `REFINE`, batches are submitted and drained through a sliding window of `taskWindowSize`, so a permit-releasing semaphore must admit at least `taskWindowSize` batches' worth of bytes or the window can't fill. `MERGE_LEVELS` has no such constraint, and the rate-limiter realization has none in either phase. + +### Output destination (no temp-file copy) + +By default `compact(Path)` writes a standalone file. To place the graph body directly inside a host container after a reserved header — eliminating a temp-file-and-copy — use either overload: + +```java +// (1) write the body into an existing file at a reserved offset; bytes in [0, startOffset) are preserved. +compactor.compact(componentPath, /* startOffset= */ headerSize); + +// (2) resource-scoped destination with a commit/abort lifecycle; returns the body length. +long bodyLength = compactor.compact(CompactionDestination.toFile(outputPath)); +``` + +`compact(CompactionDestination)` opens one `Target` per call and drives its lifecycle: + +```java +CompactionDestination dest = () -> { + FileChannel ch = FileChannel.open(componentPath, CREATE, WRITE, READ); + writeHostHeader(ch); // reserve [0, headerSize) + return new CompactionDestination.Target() { + public Path file() { return componentPath; } + public long startOffset() { return headerSize; } + public void commit(long bodyLength) throws IOException { // success: finalize the container + writeHostFooter(ch, bodyLength); + ch.force(true); + } + public void close() throws IOException { ch.close(); } // always runs; no commit ⇒ host discards the file + }; +}; +long bodyLength = compactor.compact(dest); +``` + +- `commit(bodyLength)` fires exactly once, only on success, after the body is written and forced. `close()` always runs (try-with-resources); reaching it without a prior `commit` is an unambiguous abort — discard the partial output. +- The compactor writes into `file()` at `startOffset()` using its own random-access and memory-mapped IO — the destination expresses *where*, not *how*. +- To read the committed body back (e.g. for a checksum), address the region with the `SeekableSink` primitive (package `io.github.jbellis.jvector.disk`): `SeekableSink.over(channel, target.startOffset())` gives region-relative `writeAt`/`readAt`. + ## Algorithm ### Ordinal Remapping @@ -104,7 +198,7 @@ for alpha in [1.0, 1.2]: Level 0 (base layer) stores inline vectors, FusedPQ codes, and the neighbor list. Upper levels store only the neighbor list (plus PQ codes at level 1 for cross-level searching). -Processing is batched per source and run in parallel across sources using a `ForkJoinPool`. A backpressure window keeps at most `taskWindowSize` batches in-flight at once, bounding memory use. +Processing is batched per source and run in parallel across sources on the supplied executor (a `ForkJoinPool` by default; see [Embedding API](#embedding-api)). A backpressure window keeps at most `taskWindowSize` batches in-flight at once, bounding memory use. ### Entry Node diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java new file mode 100644 index 000000000..bcc3bd412 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/FileChannelSeekableSink.java @@ -0,0 +1,70 @@ +/* + * 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.disk; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; + +/** + * {@link SeekableSink} over a {@link FileChannel}, translating region-relative positions by a fixed + * base offset. The channel is owned by the caller; {@link #close()} does not close it. + */ +final class FileChannelSeekableSink implements SeekableSink { + private final FileChannel channel; + private final long baseOffset; + + FileChannelSeekableSink(FileChannel channel, long baseOffset) { + if (channel == null) { + throw new NullPointerException("channel"); + } + if (baseOffset < 0) { + throw new IllegalArgumentException("baseOffset must be >= 0, got " + baseOffset); + } + this.channel = channel; + this.baseOffset = baseOffset; + } + + @Override + public void writeAt(long position, ByteBuffer src) throws IOException { + if (position < 0) { + throw new IllegalArgumentException("position must be >= 0, got " + position); + } + long abs = baseOffset + position; + while (src.hasRemaining()) { + abs += channel.write(src, abs); + } + } + + @Override + public int readAt(long position, ByteBuffer dst) throws IOException { + if (position < 0) { + throw new IllegalArgumentException("position must be >= 0, got " + position); + } + return channel.read(dst, baseOffset + position); + } + + @Override + public void force() throws IOException { + channel.force(false); + } + + @Override + public void close() { + // The channel is owned by the caller, per SeekableSink.over(...). + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java new file mode 100644 index 000000000..85833062f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SeekableSink.java @@ -0,0 +1,65 @@ +/* + * 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.disk; + +import io.github.jbellis.jvector.annotations.Experimental; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; + +/** + * A seekable region that supports positional reads and writes, addressed in coordinates relative + * to the region's start (0-based). An embedder uses it to hand a compactor (or other writer) a + * bounded window inside a larger container file: positions are region-relative and the + * implementation adds the container's base offset, so the writer never needs to know the absolute + * offset. + * + *

Implementations must support concurrent positional writes and reads to disjoint ranges (a + * {@link FileChannel} does). This is a generic IO primitive; the compaction extension point that + * hands one out is {@code io.github.jbellis.jvector.graph.disk.CompactionDestination}. + */ +@Experimental +public interface SeekableSink extends AutoCloseable { + + /** Write {@code src} fully at region-relative {@code position} (must be {@code >= 0}). */ + void writeAt(long position, ByteBuffer src) throws IOException; + + /** + * Read up to {@code dst.remaining()} bytes at region-relative {@code position} (must be + * {@code >= 0}); returns the number of bytes read, or {@code -1} at end of region. + */ + int readAt(long position, ByteBuffer dst) throws IOException; + + /** Force written bytes to durable storage. */ + void force() throws IOException; + + @Override + void close() throws IOException; + + /** + * Reference implementation over a {@link FileChannel} region. Every region-relative position is + * translated by {@code baseOffset}. The channel's lifecycle is owned by the caller — this + * {@link #close()} does not close the channel. + * + * @param channel the backing channel, opened for read and write + * @param baseOffset the absolute offset of the region's start within {@code channel} ({@code >= 0}) + */ + static SeekableSink over(FileChannel channel, long baseOffset) { + return new FileChannelSeekableSink(channel, baseOffset); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java index b01901832..b0deca442 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java @@ -21,7 +21,7 @@ import java.util.Collections; import java.util.List; -import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ExecutorService; /** * Bundle of inputs that {@link QuantizationCompactionStrategy} implementations need to do their work @@ -38,7 +38,7 @@ public final class CompactionContext { public final List remappers; public final int dimension; public final int maxOrdinal; - public final ForkJoinPool executor; + public final ExecutorService executor; public final int taskWindowSize; public CompactionContext( @@ -48,7 +48,7 @@ public CompactionContext( List remappers, int dimension, int maxOrdinal, - ForkJoinPool executor, + ExecutorService executor, int taskWindowSize) { this.sources = Collections.unmodifiableList(sources); this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java new file mode 100644 index 000000000..5049333ad --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java @@ -0,0 +1,83 @@ +/* + * 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.disk; + +import io.github.jbellis.jvector.annotations.Experimental; +import io.github.jbellis.jvector.disk.SeekableSink; + +import java.io.IOException; +import java.nio.file.Path; + +/** + * Embedding extension point: tells {@link OnDiskGraphIndexCompactor} WHERE to write its compacted + * graph, so the body lands directly inside the embedder's container (after a header the embedder + * reserves) — eliminating the temp-file-and-copy. Resource-scoped: the compactor uses one + * {@link Target} per {@code compact(...)} call, commits on success, and always closes. + * + * 

{@code
+ *   try (CompactionDestination.Target t = destination.open()) {
+ *       // ...compactor writes the graph into t.file() at t.startOffset()...
+ *       t.commit(bodyLength);   // success: body written & durable; embedder finalizes its footer
+ *   }                           // close() always runs; no commit() => aborted (discard partial output)
+ * }
+ * + *

The compactor needs a real file (it uses a memory-mapped read-back during refinement and a + * random-access writer), so a {@link Target} is expressed as a container {@link Path} plus a base + * offset rather than an opaque stream. The generic {@link SeekableSink} primitive addresses the same + * window in region-relative coordinates and is what an embedder uses to read the committed body back + * for its checksum, e.g. {@code SeekableSink.over(channel, target.startOffset())}. + */ +@FunctionalInterface +@Experimental +public interface CompactionDestination { + + /** Open a fresh target for one compaction. */ + Target open() throws IOException; + + /** One compaction's output region plus its commit/abort lifecycle. */ + interface Target extends AutoCloseable { + + /** The container file the graph body is written into. */ + Path file(); + + /** The byte offset within {@link #file()} at which the graph body begins ({@code >= 0}). */ + long startOffset(); + + /** + * Signalled exactly once, after the body has been fully written and forced, reporting its + * length ({@code file() size - startOffset()}). The embedder finalizes its container here + * (e.g. writes a footer/checksum). MUST be called before {@link #close()} on the success path. + */ + void commit(long bodyLength) throws IOException; + + /** + * Always runs (try-with-resources). If reached without a prior {@link #commit}, the + * compaction failed and the embedder discards the partial output; releases embedder resources. + */ + @Override + void close() throws IOException; + } + + /** + * Default standalone destination: writes to its own file at offset {@code 0} (today's + * {@code compact(Path)} behaviour). {@code commit} is a no-op marker; a {@code close} without a + * prior commit deletes the partial file. + */ + static CompactionDestination toFile(Path path) { + return new FileCompactionDestination(path); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java new file mode 100644 index 000000000..24d610d0e --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.java @@ -0,0 +1,66 @@ +/* + * 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.disk; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; + +/** + * {@link CompactionDestination} that writes a standalone graph file at offset 0. Backs + * {@link CompactionDestination#toFile(Path)}. + */ +final class FileCompactionDestination implements CompactionDestination { + private final Path path; + + FileCompactionDestination(Path path) { + if (path == null) { + throw new NullPointerException("path"); + } + this.path = path; + } + + @Override + public Target open() { + return new Target() { + private boolean committed; + + @Override + public Path file() { + return path; + } + + @Override + public long startOffset() { + return 0L; + } + + @Override + public void commit(long bodyLength) { + // Standalone file: the graph IS the whole file; compact() already wrote and flushed it. + committed = true; + } + + @Override + public void close() throws IOException { + if (!committed) { + Files.deleteIfExists(path); // abort: discard the partial file + } + } + }; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index 9f1c89d68..95bea2a9b 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -84,7 +84,7 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final int dimension; private int maxOrdinal = -1; private int numTotalNodes = 0; - private final ForkJoinPool executor; + private final Executor executor; private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; private boolean refineAfterCompaction = false; @@ -145,36 +145,34 @@ public int getTaskWindowSize() { public void setRefineAfterCompaction(boolean refineAfterCompaction) { this.refineAfterCompaction = refineAfterCompaction; } - /** - * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. - * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. - */ - @Experimental - public OnDiskGraphIndexCompactor( - List sources, - List liveNodes, - List remappers, - VectorSimilarityFunction similarityFunction, - ForkJoinPool executor) { - this(sources, null, liveNodes, remappers, similarityFunction, executor); - } /** - * Constructs a new OnDiskGraphIndexCompactor to merge multiple graph indexes. - * Initializes thread pool, validates inputs, and prepares metadata for compaction. + * Primary constructor: merges multiple graph indexes using any {@link Executor} with an + * explicit in-flight window. * * @param sourceCompressed parallel to {@code sources}, supplying the non-fused compressed * vectors (e.g. {@link io.github.jbellis.jvector.quantization.PQVectors}) * that ship alongside each graph. Pass {@code null} when sources carry * quantization inline (FUSED_PQ) or have none. Must not be combined * with sources that carry the FUSED_PQ feature. - * @param executor the pool that runs compaction batches. Its {@code getParallelism()} sets - * both CPU concurrency and the in-flight task/memory window - * ({@link #getTaskWindowSize()}), so one knob bounds two dimensions. Compaction - * is compute- and memory-bandwidth-bound; size this from physical cores — - * logical-core sizing oversubscribes hyperthreaded hosts and costs throughput. - * The compactor never owns or shuts down the pool. Pass {@code null} to use the - * shared {@link PhysicalCoreExecutor#pool()} default. + * @param executor runs compaction batches submitted via an internal + * {@link java.util.concurrent.ExecutorCompletionService} while the calling + * thread blocks for their completion. Safe to pass a {@link ForkJoinPool} + * (work-stealing + managed blocking make submit-and-block safe) or a + * caller-runs / same-thread executor (e.g. {@code Runnable::run}), which + * runs each batch synchronously on the calling thread — no worker threads, no + * separate pool. It is not safe to pass a bounded + * {@code ThreadPoolExecutor} that is also running the calling thread: the + * caller blocks in {@code take()} while its sub-tasks queue behind it, which can + * thread-starvation-deadlock (worst case a single-thread pool). Embedders that + * want their own bounded pool without a dedicated fan-out pool should pass a + * caller-runs executor and derive parallelism from the number of concurrent + * compactions instead. Pass {@code null} to use the shared + * {@link PhysicalCoreExecutor#pool()} default. The compactor never owns or shuts + * down the executor. + * @param taskWindowSize bounds the number of in-flight batches (both concurrency and peak + * write-side memory). {@code <= 0} derives it from a {@link ForkJoinPool}'s + * {@code getParallelism()}, else defaults to 1 (serial). */ @Experimental public OnDiskGraphIndexCompactor( @@ -183,22 +181,16 @@ public OnDiskGraphIndexCompactor( List liveNodes, List remappers, VectorSimilarityFunction similarityFunction, - ForkJoinPool executor) { + Executor executor, + int taskWindowSize) { checkBeforeCompact(sources, sourceCompressed, liveNodes, remappers); - if (executor != null) { - this.executor = executor; - } else { - // Default to the shared physical-core pool. Compaction (PQ encode + parallel record - // flush + refinement) is compute- and memory-bandwidth-bound, so sizing to logical - // cores oversubscribes hyperthreaded hosts and costs throughput. This pool is - // process-wide and shared with index construction and quantization; the compactor - // never owns or shuts it down. - this.executor = PhysicalCoreExecutor.pool(); - } - // Track the pool's real parallelism so task-window / backpressure sizing stays correct - // whether the executor is the shared default or a caller-injected pool. - this.taskWindowSize = this.executor.getParallelism(); + // Default to the shared physical-core pool. Compaction (PQ encode + parallel record flush + + // refinement) is compute- and memory-bandwidth-bound, so sizing to logical cores + // oversubscribes hyperthreaded hosts. This pool is process-wide and shared with index + // construction and quantization; the compactor never owns or shuts it down. + this.executor = (executor != null) ? executor : PhysicalCoreExecutor.pool(); + this.taskWindowSize = resolveWindow(this.executor, taskWindowSize); this.sources = sources; this.sourceCompressed = (sourceCompressed == null || sourceCompressed.isEmpty()) ? null : sourceCompressed; @@ -222,6 +214,49 @@ public OnDiskGraphIndexCompactor( this.similarityFunction = similarityFunction; } + /** + * Convenience {@link Executor} constructor without a non-fused compressed sidecar. Equivalent + * to the 7-arg form with {@code sourceCompressed = null}. + */ + @Experimental + public OnDiskGraphIndexCompactor( + List sources, + List liveNodes, + List remappers, + VectorSimilarityFunction similarityFunction, + Executor executor, + int taskWindowSize) { + this(sources, null, liveNodes, remappers, similarityFunction, executor, taskWindowSize); + } + + private static int resolveWindow(Executor executor, int requested) { + if (requested > 0) return requested; + if (executor instanceof ForkJoinPool) return ((ForkJoinPool) executor).getParallelism(); + return 1; // arbitrary Executor with no declared parallelism → serial window + } + + /** + * Adapts an arbitrary {@link Executor} to the {@link ExecutorService} the quantization + * strategies need for {@code invokeAll} during PQ pre-encode. A real {@code ExecutorService} + * (including a {@link ForkJoinPool}) is used directly; a plain {@code Executor} — notably a + * caller-runs {@code Runnable::run} — is wrapped so its tasks run on whatever thread the + * executor dispatches to (the calling thread, for caller-runs). No pool is created, and the + * lifecycle methods are inert since the compactor never owns the executor. + */ + private static ExecutorService asExecutorService(Executor executor) { + if (executor instanceof ExecutorService) { + return (ExecutorService) executor; + } + return new AbstractExecutorService() { + @Override public void execute(Runnable command) { executor.execute(command); } + @Override public void shutdown() { } + @Override public List shutdownNow() { return Collections.emptyList(); } + @Override public boolean isShutdown() { return false; } + @Override public boolean isTerminated() { return false; } + @Override public boolean awaitTermination(long timeout, TimeUnit unit) { return false; } + }; + } + /** * Validates that all source indexes have compatible configurations and required features * before attempting compaction. Ensures consistent dimensions, max degrees, hierarchical @@ -396,6 +431,28 @@ public void compact(Path outputPath, long startOffset) throws FileNotFoundExcept } } + /** + * No-copy compaction into an embedder-supplied {@link CompactionDestination}: writes the graph + * body into the destination's container at its reserved offset, then {@code commit}s the body + * length on success (so the embedder can finalize its footer/checksum), or — on any failure — + * closes the target without committing (so the embedder discards the partial output). The + * compactor writes into {@code target.file()} at {@code target.startOffset()} using its own + * random-access + memory-mapped IO; the destination expresses where, not how. + * + * @return the graph body length written, i.e. {@code size(file) - startOffset} + */ + @Experimental + public long compact(CompactionDestination destination) throws IOException { + try (CompactionDestination.Target target = destination.open()) { + Path file = target.file(); + long base = target.startOffset(); + compact(file, base); + long bodyLength = java.nio.file.Files.size(file) - base; + target.commit(bodyLength); + return bodyLength; + } + } + /** * Compaction entry point for graphs that ship a non-fused compressed sidecar (e.g. * {@link io.github.jbellis.jvector.quantization.PQVectors}). Writes the merged graph to @@ -480,7 +537,7 @@ private QuantizationCompactionStrategy detectSidecarStrategy() { /** Snapshot the compactor's state into a {@link CompactionContext} for strategies to consume. */ private CompactionContext buildContext() { return new CompactionContext(sources, sourceCompressed, liveNodes, remappers, - dimension, maxOrdinal, executor, taskWindowSize); + dimension, maxOrdinal, asExecutorService(executor), taskWindowSize); } /** diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java index 15543ebbc..a60cb2ed2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java @@ -172,7 +172,7 @@ private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg, // Compact Path compactPath = tempDir.resolve("compacted"); - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, vsf, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, vsf, null, -1); long t0 = System.currentTimeMillis(); compactor.compact(compactPath); long compactionMs = System.currentTimeMillis() - t0; diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java new file mode 100644 index 000000000..9fde6b162 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/disk/TestSeekableSink.java @@ -0,0 +1,85 @@ +/* + * 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.disk; + +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestSeekableSink { + + @Test + public void writesAndReadsInRegionRelativeCoordinates() throws IOException { + Path f = Files.createTempFile("sink", ".bin"); + try (FileChannel ch = FileChannel.open(f, StandardOpenOption.WRITE, StandardOpenOption.READ)) { + long base = 100; + SeekableSink sink = SeekableSink.over(ch, base); + sink.writeAt(0, ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8))); + sink.writeAt(5, ByteBuffer.wrap("WORLD".getBytes(StandardCharsets.UTF_8))); + sink.force(); + + // Region-relative read returns what was written. + ByteBuffer dst = ByteBuffer.allocate(10); + assertEquals(10, sink.readAt(0, dst)); + assertEquals("helloWORLD", new String(dst.array(), StandardCharsets.UTF_8)); + + // The bytes actually land at the absolute base offset (region-relative -> absolute). + ByteBuffer raw = ByteBuffer.allocate(10); + ch.read(raw, base); + assertEquals("helloWORLD", new String(raw.array(), StandardCharsets.UTF_8)); + + // Nothing was written before the region. + ByteBuffer before = ByteBuffer.allocate((int) base); + ch.read(before, 0); + for (byte b : before.array()) { + assertEquals("region must not write before its base", 0, b); + } + + // close() must NOT close the caller-owned channel. + sink.close(); + assertTrue("sink.close() must not close the caller's channel", ch.isOpen()); + } + Files.deleteIfExists(f); + } + + @Test + public void rejectsNegativeBaseAndPosition() throws IOException { + Path f = Files.createTempFile("sink", ".bin"); + try (FileChannel ch = FileChannel.open(f, StandardOpenOption.WRITE, StandardOpenOption.READ)) { + try { SeekableSink.over(ch, -1); fail("negative base"); } catch (IllegalArgumentException expected) { } + SeekableSink sink = SeekableSink.over(ch, 0); + try { sink.writeAt(-1, ByteBuffer.allocate(1)); fail("negative write pos"); } catch (IllegalArgumentException expected) { } + try { sink.readAt(-1, ByteBuffer.allocate(1)); fail("negative read pos"); } catch (IllegalArgumentException expected) { } + } + Files.deleteIfExists(f); + } + + @Test(expected = NullPointerException.class) + public void rejectsNullChannel() { + SeekableSink.over(null, 0); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index 409859d92..e9bc30d49 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -52,6 +52,7 @@ import java.nio.file.Path; import java.util.*; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executor; import java.util.concurrent.ForkJoinPool; import java.util.concurrent.Semaphore; import java.util.concurrent.atomic.AtomicInteger; @@ -64,6 +65,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.fail; @ThreadLeakScope(ThreadLeakScope.Scope.NONE) public class TestOnDiskGraphIndexCompactor extends RandomizedTest { @@ -278,7 +280,7 @@ public void testExactVectorValuesAfterCompaction() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("simple_compact_out"); compactor.compact(outPath); @@ -355,7 +357,7 @@ public void testExactVectorValuesWithDeletions() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("del_compact_out"); compactor.compact(outPath); @@ -433,7 +435,7 @@ public void testExactVectorValuesWithCustomRemapping() throws Exception { List.of(g0, g1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path outPath = testDirectory.resolve("remap_compact_out"); compactor.compact(outPath); @@ -490,7 +492,7 @@ public void testCompact() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); int topK = 10; // Select query vectors from the dataset @@ -596,10 +598,10 @@ public WorkLimiter.Grant acquire(long amount) throws InterruptedException { /** Loads the fixture's source graphs with every node live and identity remapping. */ private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss) throws IOException { - return newAllLiveCompactor(rss, null); + return newAllLiveCompactor(rss, (Executor) null, -1); } - private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss, ForkJoinPool executor) throws IOException { + private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss, Executor executor, int taskWindowSize) throws IOException { List graphs = new ArrayList<>(); List liveNodes = new ArrayList<>(); List remappers = new ArrayList<>(); @@ -617,19 +619,76 @@ private OnDiskGraphIndexCompactor newAllLiveCompactor(List rss, lives.set(0, numVectorsPerGraph); liveNodes.add(lives); } - return new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, executor); + return new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, executor, taskWindowSize); } @Test - public void testGetTaskWindowSizeReflectsInjectedParallelism() throws Exception { - List rss = new ArrayList<>(); + public void testTaskWindowSizeResolution() throws Exception { + List rss1 = new ArrayList<>(); + try { + assertEquals("explicit window must be honored", + 5, newAllLiveCompactor(rss1, (Executor) Runnable::run, 5).getTaskWindowSize()); + } finally { for (var rs : rss1) rs.close(); } + + List rss2 = new ArrayList<>(); + try { + assertEquals("non-FJP executor with window<=0 must default to 1 (serial)", + 1, newAllLiveCompactor(rss2, (Executor) Runnable::run, 0).getTaskWindowSize()); + } finally { for (var rs : rss2) rs.close(); } + + List rss3 = new ArrayList<>(); ForkJoinPool pool = new ForkJoinPool(3); try { - var compactor = newAllLiveCompactor(rss, pool); - assertEquals("taskWindowSize should equal the injected pool's parallelism", - 3, compactor.getTaskWindowSize()); + assertEquals("window<=0 with a ForkJoinPool must derive getParallelism()", + 3, newAllLiveCompactor(rss3, pool, -1).getTaskWindowSize()); + } finally { pool.shutdown(); for (var rs : rss3) rs.close(); } + } + + @Test + public void testCallerRunsProducesValidGraph() throws Exception { + // Caller-runs + serial must produce a correct, searchable graph of equivalent quality to + // the ForkJoinPool path. (Not asserted byte-identical: fused-PQ retraining uses parallel + // floating-point reduction, which is order-sensitive, so serialization can shift low bits.) + List rss = new ArrayList<>(); + try { + var path = testDirectory.resolve("cr_caller"); + newAllLiveCompactor(rss, (Executor) Runnable::run, 1).compact(path); + + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + var g = OnDiskGraphIndex.load(out); + assertEquals(numSources * numVectorsPerGraph, g.size(0)); + + List> queries = new ArrayList<>(); + for (int i = 0; i < numQueries; ++i) queries.add(allVecs.get(randomIntBetween(0, allVecs.size() - 1))); + List> gt = buildGT(queries, 10); + GraphSearcher searcher = new GraphSearcher(g); + List results = new ArrayList<>(); + for (VectorFloat q : queries) { + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv); + results.add(searcher.search(ssp, 10, Bits.ALL)); + } + double recall = AccuracyMetrics.recallFromSearchResults(gt, results, 10, 10); + assertTrue("caller-runs graph recall should be reasonable, got " + recall, recall >= 0.2); + searcher.close(); + } + } finally { + for (var rs : rss) rs.close(); + } + } + + @Test + public void testCallerRunsRunsOnCallingThreadOnly() throws Exception { + List rss = new ArrayList<>(); + Set taskThreads = Collections.synchronizedSet(new HashSet<>()); + Executor recordingCallerRuns = command -> { + taskThreads.add(Thread.currentThread()); + command.run(); + }; + try { + newAllLiveCompactor(rss, recordingCallerRuns, 1).compact(testDirectory.resolve("cr_thread")); + assertEquals("all batch + pre-encode work must run on the calling thread only", + Collections.singleton(Thread.currentThread()), taskThreads); } finally { - pool.shutdown(); for (var rs : rss) rs.close(); } } @@ -735,6 +794,81 @@ public void testCompactAtNonZeroStartOffset() throws Exception { for (var rs : rss) rs.close(); } + @Test + public void testCompactToFileDestination() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var path = testDirectory.resolve("dest_tofile"); + + long bodyLength = compactor.compact(CompactionDestination.toFile(path)); + assertEquals("returned body length must equal the standalone file size", + Files.size(path), bodyLength); + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + assertEquals(numSources * numVectorsPerGraph, OnDiskGraphIndex.load(out).size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactToContainerDestinationCommitsAndPreservesPrefix() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + var container = testDirectory.resolve("dest_container"); + byte[] prefix = new byte[32]; + for (int i = 0; i < prefix.length; i++) prefix[i] = (byte) (0xC0 + (i % 16)); + Files.write(container, prefix); + final long base = prefix.length; + + final long[] committed = { -1 }; + final boolean[] closed = { false }; + CompactionDestination dest = () -> new CompactionDestination.Target() { + public Path file() { return container; } + public long startOffset() { return base; } + public void commit(long bodyLength) { committed[0] = bodyLength; } + public void close() { closed[0] = true; } + }; + + long returned = compactor.compact(dest); + assertTrue("commit must have been called", committed[0] >= 0); + assertEquals("commit bodyLength must equal the returned value", committed[0], returned); + assertTrue("target must be closed", closed[0]); + assertEquals("body length must be file size minus base", Files.size(container) - base, returned); + + byte[] afterPrefix = Arrays.copyOf(Files.readAllBytes(container), prefix.length); + assertArrayEquals("the reserved prefix must be preserved", prefix, afterPrefix); + try (ReaderSupplier out = ReaderSupplierFactory.open(container)) { + assertEquals(numSources * numVectorsPerGraph, OnDiskGraphIndex.load(out, base).size(0)); + } + for (var rs : rss) rs.close(); + } + + @Test + public void testCompactToDestinationAbortsWithoutCommitOnFailure() throws Exception { + List rss = new ArrayList<>(); + var compactor = newAllLiveCompactor(rss); + // Parent directory does not exist, so the graph writer fails partway. + final var badFile = testDirectory.resolve("no_such_dir").resolve("graph"); + + final boolean[] committed = { false }; + final boolean[] closed = { false }; + CompactionDestination dest = () -> new CompactionDestination.Target() { + public Path file() { return badFile; } + public long startOffset() { return 0; } + public void commit(long bodyLength) { committed[0] = true; } + public void close() { closed[0] = true; } + }; + + try { + compactor.compact(dest); + fail("expected compaction to fail for an unwritable destination"); + } catch (Exception expected) { + // ok + } + assertFalse("commit must NOT be called on failure", committed[0]); + assertTrue("close (abort) must run on failure", closed[0]); + for (var rs : rss) rs.close(); + } + /** * Tests compaction with deleted nodes. * Verifies that deleted nodes are properly excluded from the compacted graph. @@ -778,7 +912,7 @@ public void testCompactWithDeletions() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_with_deletions"); compactor.compact(outputPath); @@ -845,7 +979,7 @@ public void testOrdinalMapping() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_with_ordinal_mapping"); compactor.compact(outputPath); @@ -923,7 +1057,7 @@ public void testDeletionsAndOrdinalMapping() throws Exception { liveNodes.add(lives); } - var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null); + var compactor = new OnDiskGraphIndexCompactor(graphs, liveNodes, remappers, similarityFunction, null, -1); var outputPath = testDirectory.resolve("test_compact_deletions_and_mapping"); compactor.compact(outputPath); @@ -1029,7 +1163,7 @@ public void testCompactWithCompressedSidecar() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("sidecar_graph_out"); Path compressedOut = testDirectory.resolve("sidecar_pq_out"); @@ -1112,7 +1246,7 @@ public void testCompactCompressedSidecarRejectsFusedPQ() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - similarityFunction, null); + similarityFunction, null, -1); org.junit.Assert.fail("expected IllegalArgumentException for FUSED_PQ + sourceCompressed"); } catch (IllegalArgumentException expected) { assertTrue("error message mentions FUSED_PQ", @@ -1155,7 +1289,7 @@ public void testCompactCompressedSidecarRejectsSizeMismatch() throws Exception { List.of(pqv0), // size 1 vs sources size 2 List.of(live, live), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); org.junit.Assert.fail("expected IllegalArgumentException for size mismatch"); } catch (IllegalArgumentException expected) { assertTrue("error message mentions size", @@ -1191,7 +1325,7 @@ public void testCompactTwoArgRequiresSourceCompressed() throws Exception { List.of(g0, g1), List.of(live, live), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("noarg_graph_out"); Path compressedOut = testDirectory.resolve("noarg_pq_out"); @@ -1251,7 +1385,7 @@ public void testCompactCompressedSidecarWithDeletions() throws Exception { List.of(pqv0, pqv1), List.of(live0, live1), List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), - vsf, null); + vsf, null, -1); Path graphOut = testDirectory.resolve("delsidecar_graph_out"); Path compressedOut = testDirectory.resolve("delsidecar_pq_out"); From c1b7fb9b163f7cca2140ed87acac25de1ae417ff Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 6 Jul 2026 22:55:30 +0000 Subject: [PATCH 07/13] compaction memory safety: drain on unwind, truncate reused outputs, bound record reads --- .../jbellis/jvector/disk/ReaderSupplier.java | 17 + .../jvector/disk/SimpleMappedReader.java | 9 + .../graph/disk/FusedCompactionStrategy.java | 9 + .../jvector/graph/disk/OnDiskGraphIndex.java | 45 ++- .../graph/disk/OnDiskGraphIndexCompactor.java | 304 +++++++++++++++--- 5 files changed, 334 insertions(+), 50 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java index 7ebb3f9b0..8827bcdb1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java @@ -41,6 +41,23 @@ public interface ReaderSupplier extends AutoCloseable { default void prefetch(long offset, long length) { } + /** + * Releases the supplier's underlying resource. Two implementation families exist, with very + * different safety under concurrency: + *

    + *
  • Coordinated (e.g. the jvector-native {@code MemorySegmentReader.Supplier}, whose + * shared-Arena close performs a liveness handshake): closing while vended readers are still + * in use degrades to {@code IllegalStateException} on those readers.
    • + *
  • Raw-release (e.g. {@link SimpleMappedReader.Supplier}, which unmaps + * immediately): closing while any vended reader is mid-read invalidates the mapped pages + * underneath it, and the JVM fails with a native fault (SIGSEGV) rather than an + * exception.
    • + *
+ * Callers must not close a supplier until every reader vended by {@link #get()} is provably + * quiescent; implementations should document which family they belong to. + * + * @throws IOException if an I/O error occurs + */ default void close() throws IOException { } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleMappedReader.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleMappedReader.java index 46d91f8e3..142d6af6c 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleMappedReader.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/SimpleMappedReader.java @@ -84,6 +84,15 @@ public SimpleMappedReader get() { return new SimpleMappedReader((MappedByteBuffer) buffer.duplicate()); } + /** + * Unmaps the shared mapping immediately (via {@code Unsafe.invokeCleaner}), with + * no coordination with outstanding readers — the raw-release family of + * {@link ReaderSupplier#close()}. Any reader vended by {@link #get()} that touches the + * mapping after this call faults natively (SIGSEGV) rather than throwing an exception, + * so close only once every vended reader is provably done. Where JDK 22+ is available, + * prefer the jvector-native {@code MemorySegmentReader}, whose close degrades to + * {@code IllegalStateException} instead. + */ @Override public void close() { if (unsafe != null) { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java index b6a44ece3..597ab2211 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java @@ -150,6 +150,15 @@ public void onAfterHeader(CompactWriter writer) throws IOException { writer.enablePqCodeCache(codeCache, cacheCodeSize); } } catch (IOException e) { + // The fallback exists for environmental failures (mapping limits, transient IO) — + // per-write encoding produces the same output, just slower. Interruption is not + // environmental: it is the caller cancelling the compaction, and absorbing it here + // would make compact() run to completion after the one interrupt a host delivers. + for (Throwable cause = e; cause != null; cause = cause.getCause()) { + if (cause instanceof InterruptedException) { + throw e; + } + } log.warn("Code pre-encode failed, falling back to per-write encoding: {}", e.getMessage()); } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index 61d8157fb..8a8da3598 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java @@ -224,7 +224,10 @@ private Int2ObjectHashMap loadInMemoryFeatures(Random /** * Load an index from the given reader supplier where header and graph are located on the same file, - * where the index starts at `offset`. + * where the index starts at `offset`. Equivalent to {@code load(readerSupplier, offset, true)}; + * for v5+ graphs the metadata is located via the footer — see the + * {@link #load(ReaderSupplier, long, boolean)} warning about suppliers whose range extends + * past the graph's end. * * @param readerSupplier the reader supplier to use to read the graph and index. * @param offset the offset in bytes from the start of the file where the index starts. @@ -236,6 +239,16 @@ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier, long offset) /** * Load an index from the given reader supplier where header and graph are located on the same file, * where the index starts at `offset`. + *

+ * Footer loading trusts the end of the supplier's range. With {@code useFooter=true} + * and a v5+ graph, metadata is located relative to the reader's {@code length()} — the file + * end, for whole-file suppliers. That is only correct when the graph is the last + * content in the supplier's range. Never footer-load a reused or embedder-owned container + * whose length extends past the graph body: stale bytes there either fail the load loudly + * or, if they end in a stale-but-still-valid footer, silently resurrect the old graph over + * the new bytes. For such containers, pass {@code useFooter=false} with the known + * {@code offset}, or use a region-bounded {@link ReaderSupplier} whose {@code length()} is + * the end of the graph's region. * * @param readerSupplier the reader supplier to use to read the graph and index. * @param offset the offset in bytes from the start of the file where the index starts. @@ -267,6 +280,9 @@ public static OnDiskGraphIndex load(ReaderSupplier readerSupplier, long offset, /** * Load an index from the given reader supplier where header and graph are located on the same file at offset 0. + * For v5+ graphs, metadata is located via the footer at the end of the supplier's range — + * the supplier must contain the graph and nothing after it; see the + * {@link #load(ReaderSupplier, long, boolean)} warning. * * @param readerSupplier the reader supplier to use to read the graph index. */ @@ -491,6 +507,21 @@ public View(RandomAccessReader reader) { this.neighbors = new int[layerInfo.stream().mapToInt(li -> li.degree).max().orElse(0)]; } + /** + * Guards every on-disk record access. A node ordinal outside the L0 record space + * ({@code idUpperBound}, which exceeds {@code size(0)} for graphs renumbered with holes) + * would otherwise become a silent wild offset into the mapped file — reading garbage (or + * faulting) instead of failing diagnosably. Both package-private offset entry points call + * this, so each access is validated exactly once, with a single branch against a final + * bound. + */ + private void requireValidNode(int node) { + if (node < 0 || node >= idUpperBound) { + throw new IllegalArgumentException( + "node ordinal " + node + " out of range [0, " + idUpperBound + ") for this graph"); + } + } + @Override public int dimension() { return dimension; @@ -509,6 +540,7 @@ public RandomAccessVectorValues copy() { // package-private: OnDiskGraphIndexCompactor uses this for in-place neighbor refinement long offsetFor(int node, FeatureId featureId) { + requireValidNode(node); Feature feature = features.get(featureId); // Separated features are just global offset + node offset @@ -525,6 +557,7 @@ long offsetFor(int node, FeatureId featureId) { // package-private: OnDiskGraphIndexCompactor uses this for in-place neighbor refinement long neighborsOffsetFor(int level, int node) { + requireValidNode(node); assert level == 0; // higher layers are in memory // skip node ID + inline features @@ -585,8 +618,14 @@ public NodesIterator getNeighborsIterator(int level, int node) { // For layer 0, read from disk reader.seek(neighborsOffsetFor(level, node)); nodeDegree = reader.readInt(); - assert nodeDegree <= neighbors.length - : String.format("Node %d neighborCount %d > M %d", node, nodeDegree, neighbors.length); + if (nodeDegree < 0 || nodeDegree > neighbors.length) { + // A real check, not an assert: an out-of-range on-disk degree means the + // block is corrupt or the metadata is stale, and the garbage ints that a + // blind read would yield become out-of-range node ids downstream. + throw new IllegalStateException(String.format( + "Corrupt neighbor block: node %d at level 0 declares degree %d outside [0, %d] (block offset %d)", + node, nodeDegree, neighbors.length, neighborsOffsetFor(level, node))); + } reader.read(neighbors, 0, nodeDegree); stored = neighbors; } else { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index 95bea2a9b..8a1636436 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -54,6 +54,22 @@ import static java.lang.Math.*; +/** + * Merges multiple {@link OnDiskGraphIndex} sources into a single compacted index, preserving the + * layer hierarchy, remapping ordinals, and — when quantization features are present — retraining + * and re-encoding codes. Entry points are the {@code compact(...)} overloads; parallelism and + * throttling are supplied by the caller via the executor constructor argument and + * {@link #setProgressLimiter}. + * + *

Source lifecycle contract. The caller owns every source's + * {@link io.github.jbellis.jvector.disk.ReaderSupplier} and must keep the suppliers — and the + * files and mappings underneath them — open and unchanged until the {@code compact(...)} call + * returns or throws. That is sufficient: on every path, including batch failure and interrupt, + * the compactor drains its in-flight work before the call unwinds, so no source access survives + * it. Closing, truncating, or rewriting a source before then is undefined behavior; for + * raw-unmapping suppliers (see {@link io.github.jbellis.jvector.disk.ReaderSupplier#close()}) it + * manifests as a native JVM fault rather than an exception.

+ */ public final class OnDiskGraphIndexCompactor implements Accountable { private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); private static final Logger log = LoggerFactory.getLogger(OnDiskGraphIndexCompactor.class); @@ -237,24 +253,112 @@ private static int resolveWindow(Executor executor, int requested) { /** * Adapts an arbitrary {@link Executor} to the {@link ExecutorService} the quantization - * strategies need for {@code invokeAll} during PQ pre-encode. A real {@code ExecutorService} - * (including a {@link ForkJoinPool}) is used directly; a plain {@code Executor} — notably a - * caller-runs {@code Runnable::run} — is wrapped so its tasks run on whatever thread the - * executor dispatches to (the calling thread, for caller-runs). No pool is created, and the - * lifecycle methods are inert since the compactor never owns the executor. + * strategies need for {@code invokeAll} during PQ pre-encode. Tasks run on whatever thread + * the executor dispatches to (the calling thread, for a caller-runs {@code Runnable::run}); + * no pool is created, and the lifecycle methods are inert since the compactor never owns the + * executor. Always returns the {@link DrainingExecutorService} wrapper — even over a real + * {@code ExecutorService} — because the wrapper's {@code invokeAll} must own the unwind + * semantics; see there. */ private static ExecutorService asExecutorService(Executor executor) { - if (executor instanceof ExecutorService) { - return (ExecutorService) executor; - } - return new AbstractExecutorService() { - @Override public void execute(Runnable command) { executor.execute(command); } - @Override public void shutdown() { } - @Override public List shutdownNow() { return Collections.emptyList(); } - @Override public boolean isShutdown() { return false; } - @Override public boolean isTerminated() { return false; } - @Override public boolean awaitTermination(long timeout, TimeUnit unit) { return false; } - }; + return new DrainingExecutorService(executor); + } + + /** + * The adapter behind {@link #asExecutorService}. Its {@code invokeAll} carries the same + * drain-before-unwind guarantee as the batch loops ({@code awaitAbandoned}): on interrupt, + * every submitted task is awaited before the {@code InterruptedException} propagates — so no + * strategy fan-out can outlive {@code compact()} and keep reading source graphs (or writing + * the pre-encode cache that {@code onAfterClose} unmaps) after the caller regains control. + * Stock {@code invokeAll} instead cancels-with-interrupt and returns immediately, abandoning + * running tasks, which do not poll the interrupt flag. + */ + private static final class DrainingExecutorService extends AbstractExecutorService { + private final Executor executor; + + DrainingExecutorService(Executor executor) { + this.executor = executor; + } + + @Override + public void execute(Runnable command) { + executor.execute(command); + } + + @Override + public List> invokeAll(Collection> tasks) throws InterruptedException { + List> futures = new ArrayList<>(tasks.size()); + try { + for (Callable task : tasks) { + RunnableFuture f = newTaskFor(task); + futures.add(f); + executor.execute(f); + } + } catch (Throwable t) { + if (awaitDone(futures)) { + Thread.currentThread().interrupt(); + } + throw t; + } + if (awaitDone(futures)) { + // Interrupt status is consumed by the exception, matching stock invokeAll; a set + // flag would also make the unwind's own channel work (e.g. the code-cache + // truncate in onAfterClose) fail with ClosedByInterruptException. + throw new InterruptedException("interrupted during invokeAll; submitted tasks were drained first"); + } + return futures; + } + + /** + * Blocks until every future is done — uninterruptibly, the same hang-beats-crash choice + * as {@code awaitAbandoned} — swallowing per-task outcomes (callers inspect the futures, + * matching the {@code invokeAll} contract). Nothing is cancelled, deliberately: + * {@code FutureTask.cancel} cannot distinguish queued from running (a running task's + * state is still NEW), so cancelling would detach a future from its still-running body + * and recreate the abandonment this wrapper exists to prevent. Queued tasks therefore + * run to completion on the caller's executor before the unwind proceeds. Returns whether + * an interrupt was received while waiting. + */ + private static boolean awaitDone(List> futures) { + boolean interrupted = Thread.interrupted(); + for (Future f : futures) { + while (!f.isDone()) { + try { + f.get(); + } catch (InterruptedException e) { + interrupted = true; + } catch (ExecutionException | CancellationException ignored) { + // the outcome stays in the future for the caller + } + } + } + return interrupted; + } + + // lifecycle is inert: the compactor never owns the executor + @Override + public void shutdown() { + } + + @Override + public List shutdownNow() { + return Collections.emptyList(); + } + + @Override + public boolean isShutdown() { + return false; + } + + @Override + public boolean isTerminated() { + return false; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) { + return false; + } } /** @@ -393,6 +497,12 @@ private void validateFeatures(List sources) { /** * Main compaction entry point. Merges all source indexes into a single output index at the * specified path, handling PQ retraining if needed, and writing header, all layers, and footer. + * Any pre-existing file at {@code outputPath} is truncated first: the destination is wholly + * jvector-owned, and a leftover longer file would keep its old footer at the file end, which + * the footer-based default load would otherwise silently trust. + *

+ * Source lifecycle: the caller keeps every source {@code ReaderSupplier} open until this call + * returns or throws — which is sufficient; see the class-level contract. */ @Experimental public void compact(Path outputPath) throws FileNotFoundException { @@ -405,8 +515,21 @@ public void compact(Path outputPath) throws FileNotFoundException { * that wraps the graph in its own container reserves its header by passing the header size as * {@code startOffset}, so jvector's body lands directly inside the container — removing the * temp-file-and-copy. jvector writes only {@code [startOffset, projectedSize)} and never reads - * or clobbers the reserved prefix. The file is opened read/write and not truncated, so a - * prefix the embedder pre-wrote survives. {@code compact(path, 0)} equals {@link #compact(Path)}. + * or clobbers the reserved prefix. + *

+ * With {@code startOffset > 0} the file is opened read/write and not truncated, so a + * prefix the embedder pre-wrote survives — and therefore a container longer than the new body + * keeps its stale tail. Never load such a container footer-first through a whole-file reader + * (the footer search trusts the file end): use + * {@code OnDiskGraphIndex.load(supplier, startOffset, false)}, or a region-bounded + * {@code ReaderSupplier} whose {@code length()} is the region end. With + * {@code startOffset == 0} the destination is wholly jvector-owned and any pre-existing file + * is truncated before writing, so a reused path cannot leave a stale tail — or a stale, + * still-valid footer — behind the new graph. {@code compact(path, 0)} equals + * {@link #compact(Path)}. + *

+ * Source lifecycle: the caller keeps every source {@code ReaderSupplier} open until this call + * returns or throws — which is sufficient; see the class-level contract. * * @param outputPath the file to write into * @param startOffset the byte offset at which jvector's output begins; {@code 0} for a @@ -417,6 +540,9 @@ public void compact(Path outputPath, long startOffset) throws FileNotFoundExcept if (startOffset < 0) { throw new IllegalArgumentException("startOffset must be >= 0, got " + startOffset); } + if (startOffset == 0) { + truncateStandaloneDestination(outputPath); + } QuantizationCompactionStrategy strategy = detectInlineStrategy(); try { compactGraphImpl(outputPath, startOffset, strategy); @@ -438,6 +564,9 @@ public void compact(Path outputPath, long startOffset) throws FileNotFoundExcept * closes the target without committing (so the embedder discards the partial output). The * compactor writes into {@code target.file()} at {@code target.startOffset()} using its own * random-access + memory-mapped IO; the destination expresses where, not how. + *

+ * Source lifecycle: the caller keeps every source {@code ReaderSupplier} open until this call + * returns or throws — which is sufficient; see the class-level contract. * * @return the graph body length written, i.e. {@code size(file) - startOffset} */ @@ -456,11 +585,16 @@ public long compact(CompactionDestination destination) throws IOException { /** * Compaction entry point for graphs that ship a non-fused compressed sidecar (e.g. * {@link io.github.jbellis.jvector.quantization.PQVectors}). Writes the merged graph to - * {@code graphPath} and the merged compressed vectors to {@code compressedPath}. + * {@code graphPath} and the merged compressed vectors to {@code compressedPath}. Both are + * wholly jvector-owned standalone destinations: any pre-existing file at either path is + * truncated first. *

* The compressor is retrained on a balanced sample of merged source vectors, then every live * node is re-encoded against the new codebook. Requires that {@code sourceCompressed} was * supplied to the constructor. + *

+ * Source lifecycle: the caller keeps every source {@code ReaderSupplier} open until this call + * returns or throws — which is sufficient; see the class-level contract. */ @Experimental public void compact(Path graphPath, Path compressedPath) throws FileNotFoundException { @@ -470,6 +604,12 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce } Objects.requireNonNull(compressedPath, "compressedPath"); + // Both outputs are wholly jvector-owned standalone files; clear stale content the same + // way compact(Path) does. The graph reload is footer-based and provably corruptible by a + // stale tail, and the sidecar writer likewise opens "rw" without truncating. + truncateStandaloneDestination(graphPath); + truncateStandaloneDestination(compressedPath); + // Graph compaction proceeds without fused-PQ retrain (validateCompressed forbids // FUSED_PQ when sourceCompressed is set), then the sidecar is written below. QuantizationCompactionStrategy inlineStrategy = detectInlineStrategy(); @@ -488,6 +628,21 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce } } + /** + * Clears any pre-existing file at a wholly jvector-owned (standalone) destination so its + * stale tail cannot survive past this compaction. Footer-based loads locate metadata + * relative to the file END, so a stale-but-still-valid footer there would silently resurrect + * the old graph's structure over the new bytes. Never applied to embedded destinations + * ({@code startOffset > 0}), whose containers belong to the embedder. + */ + private static void truncateStandaloneDestination(Path destination) { + try (FileChannel ch = FileChannel.open(destination, StandardOpenOption.WRITE, StandardOpenOption.CREATE)) { + ch.truncate(0); + } catch (IOException e) { + throw new UncheckedIOException("Failed to truncate pre-existing standalone destination " + destination, e); + } + } + /** * For compaction use. Drops the compactor's strong references to the source graphs and their * per-source live-node / remapper sidecars, and tells the strategy to release its @@ -701,9 +856,10 @@ private void refineCompactedGraph(Path outputPath, long startOffset, Quantizatio }); inFlight++; } - nodesDone += ecs.take().get(); + Future finished = ecs.take(); + inFlight--; // finished, whatever its outcome — get() below may throw + nodesDone += finished.get(); completed++; - inFlight--; WorkLimiter.Grant g = grants.poll(); if (g != null) g.close(); limiter.onProgress(Phase.REFINE, nodesDone, total); @@ -712,6 +868,13 @@ private void refineCompactedGraph(Path outputPath, long startOffset, Quantizatio nextProgress += progressStep; } } + } catch (Throwable t) { + // Same drain-before-unwind as runBatchesWithBackpressure: the enclosing + // try-with-resources releases the output supplier, the channel, and (for fused + // strategies) the pre-encode cache as this exception propagates, which must not + // happen underneath still-running refine batches. + awaitAbandoned(ecs, inFlight); + throw t; } finally { for (WorkLimiter.Grant g : grants) g.close(); } @@ -1496,40 +1659,87 @@ private void runBatchesWithBackpressure( final int total = batches.size(); int nextToSubmit = 0; int inFlight = 0; - - // initial window - while (inFlight < taskWindowSize && nextToSubmit < total) { - submitOne.accept(batches.get(nextToSubmit++)); - inFlight++; - } - int completed = 0; - while (completed < total) { - List results = ecs.take().get(); - - // Admission runs on this (orchestrating) thread, before the write. A blocking limiter - // back-pressures dispatch/consume while in-flight workers keep computing — no - // ManagedBlocker. The amount is the bytes this batch is about to write, read here before - // onComplete consumes the buffers. For the default UNLIMITED limiter both calls are no-ops. - long amount = batchBytes.applyAsLong(results); - try (WorkLimiter.Grant g = limiter.acquire(amount)) { - onComplete.accept(results); + + try { + // initial window + while (inFlight < taskWindowSize && nextToSubmit < total) { + submitOne.accept(batches.get(nextToSubmit++)); + inFlight++; } - completed++; - inFlight--; + while (completed < total) { + Future> finished = ecs.take(); + inFlight--; // finished, whatever its outcome — get() below may throw + List results = finished.get(); + + // Admission runs on this (orchestrating) thread, before the write. A blocking limiter + // back-pressures dispatch/consume while in-flight workers keep computing — no + // ManagedBlocker. The amount is the bytes this batch is about to write, read here before + // onComplete consumes the buffers. For the default UNLIMITED limiter both calls are no-ops. + long amount = batchBytes.applyAsLong(results); + try (WorkLimiter.Grant g = limiter.acquire(amount)) { + onComplete.accept(results); + } - progress[0] = Math.min(progress[0] + results.size(), progress[1]); - limiter.onProgress(stage, progress[0], progress[1]); + completed++; - if (nextToSubmit < total) { - submitOne.accept(batches.get(nextToSubmit++)); - inFlight++; + progress[0] = Math.min(progress[0] + results.size(), progress[1]); + limiter.onProgress(stage, progress[0], progress[1]); + + if (nextToSubmit < total) { + submitOne.accept(batches.get(nextToSubmit++)); + inFlight++; + } + if (completed % 10 == 0) { + log.debug("Compaction I/O progress: {}/{} batches written to disk", completed, total); + } } - if (completed % 10 == 0) { - log.debug("Compaction I/O progress: {}/{} batches written to disk", completed, total); + } catch (Throwable t) { + // A failed batch or an interrupt must not unwind while other batches are still + // running: the caller may release the source graphs the moment compact() returns or + // throws, and a read abandoned mid-flight against a since-unmapped source faults + // natively instead of throwing. Block until every submitted batch has finished, then + // let the original failure propagate. + awaitAbandoned(ecs, inFlight); + throw t; + } + } + + /** + * Blocks until {@code remaining} already-submitted batches have finished, discarding their + * results. Runs on the orchestrating thread during an exceptional unwind, before the + * exception can reach any scope that releases resources the batches still read (the caller's + * source graphs, the output supplier and channel, the pre-encode cache): a batch abandoned + * mid-read whose mapping is then closed faults natively instead of throwing. The wait is + * deliberately unbounded — a hung batch leaves a thread-dumpable hang, which is strictly + * better than the use-after-unmap crash a timeout would reintroduce; batches never wait on + * the orchestrator, so the drain cannot deadlock. An interrupt received while draining is + * remembered and re-asserted on exit, never dropped. + */ + private static void awaitAbandoned(ExecutorCompletionService ecs, int remaining) { + boolean interrupted = Thread.interrupted(); + int drained = 0; + while (drained < remaining) { + try { + Future finished = ecs.take(); + drained++; + try { + finished.get(); + } catch (ExecutionException e) { + log.debug("Discarding abandoned batch failure during compaction unwind", e.getCause()); + } catch (InterruptedException e) { + interrupted = true; + } catch (CancellationException e) { + // nothing cancels these futures today; tolerate rather than mask the unwind + } + } catch (InterruptedException e) { + interrupted = true; } } + if (interrupted) { + Thread.currentThread().interrupt(); + } } /** From f59ac3140246f7a40ab1fd8389f39de97f1477f0 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Tue, 7 Jul 2026 22:53:20 +0000 Subject: [PATCH 08/13] GraphIndexBuilder: caller-runs build/finalize via ParallelExecutor Lets an embedder run graph build/cleanup on the calling thread instead of a ForkJoinPool. Existing pool constructors are preserved as delegating overloads. --- .../jvector/graph/GraphIndexBuilder.java | 258 +++++++++++------ .../jvector/graph/ParallelExecutor.java | 118 ++++++++ .../GraphIndexBuilderCallerRunsTest.java | 273 ++++++++++++++++++ 3 files changed, 559 insertions(+), 90 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderCallerRunsTest.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 8135bba25..d5a492861 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 @@ -74,8 +74,8 @@ public class GraphIndexBuilder implements Closeable, Accountable { private final BuildScoreProvider scoreProvider; - private final ForkJoinPool simdExecutor; - private final ForkJoinPool parallelExecutor; + private final ParallelExecutor simdExecutor; + private final ParallelExecutor parallelExecutor; private final ExplicitThreadLocal searchers; @@ -237,6 +237,39 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, this(scoreProvider, dimension, List.of(M), beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, simdExecutor, parallelExecutor); } + /** + * Reads all the vectors from vector values, builds a graph connecting them by their dense + * ordinals, using the given hyperparameter settings, and returns the resulting graph. + * + * @param scoreProvider describes how to determine the similarities between vectors + * @param M the maximum number of connections a node can have + * @param beamWidth the size of the beam search to use when finding nearest neighbors. + * @param neighborOverflow the ratio of extra neighbors to allow temporarily when inserting a + * node. larger values will build more efficiently, but use more memory. + * @param alpha how aggressive pruning diverse neighbors should be. Set alpha > 1.0 to + * allow longer edges. If alpha = 1.0 then the equivalent of the lowest level of + * an HNSW graph will be created, which is usually not what you want. + * @param addHierarchy whether we want to add an HNSW-style hierarchy on top of the Vamana index. + * @param refineFinalGraph whether we do a second pass over each node in the graph to refine its connections + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + public GraphIndexBuilder(BuildScoreProvider scoreProvider, + int dimension, + int M, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) + { + this(scoreProvider, dimension, List.of(M), beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, simdExecutor, parallelExecutor); + } + /** * Reads all the vectors from vector values, builds a graph connecting them by their dense * ordinals, using the given hyperparameter settings, and returns the resulting graph. @@ -295,6 +328,41 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, boolean refineFinalGraph, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + this(scoreProvider, dimension, maxDegrees, beamWidth, neighborOverflow, alpha, addHierarchy, refineFinalGraph, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * Reads all the vectors from vector values, builds a graph connecting them by their dense + * ordinals, using the given hyperparameter settings, and returns the resulting graph. + * + * @param scoreProvider describes how to determine the similarities between vectors + * @param maxDegrees the maximum number of connections a node can have in each layer; if fewer entries + * are specified than the number of layers, the last entry is used for all remaining layers. + * @param beamWidth the size of the beam search to use when finding nearest neighbors. + * @param neighborOverflow the ratio of extra neighbors to allow temporarily when inserting a + * node. larger values will build more efficiently, but use more memory. + * @param alpha how aggressive pruning diverse neighbors should be. Set alpha > 1.0 to + * allow longer edges. If alpha = 1.0 then the equivalent of the lowest level of + * an HNSW graph will be created, which is usually not what you want. + * @param addHierarchy whether we want to add an HNSW-style hierarchy on top of the Vamana index. + * @param refineFinalGraph whether we do a second pass over each node in the graph to refine its connections + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + public GraphIndexBuilder(BuildScoreProvider scoreProvider, + int dimension, + List maxDegrees, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { if (maxDegrees.stream().anyMatch(i -> i <= 0)) { throw new IllegalArgumentException("layer degrees must be positive"); @@ -352,6 +420,27 @@ public GraphIndexBuilder(BuildScoreProvider scoreProvider, */ @Experimental public GraphIndexBuilder(BuildScoreProvider buildScoreProvider, int dimension, MutableGraphIndex mutableGraphIndex, int beamWidth, float neighborOverflow, float alpha, boolean refineFinalGraph, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { + this(buildScoreProvider, dimension, mutableGraphIndex, beamWidth, neighborOverflow, alpha, refineFinalGraph, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * Create this builder from an existing {@link io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex}, this is useful when we just loaded a graph from disk + * copy it into {@link OnHeapGraphIndex} and then start mutating it with minimal overhead of recreating the mutable {@link OnHeapGraphIndex} used in the new GraphIndexBuilder object + * + * @param buildScoreProvider the provider responsible for calculating build scores. + * @param mutableGraphIndex a mutable graph index. + * @param beamWidth the width of the beam used during the graph building process. + * @param neighborOverflow the factor determining how many additional neighbors are allowed beyond the configured limit. + * @param alpha the weight factor for balancing score computations. + * @param refineFinalGraph whether to perform a refinement step on the final graph structure. + * @param simdExecutor runs SIMD-heavy build iterations; use {@link ParallelExecutor#callerRuns()} + * to run them synchronously on the calling thread with no worker threads. + * @param parallelExecutor runs the parallel-stream build/cleanup iterations; use + * {@link ParallelExecutor#callerRuns()} for single-threaded, caller-runs finalize. + */ + @Experimental + public GraphIndexBuilder(BuildScoreProvider buildScoreProvider, int dimension, MutableGraphIndex mutableGraphIndex, int beamWidth, float neighborOverflow, float alpha, boolean refineFinalGraph, ParallelExecutor simdExecutor, ParallelExecutor parallelExecutor) { if (beamWidth <= 0) { throw new IllegalArgumentException("beamWidth must be positive"); } @@ -403,29 +492,27 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi var otherView = other.graph.getView(); // Copy each node and its neighbors from the old graph to the new one - other.parallelExecutor.submit(() -> { - IntStream.range(0, other.graph.getIdUpperBound()).parallel().forEach(i -> { - // Find the highest layer this node exists in - int maxLayer = other.graph.getMaxLevelForNode(i); - if (maxLayer < 0) { - return; - } + other.parallelExecutor.forEachInt(other.graph.getIdUpperBound(), i -> { + // Find the highest layer this node exists in + int maxLayer = other.graph.getMaxLevelForNode(i); + if (maxLayer < 0) { + return; + } - // Loop over 0..maxLayer, re-score neighbors for each layer - var sf = newProvider.searchProviderFor(i).scoreFunction(); - for (int lvl = 0; lvl <= maxLayer; lvl++) { - var oldNeighborsIt = otherView.getNeighborsIterator(lvl, i); - // Copy edges, compute new scores - var newNeighbors = new NodeArray(oldNeighborsIt.size()); - while (oldNeighborsIt.hasNext()) { - int neighbor = oldNeighborsIt.nextInt(); - // since we're using a different score provider, use insertSorted instead of addInOrder - newNeighbors.insertSorted(neighbor, sf.similarityTo(neighbor)); - } - newBuilder.graph.connectNode(lvl, i, newNeighbors); + // Loop over 0..maxLayer, re-score neighbors for each layer + var sf = newProvider.searchProviderFor(i).scoreFunction(); + for (int lvl = 0; lvl <= maxLayer; lvl++) { + var oldNeighborsIt = otherView.getNeighborsIterator(lvl, i); + // Copy edges, compute new scores + var newNeighbors = new NodeArray(oldNeighborsIt.size()); + while (oldNeighborsIt.hasNext()) { + int neighbor = oldNeighborsIt.nextInt(); + // since we're using a different score provider, use insertSorted instead of addInOrder + newNeighbors.insertSorted(neighbor, sf.similarityTo(neighbor)); } - }); - }).join(); + newBuilder.graph.connectNode(lvl, i, newNeighbors); + } + }); // Set the entry node newBuilder.graph.updateEntryNode(otherView.entryNode()); @@ -437,11 +524,9 @@ public ImmutableGraphIndex build(RandomAccessVectorValues ravv) { var vv = ravv.threadLocalSupplier(); int size = ravv.size(); - simdExecutor.submit(() -> { - IntStream.range(0, size).parallel().forEach(node -> { - addGraphNode(node, vv.get().getVector(node)); - }); - }).join(); + simdExecutor.forEachInt(size, node -> { + addGraphNode(node, vv.get().getVector(node)); + }); cleanup(); return graph; @@ -463,7 +548,8 @@ void validateEntryNode() { * Cleanup the graph by completing removal of marked-for-delete nodes, trimming * neighbor sets to the advertised degree, and updating the entry node. *

- * Uses default threadpool to process nodes in parallel. There is currently no way to restrict this to a single thread. + * Processes nodes in parallel via the builder's {@code parallelExecutor}. Construct the builder with + * {@link ParallelExecutor#callerRuns()} to run this cleanup synchronously on the calling thread instead. *

* Must be called before writing to disk. *

@@ -490,19 +576,15 @@ public void cleanup() { // It may be helpful for 2D use cases, but empirically it seems unnecessary for high-dimensional vectors. // It may bring a slight improvement in recall for small maximum degrees, // but it can be easily be compensated by using a slightly larger neighborOverflow. - parallelExecutor.submit(() -> { - graph.nodeStream(1).parallel().forEach(this::improveConnections); - }).join(); + parallelExecutor.forEach(graph.nodeStream(1), this::improveConnections); } // clean up overflowed neighbor lists - parallelExecutor.submit(() -> { - IntStream.range(0, graph.getIdUpperBound()).parallel().forEach(id -> { - for (int level = 0; level <= graph.getMaxLevel(); level++) { - graph.enforceDegree(id); - } - }); - }).join(); + parallelExecutor.forEachInt(graph.getIdUpperBound(), id -> { + for (int level = 0; level <= graph.getMaxLevel(); level++) { + graph.enforceDegree(id); + } + }); graph.setAllMutationsCompleted(); } @@ -701,67 +783,63 @@ public synchronized long removeDeletedNodes() { // strategy is proposed in "FreshDiskANN: A Fast and Accurate Graph-Based // ANN Index for Streaming Similarity Search" section 4.2. var newEdges = new ConcurrentHashMap>(); // new edges for key k are values v - parallelExecutor.submit(() -> { - IntStream.range(0, graph.getIdUpperBound()).parallel().forEach(i -> { - if (toDelete.get(i)) { - return; - } - for (var it = graph.getNeighborsIterator(level, i); it.hasNext(); ) { - var j = it.nextInt(); - if (toDelete.get(j)) { - var newEdgesForI = newEdges.computeIfAbsent(i, __ -> ConcurrentHashMap.newKeySet()); - for (var jt = graph.getNeighborsIterator(level, j); jt.hasNext(); ) { - int k = jt.nextInt(); - if (i != k && !toDelete.get(k)) { - newEdgesForI.add(k); - } + parallelExecutor.forEachInt(graph.getIdUpperBound(), i -> { + if (toDelete.get(i)) { + return; + } + for (var it = graph.getNeighborsIterator(level, i); it.hasNext(); ) { + var j = it.nextInt(); + if (toDelete.get(j)) { + var newEdgesForI = newEdges.computeIfAbsent(i, __ -> ConcurrentHashMap.newKeySet()); + for (var jt = graph.getNeighborsIterator(level, j); jt.hasNext(); ) { + int k = jt.nextInt(); + if (i != k && !toDelete.get(k)) { + newEdgesForI.add(k); } } } - }); - }).join(); + } + }); // Remove deleted nodes from neighbors lists; // Score the new edges, and connect the most diverse ones as neighbors - simdExecutor.submit(() -> { - newEdges.entrySet().stream().parallel().forEach(e -> { - // turn the new edges into a NodeArray - int node = e.getKey(); - // each deleted node has ALL of its neighbors added as candidates, so using approximate - // scoring and then re-scoring only the best options later makes sense here - var sf = scoreProvider.searchProviderFor(node).scoreFunction(); - var candidates = new NodeArray(graph.getDegree(level)); - for (var k : e.getValue()) { - candidates.insertSorted(k, sf.similarityTo(k)); - } + simdExecutor.forEach(newEdges.entrySet().stream(), e -> { + // turn the new edges into a NodeArray + int node = e.getKey(); + // each deleted node has ALL of its neighbors added as candidates, so using approximate + // scoring and then re-scoring only the best options later makes sense here + var sf = scoreProvider.searchProviderFor(node).scoreFunction(); + var candidates = new NodeArray(graph.getDegree(level)); + for (var k : e.getValue()) { + candidates.insertSorted(k, sf.similarityTo(k)); + } - // it's unlikely, but possible, that all the potential replacement edges were to nodes that have also - // been deleted. if that happens, keep the graph connected by adding random edges. - // (this is overly conservative -- really what we care about is that the end result of - // replaceDeletedNeighbors not be empty -- but we want to avoid having the node temporarily - // neighborless while concurrent searches run. empirically, this only results in a little extra work.) - if (candidates.size() == 0) { - var R = ThreadLocalRandom.current(); - // doing actual sampling-without-replacement is expensive so we'll loop a fixed number of times instead - for (int i = 0; i < 2 * graph.getDegree(level); i++) { - int randomNode = R.nextInt(graph.getIdUpperBound()); - while (toDelete.get(randomNode)) { - randomNode = R.nextInt(graph.getIdUpperBound()); - } - if (randomNode != node && !candidates.contains(randomNode) && graph.contains(level, randomNode)) { - float score = sf.similarityTo(randomNode); - candidates.insertSorted(randomNode, score); - } - if (candidates.size() == graph.getDegree(level)) { - break; - } + // it's unlikely, but possible, that all the potential replacement edges were to nodes that have also + // been deleted. if that happens, keep the graph connected by adding random edges. + // (this is overly conservative -- really what we care about is that the end result of + // replaceDeletedNeighbors not be empty -- but we want to avoid having the node temporarily + // neighborless while concurrent searches run. empirically, this only results in a little extra work.) + if (candidates.size() == 0) { + var R = ThreadLocalRandom.current(); + // doing actual sampling-without-replacement is expensive so we'll loop a fixed number of times instead + for (int i = 0; i < 2 * graph.getDegree(level); i++) { + int randomNode = R.nextInt(graph.getIdUpperBound()); + while (toDelete.get(randomNode)) { + randomNode = R.nextInt(graph.getIdUpperBound()); + } + if (randomNode != node && !candidates.contains(randomNode) && graph.contains(level, randomNode)) { + float score = sf.similarityTo(randomNode); + candidates.insertSorted(randomNode, score); + } + if (candidates.size() == graph.getDegree(level)) { + break; } } + } - // remove edges to deleted nodes and add the new connections, maintaining diversity - graph.replaceDeletedNeighbors(level, node, toDelete, candidates); - }); - }).join(); + // remove edges to deleted nodes and add the new connections, maintaining diversity + graph.replaceDeletedNeighbors(level, node, toDelete, candidates); + }); } // Generally we want to keep entryPoint update and node removal distinct, because both can be expensive, diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java new file mode 100644 index 000000000..e4ac29d6d --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java @@ -0,0 +1,118 @@ +/* + * 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 java.util.concurrent.ForkJoinPool; +import java.util.function.Consumer; +import java.util.function.IntConsumer; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +/** + * Runs a {@link GraphIndexBuilder}'s internal build/finalize iterations to completion, blocking the + * calling thread until every element has been processed. The implementation decides how the + * iteration is distributed: {@link #forkJoin(ForkJoinPool)} hosts a parallel stream on a dedicated + * pool (the historical behavior), while {@link #callerRuns()} runs everything sequentially on the + * calling thread with no worker threads and no pool. + *

+ * This is the seam that lets an embedder bound vector-graph construction to its own thread budget — + * e.g. one thread per compaction — instead of a jvector-owned all-core pool. It is the build/finalize + * counterpart to the caller-runs executor injection already available on the compaction merge path. + */ +public interface ParallelExecutor { + /** + * Runs {@code body} for each {@code i} in {@code [0, upperBound)}, blocking until all complete. + * + * @param upperBound the exclusive upper bound of the index range (may be {@code 0}) + * @param body the action to apply to each index + */ + void forEachInt(int upperBound, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream of primitive ints to iterate + * @param body the action to apply to each element + */ + void forEach(IntStream source, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream to iterate + * @param body the action to apply to each element + * @param the stream element type + */ + void forEach(Stream source, Consumer body); + + /** + * Returns an executor backed by {@code pool}: each iteration is hosted as a parallel stream on + * that pool and the calling thread blocks on the result. This reproduces the behavior of the + * {@code ForkJoinPool}-based {@link GraphIndexBuilder} constructors. + * + * @param pool the pool that hosts the parallel iterations + * @return a pool-backed {@code ParallelExecutor} + */ + static ParallelExecutor forkJoin(ForkJoinPool pool) { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + pool.submit(() -> IntStream.range(0, upperBound).parallel().forEach(body)).join(); + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + + @Override + public void forEach(Stream source, Consumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + }; + } + + /** + * Returns an executor that runs every iteration sequentially on the calling thread — no worker + * threads, no pool, and the common pool is left untouched. Graph structure and recall are + * equivalent to the {@link #forkJoin(ForkJoinPool)} path; only wall-clock and thread usage differ. + * + * @return a caller-runs {@code ParallelExecutor} + */ + static ParallelExecutor callerRuns() { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + for (int i = 0; i < upperBound; i++) { + body.accept(i); + } + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + source.forEach(body); + } + + @Override + public void forEach(Stream source, Consumer body) { + source.forEach(body); + } + }; + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderCallerRunsTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderCallerRunsTest.java new file mode 100644 index 000000000..e4cd28374 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderCallerRunsTest.java @@ -0,0 +1,273 @@ +/* + * 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 com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.graph.similarity.ScoreFunction; +import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ForkJoinPool; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/// Tests for the {@link ParallelExecutor} caller-runs build/finalize mode added to +/// {@link GraphIndexBuilder}. Verifies that {@link ParallelExecutor#callerRuns()} runs every +/// internal build/cleanup iteration on the calling thread (no worker threads), that it yields a +/// graph of equivalent quality to the {@link java.util.concurrent.ForkJoinPool} path (build, +/// cleanup, and the {@code removeDeletedNodes} path), and that the existing {@code ForkJoinPool} +/// constructors remain intact. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class GraphIndexBuilderCallerRunsTest extends RandomizedTest { + + private static final int DIM = 32; + private static final int SIZE = 2_000; + private static final int M = 16; + private static final int BEAM = 100; + private static final int TOP_K = 10; + private static final VectorSimilarityFunction VSF = VectorSimilarityFunction.EUCLIDEAN; + + private List> vectors; + private ListRandomAccessVectorValues ravv; + + private void makeData() { + vectors = TestUtil.createRandomVectors(SIZE, DIM); + ravv = new ListRandomAccessVectorValues(vectors, DIM); + } + + /// With {@code callerRuns()}, every build and cleanup iteration must execute on the calling + /// thread — a probe recording the thread of each per-node scoring call must observe only the + /// test thread, proving no worker threads (and not the common pool) were used. + @Test + public void callerRunsExecutesOnlyOnCallingThread() throws IOException { + makeData(); + Thread testThread = Thread.currentThread(); + ThreadRecordingScoreProvider probe = + new ThreadRecordingScoreProvider(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF)); + + try (var builder = new GraphIndexBuilder(probe, DIM, M, BEAM, 1.2f, 1.2f, true, true, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns())) { + // Also exercise the removeDeletedNodes (sites 5/6) path under caller-runs. + builder.build(ravv); + for (int i = 0; i < SIZE; i += 25) { + builder.markNodeDeleted(i); + } + builder.cleanup(); + } + + assertTrue("build must have invoked the score provider", probe.observedThreads.size() >= 1); + assertEquals("caller-runs must execute all build/cleanup work on the calling thread only, but saw: " + + probe.observedThreads, + Set.of(testThread), probe.observedThreads); + } + + /// Caller-runs and ForkJoinPool builds of the same dataset must have equivalent recall. Graph + /// structure is not bit-identical (concurrent insertion is order-sensitive), so the bar is + /// recall-within-tolerance, both above a solid floor. + @Test + public void callerRunsMatchesForkJoinRecall() throws IOException { + makeData(); + List> queries = TestUtil.createRandomVectors(100, DIM); + List> groundTruth = bruteForceGroundTruth(queries); + + double callerRunsRecall; + try (var builder = new GraphIndexBuilder(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF), + DIM, M, BEAM, 1.2f, 1.2f, true, true, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns())) { + callerRunsRecall = recallAtK(builder.build(ravv), queries, groundTruth); + } + + ForkJoinPool pool = new ForkJoinPool(4); + double forkJoinRecall; + try (var builder = new GraphIndexBuilder(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF), + DIM, M, BEAM, 1.2f, 1.2f, true, true, pool, pool)) { + forkJoinRecall = recallAtK(builder.build(ravv), queries, groundTruth); + } finally { + pool.shutdown(); + } + + assertTrue("caller-runs recall too low: " + callerRunsRecall, callerRunsRecall > 0.85); + assertTrue("fork-join recall too low: " + forkJoinRecall, forkJoinRecall > 0.85); + assertTrue("caller-runs and fork-join recall must be equivalent, but were " + + callerRunsRecall + " vs " + forkJoinRecall, + Math.abs(callerRunsRecall - forkJoinRecall) < 0.1); + } + + /// The removeDeletedNodes path (cleanup sites 5/6) under caller-runs must produce a valid, + /// good-recall graph over the surviving nodes, equivalent to the ForkJoinPool path. + @Test + public void callerRunsRemoveDeletedNodesMatchesForkJoin() throws IOException { + makeData(); + List> queries = TestUtil.createRandomVectors(100, DIM); + + double callerRunsRecall = buildDeleteAndMeasure(ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns(), queries); + + ForkJoinPool pool = new ForkJoinPool(4); + double forkJoinRecall; + try { + forkJoinRecall = buildDeleteAndMeasure(ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool), queries); + } finally { + pool.shutdown(); + } + + assertTrue("caller-runs post-delete recall too low: " + callerRunsRecall, callerRunsRecall > 0.80); + assertTrue("post-delete recall must be equivalent, but were " + callerRunsRecall + " vs " + forkJoinRecall, + Math.abs(callerRunsRecall - forkJoinRecall) < 0.1); + } + + /// The existing ForkJoinPool constructor must still build a good graph (back-compat). + @Test + public void forkJoinPoolConstructorStillBuildsGoodGraph() throws IOException { + makeData(); + List> queries = TestUtil.createRandomVectors(100, DIM); + List> groundTruth = bruteForceGroundTruth(queries); + + try (var builder = new GraphIndexBuilder(ravv, VSF, M, BEAM, 1.2f, 1.2f, true, true)) { + // default-pool constructor path + assertTrue(recallAtK(builder.build(ravv), queries, groundTruth) > 0.85); + } + } + + // ---- helpers ---- + + private double buildDeleteAndMeasure(ParallelExecutor simd, ParallelExecutor parallel, List> queries) throws IOException { + try (var builder = new GraphIndexBuilder(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF), + DIM, M, BEAM, 1.2f, 1.2f, true, true, simd, parallel)) { + var graph = builder.build(ravv); + Set deleted = new HashSet<>(); + for (int i = 0; i < SIZE; i += 10) { + builder.markNodeDeleted(i); + deleted.add(i); + } + builder.cleanup(); + List> gt = bruteForceGroundTruthExcluding(queries, deleted); + return recallAtK(graph, queries, gt); + } + } + + private List> bruteForceGroundTruth(List> queries) { + return bruteForceGroundTruthExcluding(queries, Set.of()); + } + + private List> bruteForceGroundTruthExcluding(List> queries, Set excluded) { + List> out = new ArrayList<>(queries.size()); + for (VectorFloat q : queries) { + int[] best = new int[TOP_K]; + float[] bestScore = new float[TOP_K]; + int count = 0; + for (int j = 0; j < vectors.size(); j++) { + if (excluded.contains(j)) { + continue; + } + float s = VSF.compare(q, vectors.get(j)); + // insert into the small top-K buffer (higher score = closer) + if (count < TOP_K) { + best[count] = j; + bestScore[count] = s; + count++; + } else { + int minIdx = 0; + for (int t = 1; t < TOP_K; t++) { + if (bestScore[t] < bestScore[minIdx]) minIdx = t; + } + if (s > bestScore[minIdx]) { + best[minIdx] = j; + bestScore[minIdx] = s; + } + } + } + Set gt = new HashSet<>(); + for (int t = 0; t < count; t++) gt.add(best[t]); + out.add(gt); + } + return out; + } + + private double recallAtK(ImmutableGraphIndex graph, List> queries, List> groundTruth) throws IOException { + int efSearch = 100; // search width >> topK, matching TestVectorGraph's recall methodology + double total = 0; + for (int qi = 0; qi < queries.size(); qi++) { + SearchResult.NodeScore[] nodes = + GraphSearcher.search(queries.get(qi), efSearch, ravv, VSF, graph, Bits.ALL).getNodes(); + Set gt = groundTruth.get(qi); + int limit = Math.min(TOP_K, nodes.length); + int hits = 0; + for (int t = 0; t < limit; t++) { + if (gt.contains(nodes[t].node)) hits++; + } + total += (double) hits / Math.max(1, gt.size()); + } + return total / queries.size(); + } + + /// A {@link BuildScoreProvider} that records the thread of every per-node scoring request, + /// delegating all work. Used to prove caller-runs never leaves the calling thread. + private static final class ThreadRecordingScoreProvider implements BuildScoreProvider { + private final BuildScoreProvider delegate; + final Set observedThreads = ConcurrentHashMap.newKeySet(); + + ThreadRecordingScoreProvider(BuildScoreProvider delegate) { + this.delegate = delegate; + } + + @Override + public boolean isExact() { + return delegate.isExact(); + } + + @Override + public VectorFloat approximateCentroid() { + return delegate.approximateCentroid(); + } + + @Override + public SearchScoreProvider searchProviderFor(VectorFloat vector) { + return delegate.searchProviderFor(vector); + } + + @Override + public SearchScoreProvider searchProviderFor(int node1) { + observedThreads.add(Thread.currentThread()); + return delegate.searchProviderFor(node1); + } + + @Override + public SearchScoreProvider diversityProviderFor(int node1) { + observedThreads.add(Thread.currentThread()); + return delegate.diversityProviderFor(node1); + } + + @Override + public ScoreFunction diversityScoreFunctionFor(int node1) { + observedThreads.add(Thread.currentThread()); + return delegate.diversityScoreFunctionFor(node1); + } + } +} From 9921fa4e2ea501dcc4d7284ca06753c8b9dd0ab8 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 6 Jul 2026 22:55:30 +0000 Subject: [PATCH 09/13] add memory-safety repro and regression-guard tests --- .../example/repro/BoundsCheckBenchTest.java | 125 +++++ .../jvector/example/repro/ChildJvm.java | 157 ++++++ .../repro/CompactorStragglerReproTest.java | 500 ++++++++++++++++++ .../example/repro/GatedReaderSupplier.java | 287 ++++++++++ .../example/repro/HostStyleMappedReader.java | 113 ++++ .../IndexFileLifecycleCorruptionTest.java | 314 +++++++++++ .../repro/MemorySafetyCrashReproTest.java | 151 ++++++ .../repro/MemorySafetyReproHarness.java | 330 ++++++++++++ .../jvector/example/repro/ReproGraphs.java | 222 ++++++++ .../disk/TestOnDiskGraphIndexViewGuards.java | 189 +++++++ 10 files changed, 2388 insertions(+) create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/BoundsCheckBenchTest.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ChildJvm.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/CompactorStragglerReproTest.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/GatedReaderSupplier.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/HostStyleMappedReader.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/IndexFileLifecycleCorruptionTest.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyCrashReproTest.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyReproHarness.java create mode 100644 jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ReproGraphs.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexViewGuards.java diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/BoundsCheckBenchTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/BoundsCheckBenchTest.java new file mode 100644 index 000000000..2f3833b18 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/BoundsCheckBenchTest.java @@ -0,0 +1,125 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.GraphSearcher; +import io.github.jbellis.jvector.graph.SearchResult; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.similarity.DefaultSearchScoreProvider; +import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Assume; +import org.junit.Test; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; + +/// A/B microbenchmark for the F3 per-record-access bounds checks (`requireValidNode` + the real +/// degree check in `OnDiskGraphIndex.View`): measures single-threaded exact-scored search +/// throughput over an on-disk graph, the hot path where every neighbor-block read and every +/// scored vector read passes the guard. +/// +/// Not part of the normal suite — gated on `REPRO_BENCH=1`. Protocol (one fresh JVM per run, the +/// graph file is built once and reused so both variants read identical bytes): +/// +/// ``` +/// REPRO_BENCH=1 mvn test -pl jvector-examples -am -Dtest='BoundsCheckBenchTest' \ +/// -Dsurefire.failIfNoSpecifiedTests=false \ +/// -DargLine="--add-modules jdk.incubator.vector --enable-native-access=ALL-UNNAMED" +/// # run WITH checks (working tree), then: +/// git stash push -- jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java +/// # run WITHOUT, pop, repeat interleaved +/// ``` +/// +/// The printed `checksum` folds every result node id and score; it must be identical across +/// variants — proof that both did exactly the same traversal work. +public class BoundsCheckBenchTest { + private static final int N = 32_768; + private static final int DIM = 64; + private static final int QUERIES = 4096; + private static final int TOP_K = 10; + private static final int WARMUP_ROUNDS = 5; + private static final int MEASURED_ROUNDS = 15; + + @Test(timeout = 1_800_000) + public void searchThroughputAB() throws Exception { + Assume.assumeTrue("set REPRO_BENCH=1 to run the bounds-check benchmark", + "1".equals(System.getenv("REPRO_BENCH"))); + + System.out.println("BENCH|provider|" + VectorizationProvider.getInstance().getClass().getSimpleName()); + + // Build once, reuse across variant runs: identical bytes for both sides of the A/B. + Path benchDir = ReproGraphs.newWorkDir("bench").getParent().resolve("bench-fixed"); + Files.createDirectories(benchDir); + Path graphPath = benchDir.resolve("bench_n" + N + "_d" + DIM + ".graph"); + if (!Files.exists(graphPath)) { + long t0 = System.nanoTime(); + ReproGraphs.buildInlineGraph(graphPath, ReproGraphs.randomVectors(N, DIM, 777), DIM); + System.out.println("BENCH|built|" + ((System.nanoTime() - t0) / 1_000_000) + "ms|" + Files.size(graphPath) + "bytes"); + } else { + System.out.println("BENCH|reused|" + Files.size(graphPath) + "bytes"); + } + + List> queries = ReproGraphs.randomVectors(QUERIES, DIM, 999); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(graphPath)) { + OnDiskGraphIndex graph = OnDiskGraphIndex.load(rs); + try (var scoringView = graph.getView(); + GraphSearcher searcher = new GraphSearcher(graph)) { + + long[] roundUs = new long[MEASURED_ROUNDS]; + long checksum = 0; + for (int round = 0; round < WARMUP_ROUNDS + MEASURED_ROUNDS; round++) { + long sink = 0; + long start = System.nanoTime(); + for (VectorFloat q : queries) { + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, VectorSimilarityFunction.EUCLIDEAN, scoringView); + SearchResult sr = searcher.search(ssp, TOP_K, Bits.ALL); + for (var node : sr.getNodes()) { + sink += node.node; + sink += Float.floatToIntBits(node.score); + } + } + long elapsedUs = (System.nanoTime() - start) / 1_000; + boolean measured = round >= WARMUP_ROUNDS; + if (measured) { + roundUs[round - WARMUP_ROUNDS] = elapsedUs; + checksum = sink; // identical every round; keep the last + } + System.out.println("BENCH|round|" + (measured ? "measured" : "warmup") + "|" + round + + "|" + elapsedUs + "us|" + (QUERIES * 1_000_000L / Math.max(1, elapsedUs)) + "qps|sink=" + sink); + } + + long[] sorted = roundUs.clone(); + Arrays.sort(sorted); + long median = sorted[sorted.length / 2]; + long best = sorted[0]; + System.out.println("BENCH|summary|median=" + median + "us|best=" + best + + "us|medianQps=" + (QUERIES * 1_000_000L / Math.max(1, median)) + + "|nsPerQuery=" + (median * 1_000L / QUERIES) + + "|checksum=" + checksum); + } + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ChildJvm.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ChildJvm.java new file mode 100644 index 000000000..ca70cfb91 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ChildJvm.java @@ -0,0 +1,157 @@ +/* + * 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.example.repro; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import java.util.stream.Stream; + +/// Launches [MemorySafetyReproHarness] scenarios in a throwaway child JVM so that scenarios whose +/// expected outcome is a native JVM crash (SIGSEGV / SIGBUS) can be asserted on from a healthy +/// parent. The child writes its `hs_err` file into the scenario's work directory; the parent +/// collects exit code, combined stdout/stderr, and the parsed crash evidence (`siginfo` line and +/// problematic frame — the discriminators the bug doc's trigger table calls for). +public final class ChildJvm { + + /// Everything the parent needs to classify one child run. + public static final class Result { + public final int exitCode; + public final String output; + public final String hsErr; // null when the child did not crash + + Result(int exitCode, String output, String hsErr) { + this.exitCode = exitCode; + this.output = output; + this.hsErr = hsErr; + } + + public boolean crashed() { + return hsErr != null; + } + + /// The `siginfo:` line from the hs_err file (signal + si_code), or "" if absent. + public String siginfo() { + return firstMatch("(?m)^siginfo:.*$"); + } + + /// The `# Problematic frame:` detail line from the hs_err file, or "" if absent. + public String problematicFrame() { + Matcher m = Pattern.compile("(?m)^# Problematic frame:\\R# (.*)$").matcher(hsErr == null ? "" : hsErr); + return m.find() ? m.group(1).trim() : ""; + } + + /// The `Current thread` line from the hs_err file, or "" if absent. + public String currentThread() { + return firstMatch("(?m)^Current thread.*$"); + } + + private String firstMatch(String regex) { + if (hsErr == null) { + return ""; + } + Matcher m = Pattern.compile(regex).matcher(hsErr); + return m.find() ? m.group().trim() : ""; + } + + public boolean hsErrContains(String needle) { + return hsErr != null && hsErr.contains(needle); + } + + /// One-line summary for assertion messages and the verdict report. + public String summary() { + if (crashed()) { + return "exit=" + exitCode + " CRASHED [" + siginfo() + "] frame=[" + problematicFrame() + "]"; + } + String outcome = ""; + for (String line : output.split("\\R")) { + if (line.startsWith("REPRO|OUTCOME|")) { + outcome = line; + } + } + return "exit=" + exitCode + " " + (outcome.isEmpty() ? "(no outcome line)" : outcome); + } + } + + private ChildJvm() { + } + + /// Runs `MemorySafetyReproHarness ` in a child JVM, with `hs_err` redirected + /// into `workDir`. Blocks up to `timeoutSeconds`, then kills and fails. + public static Result run(Path workDir, long timeoutSeconds, String... harnessArgs) throws IOException, InterruptedException { + String javaBin = Path.of(System.getProperty("java.home"), "bin", "java").toString(); + List cmd = new ArrayList<>(); + cmd.add(javaBin); + cmd.add("-cp"); + cmd.add(System.getProperty("java.class.path")); + cmd.add("-ea"); + cmd.add("-Xmx1g"); + cmd.add("-XX:-CreateCoredumpOnCrash"); + cmd.add("-XX:ErrorFile=" + workDir.resolve("hs_err_pid%p.log").toAbsolutePath()); + if (Runtime.version().feature() >= 22) { + cmd.add("--enable-native-access=ALL-UNNAMED"); + } + cmd.add(MemorySafetyReproHarness.class.getName()); + cmd.addAll(Arrays.asList(harnessArgs)); + + Process process = new ProcessBuilder(cmd).redirectErrorStream(true).start(); + StringBuilder out = new StringBuilder(); + Thread drain = new Thread(() -> { + try (BufferedReader r = new BufferedReader(new InputStreamReader(process.getInputStream(), StandardCharsets.UTF_8))) { + String line; + while ((line = r.readLine()) != null) { + synchronized (out) { + out.append(line).append('\n'); + } + } + } catch (IOException ignored) { + // stream closes when the child dies; partial output is fine + } + }, "repro-child-drain"); + drain.start(); + + if (!process.waitFor(timeoutSeconds, TimeUnit.SECONDS)) { + process.destroyForcibly(); + drain.join(5_000); + throw new IllegalStateException("child JVM did not finish within " + timeoutSeconds + "s; args=" + Arrays.toString(harnessArgs) + + "\n--- child output ---\n" + out); + } + drain.join(10_000); + + String hsErr = null; + try (Stream files = Files.list(workDir)) { + Path errFile = files.filter(p -> p.getFileName().toString().startsWith("hs_err_pid")).findFirst().orElse(null); + if (errFile != null) { + hsErr = Files.readString(errFile, StandardCharsets.ISO_8859_1); + } + } + String output; + synchronized (out) { + output = out.toString(); + } + return new Result(process.exitValue(), output, hsErr); + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/CompactorStragglerReproTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/CompactorStragglerReproTest.java new file mode 100644 index 000000000..1a89b67c1 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/CompactorStragglerReproTest.java @@ -0,0 +1,500 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor; +import io.github.jbellis.jvector.util.work.ProgressLimiter; +import io.github.jbellis.jvector.util.work.WorkStage; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import org.junit.Assume; +import org.junit.Test; + +import java.nio.file.Path; +import java.util.List; +import java.util.Locale; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +/// Guard tests for the F1 drain-on-unwind fix in `OnDiskGraphIndexCompactor` (prescription: +/// `local/memory_safety_fix_plan.md` F1; the defect these tests originally *proved* is documented +/// with its evidence in `local/memory_safety_design_brief.md` §4): when a merge batch fails or the +/// orchestrating thread is interrupted, `compact()` must block until every in-flight batch has +/// finished before the exception escapes. Otherwise control returns to a caller who is entitled +/// to release the source graphs while worker tasks are still reading them — natively fatal for +/// NIO-mapped sources; the pre-fix form of T4b reproduced the production core dump exactly that +/// way (SIGSEGV/`SEGV_MAPERR` in the merge read lineage). +/// +/// - T4a checks the drained unwind deterministically, in-process: a poisoned source read fails +/// one batch while another worker's source read is provably parked in flight; `compact()` must +/// not throw until that read has completed. +/// - T4b runs the production sequence in a child JVM: after `compact()` throws, a +/// contract-compliant host closes its host-style mapped sources — which must now be crash-free. +/// - T5 interrupts the orchestrator mid-refine (the realistic host trigger: nodetool stop / +/// shutdown / cancellation) while one refine task is provably parked in flight; `compact()` +/// must not throw until that task has finished. +/// - T5b (the F5 fix) interrupts the orchestrator during the fused pre-encode fan-out +/// (`invokeAll` on the strategies' executor adapter) while one pre-encode task is provably +/// parked in flight. Stock `invokeAll` cancels-with-interrupt and returns immediately, so +/// pre-F5 that task — which reads source graphs and writes the code cache that +/// `onAfterClose` unmaps — outlived `compact()`. +public class CompactorStragglerReproTest { + + private static void assumeLinux() { + Assume.assumeTrue("child-JVM crash detection relies on Linux mmap semantics", + System.getProperty("os.name", "").toLowerCase(Locale.ROOT).contains("linux")); + } + + private static void verdict(String theory, String text) { + System.out.println("VERDICT|" + theory + "|" + text); + } + + private static ThreadFactory workerFactory() { + return new WorkerThreadFactory(); + } + + private static final class WorkerThreadFactory implements ThreadFactory { + private final AtomicInteger n = new AtomicInteger(); + + @Override + public Thread newThread(Runnable r) { + Thread t = new Thread(r, GatedReaderSupplier.Gate.WORKER_PREFIX + n.getAndIncrement()); + t.setDaemon(true); + return t; + } + } + + private static boolean chainContains(Throwable t, String needle) { + for (Throwable c = t; c != null; c = c.getCause()) { + if (String.valueOf(c.getMessage()).contains(needle) || c.getClass().getName().contains(needle)) { + return true; + } + } + return false; + } + + private static String chain(Throwable t) { + StringBuilder sb = new StringBuilder(); + for (Throwable c = t; c != null; c = c.getCause()) { + if (sb.length() > 0) { + sb.append(" <- "); + } + sb.append(c.getClass().getSimpleName()).append('(').append(c.getMessage()).append(')'); + } + return sb.toString(); + } + + /// F1 guard: `compact()` must not return control while source reads are in flight. One batch + /// is poisoned while another worker's source read is parked — provably in flight — and the + /// drain must complete that read (released by the [GatedReaderSupplier.AutoReleaser] ~1s + /// later) before the failure escapes. Pre-F1 this test proved the inverse: `compact()` threw + /// with 3 source reads still executing and thousands more completing after the unwind. + @Test(timeout = 240_000) + public void t4a_compactDrainsInflightSourceReadsBeforeUnwinding() throws Exception { + Path dir = ReproGraphs.newWorkDir("t4a-straggler"); + final int dimension = 24; + final int perSource = 220; + Path s0 = ReproGraphs.buildInlineGraph(dir.resolve("s0.graph"), ReproGraphs.randomVectors(perSource, dimension, 301), dimension); + Path s1 = ReproGraphs.buildInlineGraph(dir.resolve("s1.graph"), ReproGraphs.randomVectors(perSource, dimension, 302), dimension); + + try (ReaderSupplier r0 = ReaderSupplierFactory.open(s0); + ReaderSupplier r1 = ReaderSupplierFactory.open(s1)) { + GatedReaderSupplier.Gate gate = new GatedReaderSupplier.Gate(); + GatedReaderSupplier g0 = new GatedReaderSupplier(r0, gate, GatedReaderSupplier.Role.POISON); + GatedReaderSupplier g1 = new GatedReaderSupplier(r1, gate, GatedReaderSupplier.Role.PARK); + OnDiskGraphIndex src0 = OnDiskGraphIndex.load(g0); + OnDiskGraphIndex src1 = OnDiskGraphIndex.load(g1); + + ExecutorService pool = Executors.newFixedThreadPool(4, workerFactory()); + try { + var compactor = new OnDiskGraphIndexCompactor( + List.of(src0, src1), + ReproGraphs.allLive(perSource, perSource), + ReproGraphs.stackedRemappers(perSource, perSource), + VectorSimilarityFunction.EUCLIDEAN, + pool, 4); + + gate.arm(); + Thread releaser = new Thread(new GatedReaderSupplier.AutoReleaser(gate, 1_000), "repro-gate-releaser"); + releaser.setDaemon(true); + releaser.start(); + + Throwable thrown = null; + int activeReadsAtThrow = -1; + try { + compactor.compact(dir.resolve("out.graph")); + } catch (Throwable t) { + gate.mark(); + activeReadsAtThrow = gate.activeWorkerReads(); + thrown = t; + } + + assertNotNull("the poisoned batch must make compact() throw", thrown); + assertTrue("compact() must have failed because of the injected poison, but failed with: " + chain(thrown), + chainContains(thrown, "REPRO_POISON")); + assertNotNull("the repro choreography must have parked a source read across the failure", + gate.parkedThread()); + assertEquals("F1: compact() must drain every in-flight source read before unwinding", + 0, activeReadsAtThrow); + assertFalse("the parked read must complete before compact() throws, never after", + gate.parkedReadResumedAfterMark()); + + pool.shutdown(); + assertTrue("workers must quiesce", pool.awaitTermination(120, TimeUnit.SECONDS)); + assertEquals("no source read may complete after compact() has thrown", + 0, gate.workerReadsCompletedAfterMark()); + verdict("T4a", "FIX HOLDS: compact() drained all in-flight source reads before unwinding" + + " (0 active at throw, 0 completed after)"); + } finally { + gate.releaseParked(); + pool.shutdownNow(); + } + } + } + + /// F1 guard: the production sequence must be crash-free end to end. In a child JVM, after + /// `compact()` throws, the host — correctly, per the documented contract — closes its + /// host-style mapped sources (a raw munmap). The drain guarantees nothing is still reading + /// them, so the child must reach the orderly NO_CRASH outcome (exit 45). Pre-F1 this exact + /// sequence died with SIGSEGV/`SEGV_MAPERR` on a worker thread inside the merge read lineage + /// — the observed production core dump. + @Test(timeout = 360_000) + public void t4b_contractCompliantSourceCloseAfterCompactThrowsIsCrashFree() throws Exception { + assumeLinux(); + Path dir = ReproGraphs.newWorkDir("t4b-straggler-crash"); + ChildJvm.Result r = ChildJvm.run(dir, 300, "compactor-straggler", dir.toString()); + + assertTrue("child must reach the compact() failure: " + r.summary(), r.output.contains("STAGE|COMPACT_THREW")); + assertNotEquals("compact() unexpectedly succeeded despite the poisoned batch", 44, r.exitCode); + assertTrue("the drain must leave zero in-flight source reads at the throw: " + r.summary(), + r.output.contains("ACTIVE_WORKER_READS=0")); + assertTrue("child must close sources only after compact() threw: " + r.summary(), + r.output.contains("STAGE|SOURCES_CLOSED")); + assertFalse("F1: contract-compliant source close after compact() throws must be crash-free: " + r.summary(), + r.crashed()); + assertEquals("child must survive to the orderly no-crash outcome: " + r.summary(), 45, r.exitCode); + assertTrue("no source read may complete after the unwind: " + r.summary(), + r.output.contains("readsCompletedAfterThrow=0")); + verdict("T4b", "FIX HOLDS: contract-compliant close after compact() threw is crash-free: " + r.summary()); + } + + /// F1 guard: the refine window loop must drain its in-flight tasks before an interrupt-driven + /// unwind escapes. One refine task is parked — provably in flight — when the orchestrator is + /// interrupted; the 1s hold before release is the regression window in which the pre-F1 code + /// threw (abandoning the parked task and up to a window of others, which then raced the + /// unwind's closing of the output supplier and channel). Post-F1, `compact()` must not throw + /// until the parked task has finished. + @Test(timeout = 600_000) + public void t5_refineLoopDrainsInflightTasksOnInterrupt() throws Exception { + Path dir = ReproGraphs.newWorkDir("t5-refine-interrupt"); + final int dimension = 16; + final int perSource = 1000; + Path s0 = ReproGraphs.buildInlineGraph(dir.resolve("s0.graph"), ReproGraphs.randomVectors(perSource, dimension, 401), dimension); + Path s1 = ReproGraphs.buildInlineGraph(dir.resolve("s1.graph"), ReproGraphs.randomVectors(perSource, dimension, 402), dimension); + + try (ReaderSupplier r0 = ReaderSupplierFactory.open(s0); + ReaderSupplier r1 = ReaderSupplierFactory.open(s1)) { + OnDiskGraphIndex src0 = OnDiskGraphIndex.load(r0); + OnDiskGraphIndex src1 = OnDiskGraphIndex.load(r1); + + ExecutorService pool = Executors.newFixedThreadPool(4, workerFactory()); + CountingExecutor counting = new CountingExecutor(pool); + try { + var compactor = new OnDiskGraphIndexCompactor( + List.of(src0, src1), + ReproGraphs.allLive(perSource, perSource), + ReproGraphs.stackedRemappers(perSource, perSource), + VectorSimilarityFunction.EUCLIDEAN, + counting, 4); + RefineSignal refineSignal = new RefineSignal(); + compactor.setProgressLimiter(refineSignal); + // Refinement is opt-in (default off); this F1 guard parks a refine task and + // interrupts mid-refine, so the refinement pass must be enabled explicitly. + compactor.setRefineAfterCompaction(true); + + AtomicReference thrown = new AtomicReference<>(); + AtomicInteger inFlightAtThrow = new AtomicInteger(-1); + AtomicLong thrownAtNanos = new AtomicLong(); + Thread orchestrator = new Thread( + new CompactRun(compactor, dir.resolve("out.graph"), counting, thrown, inFlightAtThrow, thrownAtNanos), + "repro-orchestrator"); + orchestrator.start(); + + assertTrue("refine phase must begin (merge finished)", refineSignal.refineStarted.await(480, TimeUnit.SECONDS)); + counting.armParkOnce(); + assertTrue("a refine task must park in flight", counting.awaitParked(120, TimeUnit.SECONDS)); + orchestrator.interrupt(); + // Regression window: pre-F1, compact() throws during this hold with the parked + // task still in flight; post-F1 the drain blocks on it instead. + Thread.sleep(1_000); + counting.releaseParked(); + + orchestrator.join(TimeUnit.SECONDS.toMillis(120)); + assertFalse("orchestrator must unwind after interrupt", orchestrator.isAlive()); + assertNotNull("interrupt during refine must make compact() throw", thrown.get()); + assertTrue("unwind must be caused by the interrupt, was: " + chain(thrown.get()), + chainContains(thrown.get(), "InterruptedException")); + + pool.shutdown(); + assertTrue("workers must quiesce", pool.awaitTermination(120, TimeUnit.SECONDS)); + + long graceNanos = TimeUnit.MILLISECONDS.toNanos(250); + long parkedEndMinusThrow = counting.parkedTaskEndNanos() - thrownAtNanos.get(); + assertTrue("F1: compact() must not throw while a refine task is in flight" + + " (parked task ended " + TimeUnit.NANOSECONDS.toMillis(parkedEndMinusThrow) + + "ms after the throw; expected at/before it)", + counting.parkedTaskEndNanos() <= thrownAtNanos.get() + graceNanos); + verdict("T5", "FIX HOLDS: refine unwind drained in-flight tasks; parked task finished " + + TimeUnit.NANOSECONDS.toMillis(Math.max(0, -parkedEndMinusThrow)) + + "ms before compact() threw (inFlightAtThrow=" + inFlightAtThrow.get() + ")"); + } finally { + counting.releaseParked(); + pool.shutdownNow(); + } + } + } + + /// F5 guard: the strategies' `invokeAll` fan-out — fused pre-encode here — must drain its + /// started tasks before an interrupt-driven unwind escapes `compact()`. Pre-encode tasks + /// open source views (`getVectorInto` per live node) and write the pre-encode code cache + /// that `compact()`'s finally unmaps, so pre-F5 abandonment was doubly fatal: the caller + /// could unmap sources under a live read, and jvector itself unmapped the cache under a + /// live write. The parked task ignores interrupts, exactly like real encode work, so stock + /// `invokeAll`'s cancel-with-interrupt cannot end it — only a real drain can. + @Test(timeout = 600_000) + public void t5b_preEncodeFanoutDrainsInflightTasksOnInterrupt() throws Exception { + Path dir = ReproGraphs.newWorkDir("t5b-preencode-interrupt"); + final int dimension = 32; + final int perSource = 256; + Path s0 = ReproGraphs.buildFusedGraph(dir.resolve("s0.graph"), ReproGraphs.randomVectors(perSource, dimension, 501), dimension); + Path s1 = ReproGraphs.buildFusedGraph(dir.resolve("s1.graph"), ReproGraphs.randomVectors(perSource, dimension, 502), dimension); + + try (ReaderSupplier r0 = ReaderSupplierFactory.open(s0); + ReaderSupplier r1 = ReaderSupplierFactory.open(s1)) { + OnDiskGraphIndex src0 = OnDiskGraphIndex.load(r0); + OnDiskGraphIndex src1 = OnDiskGraphIndex.load(r1); + + ExecutorService pool = Executors.newFixedThreadPool(4, workerFactory()); + CountingExecutor counting = new CountingExecutor(pool); + try { + var compactor = new OnDiskGraphIndexCompactor( + List.of(src0, src1), + ReproGraphs.allLive(perSource, perSource), + ReproGraphs.stackedRemappers(perSource, perSource), + VectorSimilarityFunction.COSINE, + counting, 4); + + AtomicReference thrown = new AtomicReference<>(); + AtomicInteger inFlightAtThrow = new AtomicInteger(-1); + AtomicLong thrownAtNanos = new AtomicLong(); + // In a fused compaction the first task through the executor is a pre-encode + // chunk (strategy.onAfterHeader -> precomputeCodes -> invokeAll, before any + // merge batch), so arming the one-shot park before compact() pins a pre-encode + // task in flight. + counting.armParkOnce(); + Thread orchestrator = new Thread( + new CompactRun(compactor, dir.resolve("out.graph"), counting, thrown, inFlightAtThrow, thrownAtNanos), + "repro-orchestrator"); + orchestrator.start(); + + assertTrue("a pre-encode task must park in flight", counting.awaitParked(120, TimeUnit.SECONDS)); + orchestrator.interrupt(); + // Regression window: with stock invokeAll, compact() throws during this hold — + // cancel-with-interrupt does not end the parked task; with the F5 drain it + // blocks until the release below. + Thread.sleep(1_000); + counting.releaseParked(); + + orchestrator.join(TimeUnit.SECONDS.toMillis(120)); + assertFalse("orchestrator must unwind after interrupt", orchestrator.isAlive()); + assertNotNull("interrupt during pre-encode must make compact() throw", thrown.get()); + assertTrue("unwind must be caused by the interrupt, was: " + chain(thrown.get()), + chainContains(thrown.get(), "InterruptedException")); + + pool.shutdown(); + assertTrue("workers must quiesce", pool.awaitTermination(120, TimeUnit.SECONDS)); + + long graceNanos = TimeUnit.MILLISECONDS.toNanos(250); + long parkedEndMinusThrow = counting.parkedTaskEndNanos() - thrownAtNanos.get(); + assertTrue("F5: compact() must not throw while a pre-encode task is in flight" + + " (parked task ended " + TimeUnit.NANOSECONDS.toMillis(parkedEndMinusThrow) + + "ms after the throw; expected at/before it)", + counting.parkedTaskEndNanos() <= thrownAtNanos.get() + graceNanos); + verdict("T5b", "FIX HOLDS: pre-encode fan-out drained before unwind; parked task finished " + + TimeUnit.NANOSECONDS.toMillis(Math.max(0, -parkedEndMinusThrow)) + + "ms before compact() threw (inFlightAtThrow=" + inFlightAtThrow.get() + ")"); + } finally { + counting.releaseParked(); + pool.shutdownNow(); + } + } + } + + /// Runs `compact()` and records the unwind instant plus how many pool tasks were mid-flight. + private static final class CompactRun implements Runnable { + private final OnDiskGraphIndexCompactor compactor; + private final Path out; + private final CountingExecutor counting; + private final AtomicReference thrown; + private final AtomicInteger inFlightAtThrow; + private final AtomicLong thrownAtNanos; + + CompactRun(OnDiskGraphIndexCompactor compactor, Path out, CountingExecutor counting, + AtomicReference thrown, AtomicInteger inFlightAtThrow, AtomicLong thrownAtNanos) { + this.compactor = compactor; + this.out = out; + this.counting = counting; + this.thrown = thrown; + this.inFlightAtThrow = inFlightAtThrow; + this.thrownAtNanos = thrownAtNanos; + } + + @Override + public void run() { + try { + compactor.compact(out); + } catch (Throwable t) { + thrownAtNanos.set(System.nanoTime()); + inFlightAtThrow.set(counting.inFlight()); + thrown.set(t); + } + } + } + + /// Signals when the compactor reports the first REFINE progress event (orchestrator-side). + private static final class RefineSignal implements ProgressLimiter { + final CountDownLatch refineStarted = new CountDownLatch(1); + + @Override + public void onProgress(WorkStage stage, long completed, long total) { + if (stage == OnDiskGraphIndexCompactor.Phase.REFINE) { + refineStarted.countDown(); + } + } + } + + /// Wraps the pool so the test can observe compactor-submitted tasks — how many are mid-flight + /// and when each finished — and, for the F1 guard, park a single task (one-shot) provably in + /// flight until released. + private static final class CountingExecutor implements Executor { + private final Executor delegate; + private final AtomicInteger inFlight = new AtomicInteger(); + private final AtomicLong lastTaskEndNanos = new AtomicLong(); + private final AtomicBoolean parkArmed = new AtomicBoolean(); + private final CountDownLatch parkedLatch = new CountDownLatch(1); + private final CountDownLatch parkRelease = new CountDownLatch(1); + private final AtomicLong parkedTaskEndNanos = new AtomicLong(); + + CountingExecutor(Executor delegate) { + this.delegate = delegate; + } + + int inFlight() { + return inFlight.get(); + } + + long lastTaskEndNanos() { + return lastTaskEndNanos.get(); + } + + long parkedTaskEndNanos() { + return parkedTaskEndNanos.get(); + } + + /// Arms the one-shot park: the next task to start executing parks until [#releaseParked()]. + void armParkOnce() { + parkArmed.set(true); + } + + boolean awaitParked(long timeout, TimeUnit unit) throws InterruptedException { + return parkedLatch.await(timeout, unit); + } + + void releaseParked() { + parkRelease.countDown(); + } + + @Override + public void execute(Runnable command) { + delegate.execute(new CountedTask(command)); + } + + private final class CountedTask implements Runnable { + private final Runnable command; + + CountedTask(Runnable command) { + this.command = command; + } + + @Override + public void run() { + boolean parkedTask = parkArmed.compareAndSet(true, false); + inFlight.incrementAndGet(); + try { + if (parkedTask) { + parkedLatch.countDown(); + // Uninterruptible on purpose: this models a real merge/encode chunk — + // CPU + mmap work that never polls the interrupt flag — so a + // cancel-with-interrupt (stock invokeAll's unwind behavior) cannot make + // it "finish" early. Interrupts are remembered and re-asserted. + boolean sawInterrupt = false; + boolean released = false; + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(60); + while (!released) { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0) { + throw new IllegalStateException("parked repro task was never released"); + } + try { + released = parkRelease.await(remaining, TimeUnit.NANOSECONDS); + } catch (InterruptedException e) { + sawInterrupt = true; + } + } + if (sawInterrupt) { + Thread.currentThread().interrupt(); + } + } + command.run(); + } finally { + inFlight.decrementAndGet(); + long now = System.nanoTime(); + lastTaskEndNanos.accumulateAndGet(now, Math::max); + if (parkedTask) { + parkedTaskEndNanos.set(now); + } + } + } + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/GatedReaderSupplier.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/GatedReaderSupplier.java new file mode 100644 index 000000000..52d02c5ad --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/GatedReaderSupplier.java @@ -0,0 +1,287 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.disk.ReaderSupplier; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +/// A [ReaderSupplier] wrapper that instruments every read a compaction worker performs against a +/// source graph, so a test can determine — deterministically — whether `compact()` can unwind +/// while source reads are still in flight. It originally proved the straggler-abandonment defect +/// (see `local/memory_safety_design_brief.md`, T4); since the F1 drain-on-unwind fix it guards +/// the fixed invariant: `compact()` returns or throws only after every worker read has finished. +/// +/// Two roles drive the choreography, coordinated through a shared [Gate]: +/// +/// - [Role#PARK]: the first worker-thread read against this source parks (before touching the +/// delegate) until the gate is released. This pins a provably in-flight source read across the +/// moment `compact()` throws. +/// - [Role#POISON]: once a reader has parked, the next worker-thread read against this source on +/// a *different* thread throws, failing that merge batch. The orchestrator sees the failure at +/// `ExecutorCompletionService.take().get()` and `compact()` unwinds — while the parked read (and +/// any other in-flight batches) are still executing. +/// +/// Gating applies only on threads named with [Gate#WORKER_PREFIX] and only after [Gate#arm()], so +/// graph loading and orchestrator-side reads pass through untouched. +public final class GatedReaderSupplier implements ReaderSupplier { + + /// What this source's reads do once the gate is armed. + public enum Role { + NONE, PARK, POISON + } + + /// Shared coordination state between the parked source, the poisoned source, and the test. + public static final class Gate { + public static final String WORKER_PREFIX = "repro-worker-"; + + final CountDownLatch parked = new CountDownLatch(1); + final CountDownLatch release = new CountDownLatch(1); + final AtomicReference parkedThread = new AtomicReference<>(); + final AtomicBoolean poisonFired = new AtomicBoolean(); + final AtomicInteger activeWorkerReads = new AtomicInteger(); + final AtomicInteger workerReadsCompletedAfterMark = new AtomicInteger(); + final AtomicBoolean parkedReadResumedAfterMark = new AtomicBoolean(); + volatile boolean armed; + volatile boolean marked; + + /// Starts gating worker reads; call after sources are loaded, right before `compact()`. + public void arm() { + armed = true; + } + + /// Records the observation point (typically: `compact()` just threw); reads that complete + /// after this are counted as post-unwind stragglers. + public void mark() { + marked = true; + } + + /// True once a worker read is parked inside the wrapped source. + public boolean awaitParked(long timeout, TimeUnit unit) throws InterruptedException { + return parked.await(timeout, unit); + } + + /// Lets the parked read proceed into the delegate. + public void releaseParked() { + release.countDown(); + } + + /// Worker reads currently inside a wrapped read call (parked ones included). + public int activeWorkerReads() { + return activeWorkerReads.get(); + } + + /// Worker reads that ran to completion after [#mark()]. + public int workerReadsCompletedAfterMark() { + return workerReadsCompletedAfterMark.get(); + } + + /// True if the parked read resumed and finished after [#mark()]. + public boolean parkedReadResumedAfterMark() { + return parkedReadResumedAfterMark.get(); + } + + public Thread parkedThread() { + return parkedThread.get(); + } + + static boolean onWorkerThread() { + return Thread.currentThread().getName().startsWith(WORKER_PREFIX); + } + } + + /// Gate choreography for the post-F1 world: waits for a read to park, holds it briefly (so a + /// concurrently injected failure provably begins the unwind while the read is in flight), + /// then releases it. The F1 drain in `compact()` blocks on the parked read until this fires, + /// so the read completes strictly before `compact()` throws. + public static final class AutoReleaser implements Runnable { + private final Gate gate; + private final long holdMillis; + + public AutoReleaser(Gate gate, long holdMillis) { + this.gate = gate; + this.holdMillis = holdMillis; + } + + @Override + public void run() { + try { + if (gate.awaitParked(60, TimeUnit.SECONDS)) { + Thread.sleep(holdMillis); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + gate.releaseParked(); + } + } + } + + private final ReaderSupplier delegate; + private final Gate gate; + private final Role role; + + public GatedReaderSupplier(ReaderSupplier delegate, Gate gate, Role role) { + this.delegate = delegate; + this.gate = gate; + this.role = role; + } + + @Override + public RandomAccessReader get() throws IOException { + return new GatedReader(delegate.get()); + } + + @Override + public void close() throws IOException { + delegate.close(); + } + + private interface IoOp { + void run() throws IOException; + } + + private final class GatedReader implements RandomAccessReader { + private final RandomAccessReader in; + + GatedReader(RandomAccessReader in) { + this.in = in; + } + + private void gated(IoOp op) throws IOException { + if (!gate.armed || !Gate.onWorkerThread()) { + op.run(); + return; + } + gate.activeWorkerReads.incrementAndGet(); + try { + applyRole(); + op.run(); + if (gate.marked) { + gate.workerReadsCompletedAfterMark.incrementAndGet(); + if (Thread.currentThread() == gate.parkedThread.get()) { + gate.parkedReadResumedAfterMark.set(true); + } + } + } finally { + gate.activeWorkerReads.decrementAndGet(); + } + } + + private void applyRole() { + if (role == Role.PARK) { + if (gate.parkedThread.compareAndSet(null, Thread.currentThread())) { + gate.parked.countDown(); + try { + // Backstop only: the choreography (AutoReleaser) releases the gate long + // before this; with the F1 drain, compact() blocks until it does. + if (!gate.release.await(15, TimeUnit.SECONDS)) { + throw new IllegalStateException("repro gate was never released"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("parked repro read interrupted", e); + } + } + } else if (role == Role.POISON) { + if (gate.parked.getCount() == 0 + && Thread.currentThread() != gate.parkedThread.get() + && gate.poisonFired.compareAndSet(false, true)) { + throw new RuntimeException("REPRO_POISON: injected merge batch failure"); + } + } + } + + @Override + public void seek(long offset) throws IOException { + in.seek(offset); + } + + @Override + public long getPosition() throws IOException { + return in.getPosition(); + } + + @Override + public int readInt() throws IOException { + int[] out = new int[1]; + gated(() -> out[0] = in.readInt()); + return out[0]; + } + + @Override + public float readFloat() throws IOException { + float[] out = new float[1]; + gated(() -> out[0] = in.readFloat()); + return out[0]; + } + + @Override + public long readLong() throws IOException { + long[] out = new long[1]; + gated(() -> out[0] = in.readLong()); + return out[0]; + } + + @Override + public void readFully(byte[] bytes) throws IOException { + gated(() -> in.readFully(bytes)); + } + + @Override + public void readFully(ByteBuffer buffer) throws IOException { + gated(() -> in.readFully(buffer)); + } + + @Override + public void readFully(float[] floats) throws IOException { + gated(() -> in.readFully(floats)); + } + + @Override + public void readFully(long[] vector) throws IOException { + gated(() -> in.readFully(vector)); + } + + @Override + public void read(int[] ints, int offset, int count) throws IOException { + gated(() -> in.read(ints, offset, count)); + } + + @Override + public void read(float[] floats, int offset, int count) throws IOException { + gated(() -> in.read(floats, offset, count)); + } + + @Override + public void close() throws IOException { + in.close(); + } + + @Override + public long length() throws IOException { + return in.length(); + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/HostStyleMappedReader.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/HostStyleMappedReader.java new file mode 100644 index 000000000..5b2d28fba --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/HostStyleMappedReader.java @@ -0,0 +1,113 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.ByteBufferReader; +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.SimpleMappedReader; +import sun.misc.Unsafe; + +import java.io.IOException; +import java.io.RandomAccessFile; +import java.lang.reflect.Field; +import java.nio.ByteOrder; +import java.nio.FloatBuffer; +import java.nio.MappedByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Path; + +/// A [RandomAccessReader] over a raw NIO `mmap` that models how an embedding host (for example a +/// Cassandra `FileHandle`-based adapter) typically reads a jvector graph file. Two properties are +/// deliberate, because they are the properties the memory-safety theories hinge on: +/// +/// - Bulk float reads go through a byte-swapping [FloatBuffer#get(float[],int,int)] view (the file +/// format is big-endian, the hardware little-endian), so an access to invalidated pages faults +/// inside `Unsafe.copySwapMemory0` / `Copy::conjoint_swap` — the exact leaf frame in the +/// production crash. (jvector's own [ByteBufferReader] reads floats one element at a time, which +/// would fault in a different leaf.) +/// - [Supplier#close()] unmaps immediately via `sun.misc.Unsafe.invokeCleaner`, exactly like +/// jvector's shipped [SimpleMappedReader.Supplier#close()] and Cassandra's `FileUtils.clean`. +/// Unlike the `Arena`-managed mapping in jvector-native's `MemorySegmentReader`, there is no +/// liveness handshake: readers vended by [Supplier#get()] are left dangling and any subsequent +/// access faults natively instead of throwing. +public final class HostStyleMappedReader extends ByteBufferReader { + private final MappedByteBuffer mbb; + + private HostStyleMappedReader(MappedByteBuffer mbb) { + super(mbb); + this.mbb = mbb; + } + + /// Bulk read through the swapped [FloatBuffer] view so a fault lands in `Copy::conjoint_swap`. + @Override + public void read(float[] floats, int offset, int count) { + FloatBuffer fb = mbb.asFloatBuffer(); + fb.get(floats, offset, count); + mbb.position(mbb.position() + count * Float.BYTES); + } + + @Override + public void readFully(float[] floats) { + read(floats, 0, floats.length); + } + + /// Vends dangling-capable readers over one shared mapping; [#close()] performs a raw, + /// immediate munmap with no coordination with outstanding readers. + public static final class Supplier implements ReaderSupplier { + private static final Unsafe UNSAFE = loadUnsafe(); + private final MappedByteBuffer buffer; + + public Supplier(Path path) throws IOException { + try (RandomAccessFile raf = new RandomAccessFile(path.toString(), "r")) { + if (raf.length() > Integer.MAX_VALUE) { + throw new IOException("file too large for a single NIO mapping: " + path); + } + this.buffer = raf.getChannel().map(FileChannel.MapMode.READ_ONLY, 0, raf.length()); + this.buffer.order(ByteOrder.BIG_ENDIAN); + } + } + + @Override + public HostStyleMappedReader get() { + MappedByteBuffer dup = (MappedByteBuffer) buffer.duplicate(); + dup.order(ByteOrder.BIG_ENDIAN); + return new HostStyleMappedReader(dup); + } + + @Override + public void close() { + if (UNSAFE != null) { + try { + UNSAFE.invokeCleaner(buffer); + } catch (IllegalArgumentException e) { + // not a cleanable direct buffer; nothing to unmap + } + } + } + + private static Unsafe loadUnsafe() { + try { + Field f = Unsafe.class.getDeclaredField("theUnsafe"); + f.setAccessible(true); + return (Unsafe) f.get(null); + } catch (Exception e) { + return null; + } + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/IndexFileLifecycleCorruptionTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/IndexFileLifecycleCorruptionTest.java new file mode 100644 index 000000000..b0cb3acc6 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/IndexFileLifecycleCorruptionTest.java @@ -0,0 +1,314 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +import java.util.Arrays; +import java.util.List; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +/// Tests for the *silent data corruption* hazards around index-file ranging (offset / region +/// loads), closing, and files rewritten to the same path with a different length. +/// +/// The originally-enabling fact: `OnDiskGraphIndexCompactor.compact(Path, long)` never truncated +/// its destination, so a path that previously held a longer file kept a stale tail — and a +/// stale-but-valid v5+ footer there hijacked the footer-based default load silently (proven by +/// the pre-F2 form of t6a; evidence in `local/memory_safety_design_brief.md` T6a). Since the F2 +/// fix, standalone destinations (`startOffset == 0`) are truncated before writing, and t6a now +/// guards that round-trip. Embedded destinations (`startOffset > 0`) remain untruncated by +/// contract, which t6c characterizes: prefix preserved, stale tail persists, header-based region +/// loads round-trip, and a footer-based whole-container load fails loudly on a junk tail. +/// Separately, mmap-based readers have no defense against the file being rewritten in place +/// (same inode: live readers silently see the new bytes) or replaced by rename (old inode: live +/// readers silently keep serving stale data) — t6b documents those as permanent mmap/inode +/// physics for hosts to respect, not fixable jvector behavior. +public class IndexFileLifecycleCorruptionTest { + + private static final int DIM = 16; + private static final int SMALL = 30; // nodes per compaction source + + private static void verdict(String theory, String text) { + System.out.println("VERDICT|" + theory + "|" + text); + } + + /// Outcome of a load attempt plus enough context to classify silent-vs-loud failures. + private static final class LoadOutcome { + final OnDiskGraphIndex graph; + final Throwable error; + + LoadOutcome(OnDiskGraphIndex graph, Throwable error) { + this.graph = graph; + this.error = error; + } + + boolean loaded() { + return graph != null; + } + + int size() { + return graph.size(0); + } + + String describe() { + if (loaded()) { + return "loaded size(0)=" + size(); + } + return "threw " + error.getClass().getName() + ": " + error.getMessage(); + } + } + + private static LoadOutcome tryLoad(ReaderSupplier rs, long offset, boolean useFooter) { + try { + return new LoadOutcome(OnDiskGraphIndex.load(rs, offset, useFooter), null); + } catch (Throwable t) { + return new LoadOutcome(null, t); + } + } + + /// Compacts two fresh 30-node sources into `out` at `offset`; returns the source vectors for + /// content verification (compacted ordinals: source0 at [0,30), source1 at [30,60)). + private static List>> compactSmallSourcesInto(Path dir, Path out, long offset) throws Exception { + List> v0 = ReproGraphs.randomVectors(SMALL, DIM, 11); + List> v1 = ReproGraphs.randomVectors(SMALL, DIM, 13); + Path s0 = ReproGraphs.buildInlineGraph(dir.resolve("cmp_src0.graph"), v0, DIM); + Path s1 = ReproGraphs.buildInlineGraph(dir.resolve("cmp_src1.graph"), v1, DIM); + try (ReaderSupplier r0 = ReaderSupplierFactory.open(s0); + ReaderSupplier r1 = ReaderSupplierFactory.open(s1)) { + OnDiskGraphIndex g0 = OnDiskGraphIndex.load(r0); + OnDiskGraphIndex g1 = OnDiskGraphIndex.load(r1); + var compactor = new OnDiskGraphIndexCompactor( + List.of(g0, g1), + ReproGraphs.allLive(SMALL, SMALL), + ReproGraphs.stackedRemappers(SMALL, SMALL), + VectorSimilarityFunction.EUCLIDEAN, + null, -1); + if (offset == 0) { + compactor.compact(out); + } else { + compactor.compact(out, offset); + } + } + return List.of(v0, v1); + } + + private static boolean vectorMatches(OnDiskGraphIndex.View view, int ordinal, VectorFloat expected) { + return Arrays.equals(ReproGraphs.readVector(view, ordinal, DIM), ReproGraphs.toArray(expected)); + } + + /// F2 guard: compacting onto a path that previously held a *longer* index must truncate the + /// stale content, so the default (footer-based) load sees exactly the new graph. Pre-F2 this + /// test proved the inverse: the stale tail kept the old graph's still-valid footer at the + /// file end and `load()` silently resurrected the old structure (size 500, expected 60) over + /// hybrid old/new bytes, with no error anywhere. + @Test(timeout = 240_000) + public void t6a_compactOntoReusedLongerPathTruncatesAndRoundTrips() throws Exception { + Path dir = ReproGraphs.newWorkDir("t6a-stale-tail"); + + // Control: the same compaction into a fresh path round-trips correctly. + Path fresh = dir.resolve("fresh_out.graph"); + List>> vecs = compactSmallSourcesInto(dir, fresh, 0); + try (ReaderSupplier rs = ReaderSupplierFactory.open(fresh)) { + OnDiskGraphIndex g = OnDiskGraphIndex.load(rs); + assertEquals("control: fresh-path compaction must load correctly", 2 * SMALL, g.size(0)); + try (var view = g.getView()) { + assertTrue("control: source-0 vector content must round-trip", vectorMatches(view, 0, vecs.get(0).get(0))); + assertTrue("control: source-1 vector content must round-trip", vectorMatches(view, SMALL, vecs.get(1).get(0))); + } + } + + // The reused path: previously held a much larger graph. + Path reused = dir.resolve("reused_out.graph"); + List> oldVecs = ReproGraphs.randomVectors(500, DIM, 7); + ReproGraphs.buildInlineGraph(reused, oldVecs, DIM); + long oldSize = Files.size(reused); + + List>> newVecs = compactSmallSourcesInto(dir, reused, 0); + long newSize = Files.size(reused); + assertTrue("F2: a standalone destination must be truncated before writing" + + " (previous file was " + oldSize + " bytes, now " + newSize + ")", newSize < oldSize); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(reused)) { + OnDiskGraphIndex g = OnDiskGraphIndex.load(rs); + assertEquals("default load of the reused path must see exactly the new compaction", + 2 * SMALL, g.size(0)); + try (var view = g.getView()) { + assertTrue("source-0 vector content must round-trip on the reused path", + vectorMatches(view, 0, newVecs.get(0).get(0))); + assertTrue("source-1 vector content must round-trip on the reused path", + vectorMatches(view, SMALL, newVecs.get(1).get(0))); + } + } + verdict("T6a", "FIX HOLDS: reused longer path truncated (" + oldSize + " -> " + newSize + + " bytes); default load round-trips the new graph exactly"); + } + + /// T6b: rewriting a mapped index file in place (same path, same inode, same length) is + /// silently visible through a live reader — no exception, just different bytes. This is the + /// unified-page-cache half of the "same location, different content/length" hazard. + /// + /// Correct behavior for a host: never rewrite a graph component in place while any reader + /// (a running compaction's source view included) may be open over it. + @Test(timeout = 240_000) + public void t6b_inPlaceRewriteUnderLiveMappingSilentlyServesNewBytes() throws Exception { + Path dir = ReproGraphs.newWorkDir("t6b-inplace"); + Path p = dir.resolve("graph.bin"); + List> vecs = ReproGraphs.randomVectors(100, DIM, 21); + ReproGraphs.buildInlineGraph(p, vecs, DIM); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(p)) { + OnDiskGraphIndex g = OnDiskGraphIndex.load(rs); + try (var view = g.getView()) { + float[] before = ReproGraphs.readVector(view, 0, DIM); + assertArrayEquals("sanity: the live view reads the original bytes", + ReproGraphs.toArray(vecs.get(0)), before, 0.0f); + + byte[] all = Files.readAllBytes(p); + for (int i = 0; i < all.length; i++) { + all[i] ^= 0x5A; + } + try (FileChannel ch = FileChannel.open(p, StandardOpenOption.WRITE)) { + ByteBuffer buf = ByteBuffer.wrap(all); + long pos = 0; + while (buf.hasRemaining()) { + pos += ch.write(buf, pos); + } + } + assertEquals("same length after in-place rewrite", all.length, Files.size(p)); + + float[] after = ReproGraphs.readVector(view, 0, DIM); + assertFalse("T6b PROOF POINT: the live reader silently serves the rewritten bytes" + + " — no exception, different data", Arrays.equals(before, after)); + verdict("T6b", "PROVEN silent corruption: in-place rewrite is immediately visible through a live mmap view"); + } + } + } + + /// T6b (rename variant): atomically replacing the file leaves live readers on the old inode — + /// they silently keep serving the *stale* graph while new readers of the same path see the new + /// one. No error surfaces on either side of the split brain. + @Test(timeout = 240_000) + public void t6b_renameReplaceUnderLiveMappingSilentlyServesStaleBytes() throws Exception { + Path dir = ReproGraphs.newWorkDir("t6b-rename"); + Path p = dir.resolve("graph.bin"); + List> vecsA = ReproGraphs.randomVectors(100, DIM, 31); + ReproGraphs.buildInlineGraph(p, vecsA, DIM); + + try (ReaderSupplier live = ReaderSupplierFactory.open(p)) { + OnDiskGraphIndex g = OnDiskGraphIndex.load(live); + try (var view = g.getView()) { + float[] before = ReproGraphs.readVector(view, 0, DIM); + + List> vecsB = ReproGraphs.randomVectors(100, DIM, 32); + Path replacement = ReproGraphs.buildInlineGraph(dir.resolve("replacement.bin"), vecsB, DIM); + Files.move(replacement, p, StandardCopyOption.REPLACE_EXISTING); + + float[] liveRead = ReproGraphs.readVector(view, 0, DIM); + assertArrayEquals("T6b PROOF POINT: after rename-replace, the live reader silently serves STALE bytes", + before, liveRead, 0.0f); + + try (ReaderSupplier freshRs = ReaderSupplierFactory.open(p)) { + OnDiskGraphIndex fresh = OnDiskGraphIndex.load(freshRs); + float[] freshRead; + try (var freshView = fresh.getView()) { + freshRead = ReproGraphs.readVector(freshView, 0, DIM); + } + assertFalse("a fresh reader of the same path must see the replacement", + Arrays.equals(before, freshRead)); + } + verdict("T6b-rename", "PROVEN silent split-brain: live reader serves the old inode, fresh reader the new file"); + } + } + } + + /// T6c: the ranged/no-copy embedding path. `compact(container, startOffset)` into a reused + /// container that is *longer* than the new body: the reserved prefix must survive (documented + /// contract — positive check), the header-based region load must round-trip (positive check), + /// and the footer-based load — which trusts the file end — must not silently succeed over the + /// stale junk tail. + @Test(timeout = 240_000) + public void t6c_offsetRegionCompactionInReusedLongerContainer() throws Exception { + Path dir = ReproGraphs.newWorkDir("t6c-container"); + Path container = dir.resolve("container.bin"); + final int prefixLen = 4096; + final int junkLen = 64_000; + byte[] prefix = new byte[prefixLen]; + Arrays.fill(prefix, (byte) 0xAB); + byte[] junk = new byte[junkLen]; + Arrays.fill(junk, (byte) 0xCA); + Files.write(container, prefix); + Files.write(container, junk, StandardOpenOption.APPEND); + + List>> vecs = compactSmallSourcesInto(dir, container, prefixLen); + + byte[] prefixAfter = new byte[prefixLen]; + try (FileChannel ch = FileChannel.open(container, StandardOpenOption.READ)) { + ByteBuffer buf = ByteBuffer.wrap(prefixAfter); + long pos = 0; + while (buf.hasRemaining()) { + int n = ch.read(buf, pos); + if (n < 0) { + break; + } + pos += n; + } + } + assertArrayEquals("documented contract: the reserved prefix must survive compaction untouched", + prefix, prefixAfter); + assertEquals("enabling condition: the longer container is not truncated, the junk tail persists", + (long) prefixLen + junkLen, Files.size(container)); + + try (ReaderSupplier rs = ReaderSupplierFactory.open(container)) { + LoadOutcome headerLoad = tryLoad(rs, prefixLen, false); + assertNull("header-based region load must work on the reused container: " + + (headerLoad.error == null ? "" : headerLoad.error.toString()), headerLoad.error); + assertEquals("header-based region load must see exactly the compacted graph", + 2 * SMALL, headerLoad.size()); + try (var view = headerLoad.graph.getView()) { + assertTrue("region-loaded source-0 content must round-trip", vectorMatches(view, 0, vecs.get(0).get(0))); + assertTrue("region-loaded source-1 content must round-trip", vectorMatches(view, SMALL, vecs.get(1).get(0))); + } + } + + try (ReaderSupplier rs = ReaderSupplierFactory.open(container)) { + LoadOutcome footerLoad = tryLoad(rs, prefixLen, true); + assertTrue("footer-based load over a stale longer tail must not silently produce the compacted graph; got: " + + footerLoad.describe(), !footerLoad.loaded() || footerLoad.size() != 2 * SMALL); + verdict("T6c", footerLoad.loaded() + ? "PROVEN silent corruption: footer load returned a bogus graph: " + footerLoad.describe() + : "footer load over the junk tail fails loudly (not silently): " + footerLoad.describe()); + } + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyCrashReproTest.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyCrashReproTest.java new file mode 100644 index 000000000..606da261c --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyCrashReproTest.java @@ -0,0 +1,151 @@ +/* + * 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.example.repro; + +import org.junit.Assume; +import org.junit.Test; + +import java.nio.file.Path; +import java.util.Locale; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertTrue; + +/// Reader-level reproductions for the compaction memory-safety theories (see +/// `local/memory_safety.md` on the cooperative-embedding branch). Each scenario runs in a child +/// JVM via [ChildJvm] because the expected outcome of several of them is a native JVM crash; the +/// parent asserts on exit status, harness outcome markers, and the hs_err evidence (`siginfo` +/// si_code and the problematic frame — the discriminators the bug doc's trigger table needs). +/// +/// The matrix is reader-kind x invalidation-trigger: +/// +/// | theory | reader | trigger | expectation under test | +/// |--------|---------------------------------|----------------|----------------------------------------------| +/// | T1 | Arena (`MemorySegmentReader`) | clean close | Java exception, never a native crash | +/// | T2 | Arena (`MemorySegmentReader`) | truncate | catastrophic-or-loud (SIGBUS crash or error) | +/// | T2b | host-style NIO mmap | truncate | catastrophic-or-loud (SIGBUS crash or error) | +/// | T3 | host-style NIO mmap | clean unmap | native SIGSEGV in the swap-copy leaf | +/// | T3b | `SimpleMappedReader` (shipped) | supplier close | native crash (its close IS a raw unmap) | +/// +/// T1 validates the bug doc's §3 deduction for the arena reader. T3/T3b test the refinement that +/// the deduction does NOT extend to NIO-mapped readers: for those, a clean close of a source +/// while a read is in flight is itself sufficient to produce the production core dump — no file +/// truncation or rewrite required. +public class MemorySafetyCrashReproTest { + + private static void assumeLinux() { + Assume.assumeTrue("crash reproductions rely on Linux mmap semantics and hs_err parsing", + System.getProperty("os.name", "").toLowerCase(Locale.ROOT).contains("linux")); + } + + private static void assumeArenaReader() { + Assume.assumeTrue("jvector-native MemorySegmentReader (JDK 22+) not available on this classpath", + MemorySafetyReproHarness.arenaReaderAvailable()); + } + + private static ChildJvm.Result runReader(String label, String kind, String trigger) throws Exception { + Path dir = ReproGraphs.newWorkDir(label); + return ChildJvm.run(dir, 180, "reader", kind, trigger, dir.toString()); + } + + private static void verdict(String theory, String text) { + System.out.println("VERDICT|" + theory + "|" + text); + } + + /// T1: closing the Arena-managed supplier while reads are in flight must degrade to a Java + /// exception (or a refused close) — never a native crash. This is the safe half of the + /// reader matrix and validates the doc's "clean close cannot core-dump" deduction *for this + /// reader implementation*. + @Test(timeout = 240_000) + public void t1_arenaReaderCloseUnderLiveReadsIsSafe() throws Exception { + assumeLinux(); + assumeArenaReader(); + ChildJvm.Result r = runReader("t1-arena-close", "arena", "close"); + + assertFalse("T1 DISPROVEN: clean Arena close crashed natively: " + r.summary(), r.crashed()); + assertTrue("arena close must stop readers with an exception or refuse to close: " + r.summary(), + r.exitCode == 0 || r.exitCode == 43); + verdict("T1", "PROVEN safe: " + r.summary()); + } + + /// T2: truncating the file under the Arena reader's live byte-swapping reads. The doc's §5a + /// predicts a SIGBUS crash in `Copy::conjoint_swap`; modern JDKs may instead rescue the fault + /// into `java.lang.InternalError`. Either way the read must not silently keep succeeding — + /// which outcome actually occurs decides how the doc's trigger table should read for this + /// reader, so the verdict line records it. + @Test(timeout = 240_000) + public void t2_arenaReaderTruncationIsCatastrophicOrLoud() throws Exception { + assumeLinux(); + assumeArenaReader(); + ChildJvm.Result r = runReader("t2-arena-truncate", "arena", "truncate"); + + assertNotEquals("T2 DISPROVEN: truncation had no effect on live reads: " + r.summary(), 42, r.exitCode); + boolean sigbusCrash = r.crashed() && r.hsErrContains("SIGBUS"); + boolean survivedWithError = !r.crashed() && r.exitCode == 0; + assertTrue("truncation under live mmap reads must crash with SIGBUS or surface a Java error: " + r.summary(), + sigbusCrash || survivedWithError); + verdict("T2", (sigbusCrash ? "SIGBUS crash (doc §5a expectation holds): " : "no crash — fault surfaced as a Java error (doc §5a needs revision for this JDK): ") + + r.summary()); + } + + /// T2b: the same truncation, through the host-style NIO reader. + @Test(timeout = 240_000) + public void t2b_hostStyleReaderTruncationIsCatastrophicOrLoud() throws Exception { + assumeLinux(); + ChildJvm.Result r = runReader("t2b-hoststyle-truncate", "hoststyle", "truncate"); + + assertNotEquals("T2b DISPROVEN: truncation had no effect on live reads: " + r.summary(), 42, r.exitCode); + boolean sigbusCrash = r.crashed() && r.hsErrContains("SIGBUS"); + boolean survivedWithError = !r.crashed() && r.exitCode == 0; + assertTrue("truncation under live NIO mmap reads must crash with SIGBUS or surface a Java error: " + r.summary(), + sigbusCrash || survivedWithError); + verdict("T2b", (sigbusCrash ? "SIGBUS crash: " : "no crash — fault surfaced as a Java error: ") + r.summary()); + } + + /// T3: a *clean close* (raw munmap, no handshake) of the host-style NIO mapping under live + /// byte-swapping reads must crash natively with SIGSEGV in the same swap-copy leaf as the + /// production core dump. This is the refinement of the bug doc: for host-style readers, early + /// release of a source is sufficient — no truncation or in-place rewrite is required. + @Test(timeout = 240_000) + public void t3_hostStyleReaderCleanUnmapCrashesNatively() throws Exception { + assumeLinux(); + ChildJvm.Result r = runReader("t3-hoststyle-close", "hoststyle", "close"); + + assertTrue("T3 DISPROVEN: clean unmap under live reads did not crash: " + r.summary(), r.crashed()); + assertTrue("expected SIGSEGV (unmapped pages), got: " + r.siginfo() + " / " + r.summary(), + r.hsErrContains("SIGSEGV")); + assertTrue("expected the byte-swapping copy leaf in the crash evidence: frame=" + r.problematicFrame(), + r.hsErrContains("copySwapMemory") || r.hsErrContains("conjoint_swap") || r.hsErrContains("FloatBufferS")); + verdict("T3", "PROVEN: clean unmap of a host-style mapping crashes natively: " + r.summary()); + } + + /// T3b: jvector's own shipped [io.github.jbellis.jvector.disk.SimpleMappedReader] has the same + /// property — its `Supplier.close()` is `Unsafe.invokeCleaner`, a raw munmap. Any caller that + /// closes it while a straggler read is in flight gets a native crash, not an exception. + @Test(timeout = 240_000) + public void t3b_simpleMappedReaderSupplierCloseCrashesNatively() throws Exception { + assumeLinux(); + ChildJvm.Result r = runReader("t3b-simplemapped-close", "simplemapped", "close"); + + assertTrue("T3b DISPROVEN: SimpleMappedReader.Supplier.close() under live reads did not crash: " + r.summary(), + r.crashed()); + assertTrue("expected a memory fault signal, got: " + r.siginfo(), + r.hsErrContains("SIGSEGV") || r.hsErrContains("SIGBUS")); + verdict("T3b", "PROVEN: shipped SimpleMappedReader close is a raw unmap and crashes live readers: " + r.summary()); + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyReproHarness.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyReproHarness.java new file mode 100644 index 000000000..0a34e405a --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/MemorySafetyReproHarness.java @@ -0,0 +1,330 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.SimpleMappedReader; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; + +import java.lang.reflect.Constructor; +import java.nio.channels.FileChannel; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Random; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +/// Child-JVM entry point for the memory-safety reproductions. Every scenario that may end in a +/// native JVM crash runs here, launched by [ChildJvm] from the actual tests, so the crash kills +/// this process instead of the test suite. +/// +/// Reader-level scenarios (`reader `) spin reader threads over a raw +/// big-endian float file through one of three [ReaderSupplier] implementations, then invalidate +/// the mapping underneath them: +/// +/// - kind `arena`: jvector-native's `MemorySegmentReader` (Arena-managed mapping, loaded +/// reflectively exactly like `ReaderSupplierFactory` does) +/// - kind `hoststyle`: [HostStyleMappedReader] (raw NIO mapping, bulk byte-swapping reads, no +/// liveness handshake — models a Cassandra `FileHandle`-style adapter) +/// - kind `simplemapped`: jvector's shipped [SimpleMappedReader] (raw NIO mapping, per-element +/// reads, `invokeCleaner` on supplier close) +/// +/// and trigger `close` (clean supplier close / unmap) or `truncate` (shrink the file in place). +/// +/// The compactor scenario (`compactor-straggler `) exercises the full production stack: +/// a cross-source merge on a worker pool whose batches read sources through [GatedReaderSupplier]; +/// one batch is poisoned so `compact()` throws, and a contract-compliant "host" then closes the +/// source suppliers (raw unmap). Before the F1 drain-on-unwind fix this crashed the JVM (the +/// parked straggler read resumed into the dead mapping — SIGSEGV/SEGV_MAPERR, see +/// `local/memory_safety_design_brief.md` T4b); with the fix, `compact()` drains the parked read +/// before throwing and the scenario must end at `OUTCOME|NO_CRASH` (exit 45). +/// +/// Exit codes: 0 = readers stopped with a Java exception; 42 = trigger had no effect; +/// 43 = close refused with an exception while reads continued; 44 = compact() returned normally; +/// 45 = no crash after sources closed; 99 = harness failure. A native crash never reaches an +/// orderly exit — the parent detects it via the exit status and the hs_err file. +public final class MemorySafetyReproHarness { + private static final String ARENA_SUPPLIER_CLASS = "io.github.jbellis.jvector.disk.MemorySegmentReader$Supplier"; + + private MemorySafetyReproHarness() { + } + + public static void main(String[] args) { + try { + String mode = args[0]; + int code; + if ("reader".equals(mode)) { + code = readerScenario(args[1], args[2], Path.of(args[3])); + } else if ("compactor-straggler".equals(mode)) { + code = compactorStragglerScenario(Path.of(args[1])); + } else { + throw new IllegalArgumentException("unknown mode: " + mode); + } + System.exit(code); + } catch (Throwable t) { + t.printStackTrace(System.out); + stage("FATAL|" + t); + System.exit(99); + } + } + + private static void stage(String message) { + System.out.println("REPRO|" + message); + System.out.flush(); + } + + /// True when the arena-based reader (jvector-native, JDK 22+) is loadable in this JVM. + public static boolean arenaReaderAvailable() { + try { + Class.forName(ARENA_SUPPLIER_CLASS); + return true; + } catch (Throwable t) { + return false; + } + } + + static ReaderSupplier openSupplier(String kind, Path file) throws Exception { + switch (kind) { + case "arena": { + Constructor ctor = Class.forName(ARENA_SUPPLIER_CLASS).getConstructor(Path.class); + return (ReaderSupplier) ctor.newInstance(file); + } + case "hoststyle": + return new HostStyleMappedReader.Supplier(file); + case "simplemapped": + return new SimpleMappedReader.Supplier(file); + default: + throw new IllegalArgumentException("unknown reader kind: " + kind); + } + } + + // ---------------------------------------------------------------- reader-level scenarios + + private static int readerScenario(String kind, String trigger, Path workDir) throws Exception { + final int dimension = 64; + final int count = 100_000; + Path file = ReproGraphs.writeBigEndianFloats(workDir.resolve("floats.bin"), count, dimension, 42L); + ReaderSupplier supplier = openSupplier(kind, file); + + final AtomicLong reads = new AtomicLong(); + final List readerErrors = Collections.synchronizedList(new ArrayList<>()); + final int readerCount = 2; + final CountDownLatch stopped = new CountDownLatch(readerCount); + + for (int i = 0; i < readerCount; i++) { + final long seed = 1000 + i; + Thread t = new Thread(new ReaderLoop(supplier, count, dimension, seed, reads, readerErrors, stopped), + "repro-reader-" + i); + t.setDaemon(true); + t.start(); + } + + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30); + while (reads.get() < 30_000 && System.nanoTime() < deadline && stopped.getCount() == readerCount) { + Thread.sleep(10); + } + stage("WARMED|reads=" + reads.get()); + + Throwable closeError = null; + if ("close".equals(trigger)) { + try { + supplier.close(); + } catch (Throwable t) { + closeError = t; + } + } else if ("truncate".equals(trigger)) { + try (FileChannel ch = FileChannel.open(file, StandardOpenOption.WRITE)) { + ch.truncate(4096); + } + } else { + throw new IllegalArgumentException("unknown trigger: " + trigger); + } + stage("TRIGGERED|" + trigger + (closeError == null ? "" : "|closeThrew=" + closeError.getClass().getName())); + + boolean allStopped = stopped.await(8, TimeUnit.SECONDS); + if (allStopped) { + Throwable first = readerErrors.isEmpty() ? null : readerErrors.get(0); + stage("OUTCOME|SURVIVED_EXCEPTION|" + describe(first)); + return 0; + } + long before = reads.get(); + Thread.sleep(300); + boolean stillFlowing = reads.get() > before; + if (closeError != null) { + stage("OUTCOME|CLOSE_REFUSED|" + describe(closeError) + "|readsStillFlowing=" + stillFlowing); + return 43; + } + stage("OUTCOME|NO_EFFECT|readsStillFlowing=" + stillFlowing); + return 42; + } + + /// The hot read loop: random-position bulk float reads, the access pattern of a graph search. + private static final class ReaderLoop implements Runnable { + private final ReaderSupplier supplier; + private final int count; + private final int dimension; + private final long seed; + private final AtomicLong reads; + private final List errors; + private final CountDownLatch stopped; + + ReaderLoop(ReaderSupplier supplier, int count, int dimension, long seed, + AtomicLong reads, List errors, CountDownLatch stopped) { + this.supplier = supplier; + this.count = count; + this.dimension = dimension; + this.seed = seed; + this.reads = reads; + this.errors = errors; + this.stopped = stopped; + } + + @Override + public void run() { + try { + RandomAccessReader reader = supplier.get(); + float[] dst = new float[dimension]; + Random rnd = new Random(seed); + while (true) { + reader.seek((long) rnd.nextInt(count) * dimension * Float.BYTES); + reader.read(dst, 0, dimension); + reads.incrementAndGet(); + } + } catch (Throwable t) { + errors.add(t); + } finally { + stopped.countDown(); + } + } + } + + private static String describe(Throwable t) { + if (t == null) { + return "none"; + } + String msg = String.valueOf(t.getMessage()); + return t.getClass().getName() + ": " + (msg.length() > 160 ? msg.substring(0, 160) : msg); + } + + // ---------------------------------------------------------------- compactor-level scenario + + /// Exercises the production crash sequence end to end. Sequence markers narrate each step so + /// the parent can assert ordering around COMPACT_THREW and SOURCES_CLOSED — the actions of a + /// host that followed the documented lifecycle contract ("keep sources alive until compact() + /// returns or throws"). Pre-F1 this crashed after SOURCES_CLOSED; post-F1 the drain leaves + /// nothing in flight at the throw and the scenario must survive to exit 45. + private static int compactorStragglerScenario(Path workDir) throws Exception { + final int dimension = 24; + final int perSource = 220; + + Path s0 = ReproGraphs.buildInlineGraph(workDir.resolve("src0.graph"), + ReproGraphs.randomVectors(perSource, dimension, 101), dimension); + Path s1 = ReproGraphs.buildInlineGraph(workDir.resolve("src1.graph"), + ReproGraphs.randomVectors(perSource, dimension, 202), dimension); + stage("STAGE|SOURCES_BUILT"); + + HostStyleMappedReader.Supplier h0 = new HostStyleMappedReader.Supplier(s0); + HostStyleMappedReader.Supplier h1 = new HostStyleMappedReader.Supplier(s1); + GatedReaderSupplier.Gate gate = new GatedReaderSupplier.Gate(); + GatedReaderSupplier g0 = new GatedReaderSupplier(h0, gate, GatedReaderSupplier.Role.POISON); + GatedReaderSupplier g1 = new GatedReaderSupplier(h1, gate, GatedReaderSupplier.Role.PARK); + + OnDiskGraphIndex src0 = OnDiskGraphIndex.load(g0); + OnDiskGraphIndex src1 = OnDiskGraphIndex.load(g1); + + ExecutorService pool = Executors.newFixedThreadPool(4, new WorkerThreadFactory()); + var compactor = new OnDiskGraphIndexCompactor( + List.of(src0, src1), + ReproGraphs.allLive(perSource, perSource), + ReproGraphs.stackedRemappers(perSource, perSource), + VectorSimilarityFunction.EUCLIDEAN, + pool, 4); + + gate.arm(); + // Releases the parked read ~1s after it parks: long enough that the poison provably + // fires (and the unwind begins) while the read is in flight, short enough that the F1 + // drain — which blocks compact() on the parked read — completes promptly. + Thread releaser = new Thread(new GatedReaderSupplier.AutoReleaser(gate, 1_000), "repro-gate-releaser"); + releaser.setDaemon(true); + releaser.start(); + stage("STAGE|COMPACT_START"); + try { + compactor.compact(workDir.resolve("out.graph")); + stage("OUTCOME|COMPACT_RETURNED_NORMALLY"); + return 44; + } catch (Throwable t) { + gate.mark(); + stage("STAGE|COMPACT_THREW|" + rootChain(t)); + } + + stage("STAGE|PARKED=" + (gate.parkedThread() != null) + "|ACTIVE_WORKER_READS=" + gate.activeWorkerReads()); + + // The contract-compliant host: compact() has thrown, so the caller now releases its + // sources. For a FileHandle-style mapping that release is a raw munmap. + g0.close(); + g1.close(); + stage("STAGE|SOURCES_CLOSED"); + + gate.releaseParked(); + stage("STAGE|GATE_RELEASED"); + + // Pre-F1 the abandoned straggler resumed into the unmapped pages here and the JVM died + // within this window; with the drain in place nothing is in flight and it must pass + // quietly. + for (int i = 0; i < 5; i++) { + Thread.sleep(1_000); + } + stage("OUTCOME|NO_CRASH|activeWorkerReads=" + gate.activeWorkerReads() + + "|readsCompletedAfterThrow=" + gate.workerReadsCompletedAfterMark()); + return 45; + } + + /// Names pool threads with the prefix [GatedReaderSupplier.Gate#WORKER_PREFIX] so the gates + /// only engage on merge workers, never on the orchestrator or the graph-loading main thread. + private static final class WorkerThreadFactory implements ThreadFactory { + private final AtomicInteger n = new AtomicInteger(); + + @Override + public Thread newThread(Runnable r) { + Thread t = new Thread(r, GatedReaderSupplier.Gate.WORKER_PREFIX + n.getAndIncrement()); + t.setDaemon(true); + return t; + } + } + + private static String rootChain(Throwable t) { + StringBuilder sb = new StringBuilder(); + for (Throwable c = t; c != null; c = c.getCause()) { + if (sb.length() > 0) { + sb.append(" <- "); + } + sb.append(c.getClass().getSimpleName()); + } + return sb.toString(); + } +} diff --git a/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ReproGraphs.java b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ReproGraphs.java new file mode 100644 index 000000000..936a21195 --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/ReproGraphs.java @@ -0,0 +1,222 @@ +/* + * 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.example.repro; + +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OrdinalMapper; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.quantization.PQVectors; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; + +import java.io.BufferedOutputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.concurrent.ForkJoinPool; +import java.util.function.IntFunction; + +import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; + +/// Shared fixtures for the memory-safety reproduction tests: tiny on-disk graphs with the +/// `INLINE_VECTORS` feature (the shape [io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor] +/// requires of its sources), raw big-endian float files for reader-level scenarios, and the +/// live-nodes / remapper boilerplate the compactor constructor needs. Graph construction mirrors +/// `TestOnDiskGraphIndexCompactor.buildSimpleSourceGraph` in jvector-tests. +public final class ReproGraphs { + public static final VectorTypeSupport VTS = VectorizationProvider.getInstance().getVectorTypeSupport(); + + private ReproGraphs() { + } + + /// Creates a fresh scratch directory under the build's `target/` tree (never `/tmp`), so the + /// build owns cleanup and tests never need to delete anything. + public static Path newWorkDir(String label) throws IOException { + Path moduleTarget; + if (Files.isDirectory(Path.of("jvector-examples", "target"))) { + moduleTarget = Path.of("jvector-examples", "target"); // CWD = repo root (surefire config) + } else if (Files.isDirectory(Path.of("target"))) { + moduleTarget = Path.of("target"); // CWD = module dir + } else { + moduleTarget = Path.of(System.getProperty("java.io.tmpdir")); + } + Path base = moduleTarget.resolve("repro-tmp"); + Files.createDirectories(base); + return Files.createTempDirectory(base, label + "-"); + } + + /// Deterministic random vectors. + public static List> randomVectors(int count, int dimension, long seed) { + Random r = new Random(seed); + List> out = new ArrayList<>(count); + for (int i = 0; i < count; i++) { + VectorFloat v = VTS.createFloatVector(dimension); + for (int d = 0; d < dimension; d++) { + v.set(d, r.nextFloat()); + } + out.add(v); + } + return out; + } + + /// Builds a single-layer graph with inline full-resolution vectors and writes it to + /// `outputPath` (v5+ format, footer at the logical end of the file). + public static Path buildInlineGraph(Path outputPath, List> vecs, int dimension) throws IOException { + RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(vecs, dimension); + var bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, io.github.jbellis.jvector.vector.VectorSimilarityFunction.EUCLIDEAN); + var builder = new GraphIndexBuilder(bsp, dimension, 8, 32, 1.2f, 1.2f, false, true, + ForkJoinPool.commonPool(), ForkJoinPool.commonPool()); + for (int i = 0; i < vecs.size(); i++) { + builder.addGraphNode(i, vecs.get(i)); + } + builder.cleanup(); + var graph = builder.getGraph(); + + var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withMapper(new OrdinalMapper.IdentityMapper(vecs.size() - 1)) + .with(new InlineVectors(dimension)); + var writer = writerBuilder.build(); + + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + writeSuppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravv.getVector(ordinal))); + for (int node = 0; node < vecs.size(); node++) { + var stateMap = new EnumMap(FeatureId.class); + stateMap.put(FeatureId.INLINE_VECTORS, writeSuppliers.get(FeatureId.INLINE_VECTORS).apply(node)); + writer.writeInline(node, stateMap); + } + writer.write(writeSuppliers); + return outputPath; + } + + /// Builds a single-layer graph with inline full-resolution vectors AND the `FUSED_PQ` + /// feature — the source shape that makes the compactor pick `FusedCompactionStrategy` and + /// therefore run the pre-encode `invokeAll` fan-out. Mirrors + /// `TestOnDiskGraphIndexCompactor.buildFusedPQ` in jvector-tests. + public static Path buildFusedGraph(Path outputPath, List> vecs, int dimension) throws IOException { + RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(vecs, dimension); + ProductQuantization pq = ProductQuantization.compute(ravv, 8, 256, true, UNWEIGHTED, + ForkJoinPool.commonPool(), ForkJoinPool.commonPool()); + PQVectors pqv = (PQVectors) pq.encodeAll(ravv, ForkJoinPool.commonPool()); + var bsp = BuildScoreProvider.pqBuildScoreProvider(io.github.jbellis.jvector.vector.VectorSimilarityFunction.COSINE, pqv); + var builder = new GraphIndexBuilder(bsp, dimension, 16, 100, 1.2f, 1.2f, false, true, + ForkJoinPool.commonPool(), ForkJoinPool.commonPool()); + var graph = builder.getGraph(); + + var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withMapper(new OrdinalMapper.IdentityMapper(vecs.size() - 1)) + .with(new InlineVectors(dimension)) + .with(new FusedPQ(graph.maxDegree(), pq)); + var writer = writerBuilder.build(); + + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + writeSuppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravv.getVector(ordinal))); + for (int node = 0; node < ravv.size(); node++) { + var stateMap = new EnumMap(FeatureId.class); + stateMap.put(FeatureId.INLINE_VECTORS, new InlineVectors.State(ravv.getVector(node))); + writer.writeInline(node, stateMap); + builder.addGraphNode(node, ravv.getVector(node)); + } + builder.cleanup(); + + writeSuppliers.put(FeatureId.FUSED_PQ, ordinal -> new FusedPQ.State(graph.getView(), pqv, ordinal)); + writer.write(writeSuppliers); + return outputPath; + } + + /// A raw file of `count * dimension` big-endian floats — the shape a reader-level scenario + /// scans without any graph structure on top. + public static Path writeBigEndianFloats(Path path, int count, int dimension, long seed) throws IOException { + Random r = new Random(seed); + try (DataOutputStream out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(path), 1 << 20))) { + for (long i = 0; i < (long) count * dimension; i++) { + out.writeFloat(r.nextFloat()); + } + } + return path; + } + + /// All-live bitsets, one per source size. + public static List allLive(int... sizes) { + List out = new ArrayList<>(sizes.length); + for (int size : sizes) { + FixedBitSet bits = new FixedBitSet(size); + bits.set(0, size); + out.add(bits); + } + return out; + } + + /// Identity remappers that stack each source's ordinals after the previous source's + /// (source 0 keeps `[0, n0)`, source 1 gets `[n0, n0+n1)`, ...). + public static List stackedRemappers(int... sizes) { + List out = new ArrayList<>(sizes.length); + int base = 0; + for (int size : sizes) { + Map map = new HashMap<>(); + for (int i = 0; i < size; i++) { + map.put(i, base + i); + } + out.add(new OrdinalMapper.MapMapper(map)); + base += size; + } + return out; + } + + /// Copies one full-resolution vector out of a graph view as a plain float array. + public static float[] readVector(OnDiskGraphIndex.View view, int ordinal, int dimension) { + VectorFloat buf = VTS.createFloatVector(dimension); + view.getVectorInto(ordinal, buf, 0); + float[] out = new float[dimension]; + for (int d = 0; d < dimension; d++) { + out[d] = buf.get(d); + } + return out; + } + + /// Converts a [VectorFloat] to a plain float array for comparisons. + public static float[] toArray(VectorFloat v) { + float[] out = new float[v.length()]; + for (int d = 0; d < v.length(); d++) { + out[d] = v.get(d); + } + return out; + } + + /// Loads a graph over the supplier, reflectively matching how production callers hold sources. + public static OnDiskGraphIndex load(ReaderSupplier rs) { + return OnDiskGraphIndex.load(rs); + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexViewGuards.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexViewGuards.java new file mode 100644 index 000000000..caef4affb --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexViewGuards.java @@ -0,0 +1,189 @@ +/* + * 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.disk; + +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.GraphIndexBuilder; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +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.After; +import org.junit.Before; +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardOpenOption; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ForkJoinPool; +import java.util.function.IntFunction; + +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/// Guard tests for the F3 hardening of `OnDiskGraphIndex.View` (prescription: +/// `local/memory_safety_fix_plan.md` F3): out-of-range node ordinals must fail with +/// `IllegalArgumentException` at the offset-computation entry points instead of becoming silent +/// wild offsets into the mapped file, and a corrupt (or stale-metadata) on-disk neighbor degree +/// must fail with a stack-bearing `IllegalStateException` at first touch instead of passing an +/// assert-only guard in production (`-da`) runs and turning into garbage node ids downstream. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class TestOnDiskGraphIndexViewGuards extends RandomizedTest { + private static final VectorTypeSupport VTS = VectorizationProvider.getInstance().getVectorTypeSupport(); + private static final int DIM = 4; + private static final int N = 8; + + private Path testDirectory; + private Path graphPath; + + @Before + public void setup() throws IOException { + testDirectory = Files.createTempDirectory("jvector_test"); + graphPath = buildInlineGraph(testDirectory.resolve("guards.graph"), TestUtil.createRandomVectors(N, DIM)); + } + + @After + public void tearDown() { + TestUtil.deleteQuietly(testDirectory); + } + + /// Single-layer graph with inline vectors, mirroring + /// `TestOnDiskGraphIndexCompactor.buildSimpleSourceGraph`. + private Path buildInlineGraph(Path outputPath, List> vecs) throws IOException { + RandomAccessVectorValues ravv = new ListRandomAccessVectorValues(vecs, DIM); + var bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, VectorSimilarityFunction.EUCLIDEAN); + var builder = new GraphIndexBuilder(bsp, DIM, 4, 20, 1.2f, 1.2f, false, true, + ForkJoinPool.commonPool(), ForkJoinPool.commonPool()); + for (int i = 0; i < vecs.size(); i++) { + builder.addGraphNode(i, vecs.get(i)); + } + builder.cleanup(); + var graph = builder.getGraph(); + + var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, outputPath) + .withMapper(new OrdinalMapper.IdentityMapper(vecs.size() - 1)) + .with(new InlineVectors(DIM)); + var writer = writerBuilder.build(); + Map> writeSuppliers = new EnumMap<>(FeatureId.class); + writeSuppliers.put(FeatureId.INLINE_VECTORS, ordinal -> new InlineVectors.State(ravv.getVector(ordinal))); + for (int node = 0; node < vecs.size(); node++) { + var stateMap = new EnumMap(FeatureId.class); + stateMap.put(FeatureId.INLINE_VECTORS, writeSuppliers.get(FeatureId.INLINE_VECTORS).apply(node)); + writer.writeInline(node, stateMap); + } + writer.write(writeSuppliers); + return outputPath; + } + + private static void expectOutOfRange(Runnable access) { + try { + access.run(); + fail("expected IllegalArgumentException for an out-of-range node ordinal"); + } catch (IllegalArgumentException e) { + assertTrue("message must identify the bad ordinal and the valid range, was: " + e.getMessage(), + e.getMessage().contains("out of range")); + } + } + + private static void expectCorruptDegree(Runnable access) { + try { + access.run(); + fail("expected IllegalStateException for a corrupt on-disk degree"); + } catch (IllegalStateException e) { + assertTrue("message must identify the corrupt degree, was: " + e.getMessage(), + e.getMessage().contains("degree")); + } + } + + /// Out-of-range ordinals — negative and one-past-the-end — must throw at both record-access + /// entry points (vector reads and neighbor reads), while in-range accesses keep working. + @Test + public void outOfRangeNodeOrdinalsThrow() throws Exception { + try (ReaderSupplier rs = ReaderSupplierFactory.open(graphPath)) { + var graph = OnDiskGraphIndex.load(rs); + try (var view = graph.getView()) { + VectorFloat buf = VTS.createFloatVector(DIM); + expectOutOfRange(() -> view.getVectorInto(-1, buf, 0)); + expectOutOfRange(() -> view.getVectorInto(N, buf, 0)); + expectOutOfRange(() -> view.getNeighborsIterator(0, -1)); + expectOutOfRange(() -> view.getNeighborsIterator(0, N)); + + // in-range accesses are unaffected by the guard + view.getVectorInto(0, buf, 0); + assertNotNull(view.getNeighborsIterator(0, N - 1)); + } + } + } + + /// A degree field patched to garbage — huge or negative — must throw at first touch. The + /// degree offsets are located through the package-private `neighborsOffsetFor` before the + /// file is patched on disk (big-endian, matching the format) and reopened. + @Test + public void corruptOnDiskDegreeThrows() throws Exception { + long hugeDegreeOffset; + long negativeDegreeOffset; + try (ReaderSupplier rs = ReaderSupplierFactory.open(graphPath)) { + var graph = OnDiskGraphIndex.load(rs); + try (var view = graph.getView()) { + hugeDegreeOffset = view.neighborsOffsetFor(0, 2); + negativeDegreeOffset = view.neighborsOffsetFor(0, 5); + } + } + + try (FileChannel ch = FileChannel.open(graphPath, StandardOpenOption.WRITE)) { + patchInt(ch, hugeDegreeOffset, Integer.MAX_VALUE / 2); + patchInt(ch, negativeDegreeOffset, -17); + } + + try (ReaderSupplier rs = ReaderSupplierFactory.open(graphPath)) { + var graph = OnDiskGraphIndex.load(rs); + try (var view = graph.getView()) { + expectCorruptDegree(() -> view.getNeighborsIterator(0, 2)); + expectCorruptDegree(() -> view.getNeighborsIterator(0, 5)); + + // an unpatched node still reads fine + assertNotNull(view.getNeighborsIterator(0, 0)); + } + } + } + + private static void patchInt(FileChannel ch, long offset, int value) throws IOException { + ByteBuffer bb = ByteBuffer.allocate(Integer.BYTES); // big-endian by default, matching the on-disk format + bb.putInt(value); + bb.flip(); + while (bb.hasRemaining()) { + offset += ch.write(bb, offset); + } + } +} From 0a63642e4349f20575a197575ed2e4310e39edbc Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 8 Jul 2026 05:33:47 +0000 Subject: [PATCH 10/13] EmbeddedExecutionContext: pass the pool once, no work escapes to default pools Facade carrying an embedder's compute/IO executors so build, PQ/NVQ, and compaction all run on it; also closes the PQRetrainer leak that hardcoded the default pools. --- .../graph/EmbeddedExecutionContext.java | 199 +++++++++++++ .../jvector/graph/disk/CompactionContext.java | 25 ++ .../graph/disk/OnDiskGraphIndexCompactor.java | 8 +- .../jvector/graph/disk/PQRetrainer.java | 25 +- .../jvector/graph/disk/feature/FusedPQ.java | 2 +- .../jvector/quantization/PQVectors.java | 2 +- .../graph/EmbeddedExecutionContextTest.java | 261 ++++++++++++++++++ 7 files changed, 516 insertions(+), 6 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java new file mode 100644 index 000000000..63c314769 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java @@ -0,0 +1,199 @@ +/* + * 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.annotations.Experimental; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor; +import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OrdinalMapper; +import io.github.jbellis.jvector.graph.disk.PQRetrainer; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.quantization.NVQVectors; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.PQVectors; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; + +import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; + +/** + * A single carrier for the execution resources an embedder supplies to jvector, so the pool is + * passed once and the "no parallel work escapes to {@link io.github.jbellis.jvector.util.PhysicalCoreExecutor#pool()} + * or {@link ForkJoinPool#commonPool()}" guarantee lives in one place. Construct the context + * with a bounded pool, then obtain builders/compactors/writers and run quantization through it; + * every operation routes only through the carried executors and never through a default + * overload. + * + *

Two executors. Graph and PQ/NVQ work is CPU-bound and runs on {@code compute}; the + * optional parallel graph write is IO-bound and runs on {@code io}, which defaults to + * {@code compute}. Passing one pool via {@link #of(ForkJoinPool)} is the simple, correct default; + * an embedder that enables parallel writes can supply a separate IO pool without touching any + * other call site. + * + *

Why a {@code ForkJoinPool}, not a {@link ParallelExecutor}. Product/NV quantization is + * expressed with the parallel-stream-on-a-pool idiom, which only confines work when the pool is a + * {@code ForkJoinPool}. The context therefore carries a concrete compute {@code ForkJoinPool} so it + * can serve PQ (needs a {@code ForkJoinPool}), the graph builder (via + * {@link ParallelExecutor#forkJoin}), and the compactor (via {@link java.util.concurrent.Executor}) + * uniformly. A caller-runs, single-threaded mode is not offered here — PQ requires a real + * pool; for a build-only caller-runs graph, use {@link ParallelExecutor#callerRuns()} on + * {@link GraphIndexBuilder} directly. + * + *

Lifecycle. The context neither owns nor shuts down the embedder's pool(s); the embedder + * supplies and disposes them. + * + *

Throttling is per-operation, not stored here. The compaction merge also accepts a + * {@link io.github.jbellis.jvector.util.work.ProgressLimiter} (progress + throttle), but that is + * scoped to a single host operation. Set it on the compactor returned by {@link #newCompactor} + * per call — do not expect the context to carry it. + */ +@Experimental +public final class EmbeddedExecutionContext { + private final ForkJoinPool compute; + private final ExecutorService io; + private final ParallelExecutor computeParallel; + + /** + * @param compute the pool for all CPU-bound build / cleanup / merge / PQ / NVQ work + * @param io the pool for the IO-bound parallel graph writer; {@code null} defaults to + * {@code compute} + */ + public EmbeddedExecutionContext(ForkJoinPool compute, ExecutorService io) { + this.compute = Objects.requireNonNull(compute, "compute"); + this.io = (io != null) ? io : compute; + this.computeParallel = ParallelExecutor.forkJoin(compute); + } + + /** One pool for everything (compute and IO). */ + public static EmbeddedExecutionContext of(ForkJoinPool pool) { + return new EmbeddedExecutionContext(pool, null); + } + + /** Separate compute and IO pools. */ + public static EmbeddedExecutionContext of(ForkJoinPool compute, ExecutorService io) { + return new EmbeddedExecutionContext(compute, io); + } + + /** The compute pool, as a {@link ParallelExecutor} (both simd and parallel roles). */ + public ParallelExecutor parallelExecutor() { + return computeParallel; + } + + /** The compute pool. */ + public ForkJoinPool computePool() { + return compute; + } + + /** The IO pool (equal to the compute pool unless a separate one was supplied). */ + public ExecutorService ioExecutor() { + return io; + } + + // ---- graph construction ---- + + /** + * A {@link GraphIndexBuilder} wired to the compute pool for both its executor roles. Equivalent + * to the pool-taking builder constructor with {@code compute} passed for both executors. + */ + public GraphIndexBuilder newBuilder(BuildScoreProvider scoreProvider, + int dimension, + int M, + int beamWidth, + float neighborOverflow, + float alpha, + boolean addHierarchy, + boolean refineFinalGraph) { + return new GraphIndexBuilder(scoreProvider, dimension, M, beamWidth, neighborOverflow, alpha, + addHierarchy, refineFinalGraph, computeParallel, computeParallel); + } + + // ---- compaction ---- + + /** + * An {@link OnDiskGraphIndexCompactor} wired to the compute pool. The caller sets a + * {@link io.github.jbellis.jvector.util.work.ProgressLimiter} on the returned compactor per + * operation if throttling/progress is wanted (see the class note). + */ + public OnDiskGraphIndexCompactor newCompactor(List sources, + List liveNodes, + List remappers, + VectorSimilarityFunction similarityFunction, + int taskWindowSize) { + return new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, similarityFunction, compute, taskWindowSize); + } + + // ---- parallel (IO-bound) graph writer ---- + + /** + * A parallel graph-writer builder wired to the IO pool. Chain feature/offset configuration on + * the returned builder and call {@code build()}. + */ + public OnDiskParallelGraphIndexWriter.Builder newParallelWriter(ImmutableGraphIndex graph, Path outputPath) throws IOException { + return new OnDiskParallelGraphIndexWriter.Builder(graph, outputPath).withExecutor(io); + } + + // ---- product quantization ---- + + /** Trains PQ on the compute pool (isotropic / unweighted). */ + public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter) { + return trainPQ(ravv, M, clusterCount, globallyCenter, UNWEIGHTED); + } + + /** Trains PQ on the compute pool with an explicit anisotropic threshold. */ + public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter, float anisotropicThreshold) { + return ProductQuantization.compute(ravv, M, clusterCount, globallyCenter, anisotropicThreshold, compute, compute); + } + + /** Refines an existing PQ codebook on the compute pool (one Lloyd's round, unweighted). */ + public ProductQuantization refinePQ(ProductQuantization base, RandomAccessVectorValues ravv) { + return base.refine(ravv, 1, UNWEIGHTED, compute, compute); + } + + /** + * Retrains PQ across compaction sources on the compute pool — the direct-call counterpart to + * the internal compaction retrain, and the entry point that closes the historical retrain leak. + */ + public ProductQuantization retrainPQ(PQRetrainer retrainer, VectorSimilarityFunction similarityFunction) { + return retrainer.retrain(similarityFunction, compute, compute); + } + + /** As {@link #retrainPQ(PQRetrainer, VectorSimilarityFunction)} with an explicit base PQ. */ + public ProductQuantization retrainPQ(PQRetrainer retrainer, VectorSimilarityFunction similarityFunction, ProductQuantization basePQ) { + return retrainer.retrain(similarityFunction, basePQ, compute, compute); + } + + /** Encodes all vectors with PQ on the compute pool. */ + public PQVectors encodePQ(ProductQuantization pq, RandomAccessVectorValues ravv) { + return pq.encodeAll(ravv, compute); + } + + // ---- non-uniform vector quantization ---- + + /** Encodes all vectors with NVQ on the compute pool. */ + public NVQVectors encodeNVQ(NVQuantization nvq, RandomAccessVectorValues ravv) { + return nvq.encodeAll(ravv, compute); + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java index b0deca442..c7cf2e7d5 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java @@ -18,10 +18,12 @@ import io.github.jbellis.jvector.quantization.CompressedVectors; import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import java.util.Collections; import java.util.List; import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; /** * Bundle of inputs that {@link QuantizationCompactionStrategy} implementations need to do their work @@ -39,8 +41,16 @@ public final class CompactionContext { public final int dimension; public final int maxOrdinal; public final ExecutorService executor; + /** + * A concrete {@link ForkJoinPool} for PQ operations (retrain/refine), which require one and + * cannot use the abstract {@link #executor}. Set to the compaction's own pool when it supplied + * a {@code ForkJoinPool}, so PQ retrain stays on that pool rather than leaking to an all-core + * default; otherwise the shared {@link PhysicalCoreExecutor#pool()}. + */ + public final ForkJoinPool computePool; public final int taskWindowSize; + /** Back-compat: {@code computePool} defaults to the shared physical-core pool. */ public CompactionContext( List sources, List sourceCompressed, @@ -50,6 +60,20 @@ public CompactionContext( int maxOrdinal, ExecutorService executor, int taskWindowSize) { + this(sources, sourceCompressed, liveNodes, remappers, dimension, maxOrdinal, executor, + PhysicalCoreExecutor.pool(), taskWindowSize); + } + + public CompactionContext( + List sources, + List sourceCompressed, + List liveNodes, + List remappers, + int dimension, + int maxOrdinal, + ExecutorService executor, + ForkJoinPool computePool, + int taskWindowSize) { this.sources = Collections.unmodifiableList(sources); this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed); this.liveNodes = Collections.unmodifiableList(liveNodes); @@ -57,6 +81,7 @@ public CompactionContext( this.dimension = dimension; this.maxOrdinal = maxOrdinal; this.executor = executor; + this.computePool = computePool; this.taskWindowSize = taskWindowSize; } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index 8a1636436..3eada3b49 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -691,8 +691,14 @@ private QuantizationCompactionStrategy detectSidecarStrategy() { /** Snapshot the compactor's state into a {@link CompactionContext} for strategies to consume. */ private CompactionContext buildContext() { + // PQ retrain needs a concrete ForkJoinPool; hand it the compaction's own pool when we have + // one (so retrain stays pool-bounded), else the shared physical-core pool. The abstract + // executor is still wrapped for the strategies' drain-safe invokeAll. + ForkJoinPool computePool = (executor instanceof ForkJoinPool) + ? (ForkJoinPool) executor + : PhysicalCoreExecutor.pool(); return new CompactionContext(sources, sourceCompressed, liveNodes, remappers, - dimension, maxOrdinal, asExecutorService(executor), taskWindowSize); + dimension, maxOrdinal, asExecutorService(executor), computePool, taskWindowSize); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index 616263951..83ed8b5d6 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -78,8 +78,18 @@ public PQRetrainer(List sources, List liveNodes, * performs no I/O. */ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction) { + return retrain(similarityFunction, PhysicalCoreExecutor.pool(), ForkJoinPool.commonPool()); + } + + /** + * As {@link #retrain(VectorSimilarityFunction)}, but runs the PQ refinement on the supplied + * pools instead of the default {@link PhysicalCoreExecutor#pool()} / {@code commonPool()} — so a + * compaction that supplies its own bounded pool keeps all retrain work on it rather than leaking + * to an all-core pool. + */ + public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { FusedPQ fpq = (FusedPQ) sources.get(0).getFeatures().get(FeatureId.FUSED_PQ); - return retrain(similarityFunction, fpq.getPQ()); + return retrain(similarityFunction, fpq.getPQ(), simdExecutor, parallelExecutor); } /** @@ -88,6 +98,15 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction) * non-fused source (e.g. a sidecar {@code CompressedVectors}) rather than the FUSED_PQ feature. */ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ) { + return retrain(similarityFunction, basePQ, PhysicalCoreExecutor.pool(), ForkJoinPool.commonPool()); + } + + /** + * As {@link #retrain(VectorSimilarityFunction, ProductQuantization)}, but runs the PQ refinement + * on the supplied pools. This is the overload that keeps a pool-bounded compaction from leaking + * PQ-retrain work to {@link PhysicalCoreExecutor#pool()} / {@code commonPool()}. + */ + public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { log.info("Training PQ using balanced sampling across sources"); List samples = sampleBalanced(ProductQuantization.MAX_PQ_TRAINING_SET_SIZE); @@ -115,8 +134,8 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization result = basePQ.refine(ravv, ProductQuantization.K_MEANS_ITERATIONS, -1.0f, // UNWEIGHTED / isotropic - PhysicalCoreExecutor.pool(), - ForkJoinPool.commonPool()); + simdExecutor, + parallelExecutor); log.info("PQ refinement complete in {}ms", (System.nanoTime() - t1) / 1_000_000L); return result; } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java index 3dc404481..6b00a1213 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java @@ -131,7 +131,7 @@ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ProductQuantization basePQ = this.pq; VectorCompressorRetrainer retrainer = vsf -> new PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension) - .retrain(vsf, basePQ); + .retrain(vsf, basePQ, ctx.computePool, ctx.computePool); return new FusedCompactionStrategy(ctx, this, retrainer); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java index 760eb38ee..659110063 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java @@ -424,7 +424,7 @@ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ProductQuantization basePQ = this.pq; io.github.jbellis.jvector.graph.disk.VectorCompressorRetrainer retrainer = vsf -> new io.github.jbellis.jvector.graph.disk.PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension) - .retrain(vsf, basePQ); + .retrain(vsf, basePQ, ctx.computePool, ctx.computePool); return new io.github.jbellis.jvector.graph.disk.SidecarCompactionStrategy(ctx, this, retrainer); } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java new file mode 100644 index 000000000..2ff9e7b7d --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java @@ -0,0 +1,261 @@ +/* + * 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 com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.disk.ReaderSupplier; +import io.github.jbellis.jvector.disk.ReaderSupplierFactory; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexCompactor; +import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndexWriter; +import io.github.jbellis.jvector.graph.disk.OrdinalMapper; +import io.github.jbellis.jvector.graph.disk.feature.Feature; +import io.github.jbellis.jvector.graph.disk.feature.FeatureId; +import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; +import io.github.jbellis.jvector.graph.disk.feature.InlineVectors; +import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider; +import io.github.jbellis.jvector.quantization.NVQuantization; +import io.github.jbellis.jvector.quantization.PQVectors; +import io.github.jbellis.jvector.quantization.ProductQuantization; +import io.github.jbellis.jvector.util.Bits; +import io.github.jbellis.jvector.util.FixedBitSet; +import io.github.jbellis.jvector.vector.VectorSimilarityFunction; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.EnumMap; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.ForkJoinWorkerThread; +import java.util.function.IntFunction; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/// Acceptance tests for {@link EmbeddedExecutionContext}: the "no parallel work escapes to +/// {@code commonPool} / {@code PhysicalCoreExecutor.pool()}" guarantee. The context is built with a +/// single-thread {@link ForkJoinPool} whose worker is distinctively named; a thread-recording +/// {@link RandomAccessVectorValues} observes which threads execute each operation's parallel bodies. +/// Every operation that reads vectors through the context (build, cleanup, PQ train/refine/encode, +/// NVQ encode) must run only on that named worker (or the calling thread) — never on a foreign +/// pool. A second test drives the graph-build and fused-compaction (PQ retrain) paths end to end +/// through the context and asserts the output is correct. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class EmbeddedExecutionContextTest extends RandomizedTest { + + private static final int DIM = 32; + private static final int SIZE = 1_000; + private static final VectorSimilarityFunction VSF = VectorSimilarityFunction.EUCLIDEAN; + private static final String WORKER_PREFIX = "noescape-worker-"; + + private static ForkJoinPool namedSingleThreadPool() { + ForkJoinPool.ForkJoinWorkerThreadFactory factory = p -> { + ForkJoinWorkerThread t = ForkJoinPool.defaultForkJoinWorkerThreadFactory.newThread(p); + t.setName(WORKER_PREFIX + t.getPoolIndex()); + return t; + }; + return new ForkJoinPool(1, factory, null, false); + } + + /// Every parallel body that reads vectors through the context must run on the supplied pool's + /// named worker (or the calling thread) — proving nothing escaped to the common pool or the + /// physical-core pool. + @Test + public void noParallelWorkEscapesTheSuppliedPool() throws Exception { + List> vectors = TestUtil.createRandomVectors(SIZE, DIM); + Set observed = ConcurrentHashMap.newKeySet(); + RecordingRAVV ravv = new RecordingRAVV(new ListRandomAccessVectorValues(vectors, DIM), observed); + // A separate, non-recording view for constructing inputs whose reads must NOT be attributed + // to the context (they run outside it). + ListRandomAccessVectorValues plain = new ListRandomAccessVectorValues(vectors, DIM); + + ForkJoinPool pool = namedSingleThreadPool(); + try { + var ctx = EmbeddedExecutionContext.of(pool); + + // build + cleanup (reads ravv via the build score provider and the build supplier) + try (var builder = ctx.newBuilder(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF), + DIM, 16, 100, 1.2f, 1.2f, true, true)) { + builder.build(ravv); + } + + // PQ train / refine / encode (all read ravv) + ProductQuantization pq = ctx.trainPQ(ravv, 8, 256, true); + ctx.refinePQ(pq, ravv); + ctx.encodePQ(pq, ravv); + + // NVQ encode (the quantization is trained off the plain view; only the encode reads ravv) + NVQuantization nvq = NVQuantization.compute(plain, 2); + ctx.encodeNVQ(nvq, ravv); + } finally { + pool.shutdown(); + } + + assertFalse("expected some parallel work to have run on the supplied pool", observed.isEmpty()); + boolean sawPoolWorker = false; + String mainName = Thread.currentThread().getName(); + for (String name : observed) { + boolean onPool = name.startsWith(WORKER_PREFIX); + sawPoolWorker |= onPool; + assertTrue("work escaped the supplied pool to thread: " + name + " (observed=" + observed + ")", + onPool || name.equals(mainName)); + } + assertTrue("no work actually ran on the supplied pool worker (observed=" + observed + ")", sawPoolWorker); + } + + /// End-to-end through the context: build two FusedPQ source graphs, then merge them with a + /// context-created compactor (which retrains PQ internally on the context's pool). Asserts the + /// merged graph is complete and searchable — exercising newBuilder, trainPQ, encodePQ, + /// newCompactor, and the (leak-closed) internal PQ retrain path. + @Test + public void graphBuildAndFusedCompactionThroughContext() throws Exception { + Path dir = Files.createTempDirectory("jvector_ctx_test"); + ForkJoinPool pool = namedSingleThreadPool(); + try { + var ctx = EmbeddedExecutionContext.of(pool); + int perSource = 300; + + List> v0 = TestUtil.createRandomVectors(perSource, DIM); + List> v1 = TestUtil.createRandomVectors(perSource, DIM); + Path s0 = buildFusedSource(ctx, dir.resolve("s0.graph"), v0); + Path s1 = buildFusedSource(ctx, dir.resolve("s1.graph"), v1); + + try (ReaderSupplier r0 = ReaderSupplierFactory.open(s0); + ReaderSupplier r1 = ReaderSupplierFactory.open(s1)) { + OnDiskGraphIndex g0 = OnDiskGraphIndex.load(r0); + OnDiskGraphIndex g1 = OnDiskGraphIndex.load(r1); + + FixedBitSet live0 = new FixedBitSet(perSource); + live0.set(0, perSource); + FixedBitSet live1 = new FixedBitSet(perSource); + live1.set(0, perSource); + Map map0 = new HashMap<>(); + Map map1 = new HashMap<>(); + for (int i = 0; i < perSource; i++) { + map0.put(i, i); + map1.put(i, perSource + i); + } + + OnDiskGraphIndexCompactor compactor = ctx.newCompactor( + List.of(g0, g1), List.of(live0, live1), + List.of(new OrdinalMapper.MapMapper(map0), new OrdinalMapper.MapMapper(map1)), + VSF, -1); + Path out = dir.resolve("merged.graph"); + compactor.compact(out); + + try (ReaderSupplier rOut = ReaderSupplierFactory.open(out)) { + OnDiskGraphIndex merged = OnDiskGraphIndex.load(rOut); + assertEquals("merged graph must contain all live nodes", 2 * perSource, merged.size(0)); + // sanity: the merged graph is searchable and returns the exact match for a source vector + List> all = new ArrayList<>(v0); + all.addAll(v1); + var mergedRavv = new ListRandomAccessVectorValues(all, DIM); + int hits = 0; + for (int q = 0; q < 50; q++) { + var res = GraphSearcher.search(all.get(q), 50, mergedRavv, VSF, merged, Bits.ALL); + for (var ns : res.getNodes()) { + if (ns.node == q) { hits++; break; } + } + } + assertTrue("merged graph should find self for most source vectors, got " + hits + "/50", hits >= 45); + } + } + } finally { + pool.shutdown(); + TestUtil.deleteQuietly(dir); + } + } + + /// Builds a single-layer graph with INLINE_VECTORS + FUSED_PQ, using the context for PQ, and + /// writes it. Mirrors the fused-source shape the compactor requires. + private Path buildFusedSource(EmbeddedExecutionContext ctx, Path out, List> vecs) throws IOException { + var ravv = new ListRandomAccessVectorValues(vecs, DIM); + ProductQuantization pq = ctx.trainPQ(ravv, 8, 256, true); + PQVectors pqv = ctx.encodePQ(pq, ravv); + var bsp = BuildScoreProvider.pqBuildScoreProvider(VSF, pqv); + try (var builder = ctx.newBuilder(bsp, DIM, 16, 100, 1.2f, 1.2f, true, true)) { + var graph = builder.getGraph(); + var writerBuilder = new OnDiskGraphIndexWriter.Builder(graph, out) + .withMapper(new OrdinalMapper.IdentityMapper(vecs.size() - 1)) + .with(new InlineVectors(DIM)) + .with(new FusedPQ(graph.maxDegree(), pq)); + var writer = writerBuilder.build(); + Map> suppliers = new EnumMap<>(FeatureId.class); + suppliers.put(FeatureId.INLINE_VECTORS, ord -> new InlineVectors.State(ravv.getVector(ord))); + for (int node = 0; node < ravv.size(); node++) { + var stateMap = new EnumMap(FeatureId.class); + stateMap.put(FeatureId.INLINE_VECTORS, new InlineVectors.State(ravv.getVector(node))); + writer.writeInline(node, stateMap); + builder.addGraphNode(node, ravv.getVector(node)); + } + builder.cleanup(); + suppliers.put(FeatureId.FUSED_PQ, ord -> new FusedPQ.State(graph.getView(), pqv, ord)); + writer.write(suppliers); + } + return out; + } + + /// A {@link RandomAccessVectorValues} that records the thread of every {@link #getVector} call, + /// funnelling all thread-local copies through the same shared set, then delegates. + private static final class RecordingRAVV implements RandomAccessVectorValues { + private final RandomAccessVectorValues delegate; + private final Set observed; + + RecordingRAVV(RandomAccessVectorValues delegate, Set observed) { + this.delegate = delegate; + this.observed = observed; + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public int dimension() { + return delegate.dimension(); + } + + @Override + public VectorFloat getVector(int nodeId) { + observed.add(Thread.currentThread().getName()); + return delegate.getVector(nodeId); + } + + @Override + public boolean isValueShared() { + return delegate.isValueShared(); + } + + @Override + public RandomAccessVectorValues copy() { + return new RecordingRAVV(delegate.copy(), observed); + } + } +} From 5a877b369e1d22030ecb81ca72cc4bd4f1d4c479 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 8 Jul 2026 21:21:40 +0000 Subject: [PATCH 11/13] Quantization: accept ParallelExecutor so PQ/NVQ/BQ can run caller-runs Widens the compress/train/encode entry points from ForkJoinPool to ParallelExecutor (keeping ForkJoinPool overloads), so encode/train can run on the calling thread. Encoding is byte-identical across executors; training is quality-equivalent (k-means seeds from ThreadLocalRandom). --- .../quantization/BinaryQuantization.java | 33 +-- .../jvector/quantization/NVQuantization.java | 27 +-- .../jvector/quantization/PQVectors.java | 40 ++-- .../quantization/ProductQuantization.java | 120 ++++++++--- .../quantization/VectorCompressor.java | 32 ++- .../QuantizationCallerRunsTest.java | 202 ++++++++++++++++++ 6 files changed, 376 insertions(+), 78 deletions(-) create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java index f0d660301..24d955650 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.vector.VectorizationProvider; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -26,7 +27,6 @@ import java.io.IOException; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.IntStream; /** * Binary Quantization of float vectors: each float is compressed to a single bit, @@ -57,28 +57,35 @@ public static BinaryQuantization compute(RandomAccessVectorValues ravv, ForkJoin return new BinaryQuantization(ravv.dimension()); } + /** {@link ParallelExecutor} overload; the executor is unused (BQ needs no training). */ + @Deprecated + public static BinaryQuantization compute(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { + return new BinaryQuantization(ravv.dimension()); + } + @Override public CompressedVectors createCompressedVectors(Object[] compressedVectors) { return new ImmutableBQVectors(this, (long[][]) compressedVectors); } @Override - public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { var ravvCopy = ravv.threadLocalSupplier(); - var cv = simdExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return v == null - ? new long[compressedVectorSize() / Long.BYTES] - : encode(v); - }) - .toArray(long[][]::new)) - .join(); + // Distinct-index array fill; the executor's completion barrier publishes the writes. + long[][] cv = new long[ravv.size()][]; + simdExecutor.forEachInt(ravv.size(), i -> { + var localRavv = ravvCopy.get(); + VectorFloat v = localRavv.getVector(i); + cv[i] = v == null ? new long[compressedVectorSize() / Long.BYTES] : encode(v); + }); return new ImmutableBQVectors(this, cv); } + @Override + public CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Encodes the input vector * diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java index aef0325b9..9adf3ec14 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.annotations.VisibleForTesting; import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; @@ -33,8 +34,6 @@ import java.util.List; import java.util.Objects; import java.util.concurrent.ForkJoinPool; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static io.github.jbellis.jvector.vector.VectorUtil.sub; @@ -179,18 +178,20 @@ public CompressedVectors createCompressedVectors(Object[] compressedVectors) { * Encodes the given vectors in parallel using NVQ. */ @Override - public NVQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { + public NVQVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { var ravvCopy = ravv.threadLocalSupplier(); - return new NVQVectors(this, - parallelExecutor.submit(() -> IntStream.range(0, ravv.size()) - .parallel() - .mapToObj(i -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(i); - return encode(v); - }) - .toArray(QuantizedVector[]::new)) - .join()); + // Distinct-index array fill; the executor's completion barrier publishes the writes. + QuantizedVector[] encoded = new QuantizedVector[ravv.size()]; + parallelExecutor.forEachInt(ravv.size(), i -> { + var localRavv = ravvCopy.get(); + encoded[i] = encode(localRavv.getVector(i)); + }); + return new NVQVectors(this, encoded); + } + + @Override + public NVQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(parallelExecutor)); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java index 659110063..3ca8e60aa 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.CompactionContext; import io.github.jbellis.jvector.graph.disk.QuantizationCompactionStrategy; @@ -91,10 +92,15 @@ public static PQVectors load(RandomAccessReader in, long offset) throws IOExcept * @param simdExecutor the ForkJoinPool to use for SIMD operations * @return the PQVectors instance */ - public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { return encodeAndBuild(pq, vectorCount, IntUnaryOperator.identity(), ravv, simdExecutor); } + /** {@link ForkJoinPool} overload of {@link #encodeAndBuild(ProductQuantization, int, RandomAccessVectorValues, ParallelExecutor)}. */ + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAndBuild(pq, vectorCount, ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Build a PQVectors instance from the given RandomAccessVectorValues. The vectors are encoded in parallel * and split into chunks to avoid exceeding the maximum array size. @@ -107,6 +113,15 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect * @return the PQVectors instance */ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAndBuild(pq, vectorCount, ordinalsMapping, ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + + /** + * As {@link #encodeAndBuild(ProductQuantization, int, IntUnaryOperator, RandomAccessVectorValues, ForkJoinPool)}, + * but {@code simdExecutor} may be {@link ParallelExecutor#callerRuns()} to encode synchronously on + * the calling thread. Encoding is per-vector independent, so the output is identical across paths. + */ + public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vectorCount, IntUnaryOperator ordinalsMapping, RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { // Verify that the mapped ordinals are packed, ie. the total range of ordinals does not exceed // the total vector count as they are mapped from 0 --> vector count - 1. // This assumes no duplicates are included. @@ -134,19 +149,16 @@ public static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vect // The changes are concurrent, but because they are coordinated and do not overlap, we can use parallel streams // and then we are guaranteed safe publication because we join the thread after completion. var ravvCopy = ravv.threadLocalSupplier(); - simdExecutor.submit(() -> IntStream.range(0, vectorCount) - .parallel() - .forEach(ordinal -> { - // Retrieve the slice and mutate it. - var localRavv = ravvCopy.get(); - var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); - var vector = localRavv.getVector(ordinalsMapping.applyAsInt(ordinal)); - if (vector != null) - pq.encodeTo(vector, slice); - else - slice.zero(); - })) - .join(); + simdExecutor.forEachInt(vectorCount, ordinal -> { + // Retrieve the slice and mutate it. + var localRavv = ravvCopy.get(); + var slice = PQVectors.get(chunks, ordinal, layout.fullChunkVectors, pq.getSubspaceCount()); + var vector = localRavv.getVector(ordinalsMapping.applyAsInt(ordinal)); + if (vector != null) + pq.encodeTo(vector, slice); + else + slice.zero(); + }); return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java index 6d9d23879..02b2b52d1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.annotations.VisibleForTesting; import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.disk.RandomAccessReader; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.Accountable; @@ -36,12 +37,9 @@ import java.util.List; import java.util.Objects; import java.util.SplittableRandom; -import java.util.concurrent.Callable; import java.util.concurrent.ForkJoinPool; import java.util.concurrent.atomic.AtomicReference; import java.util.logging.Logger; -import java.util.stream.Collectors; -import java.util.stream.IntStream; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; import static io.github.jbellis.jvector.util.MathUtil.square; @@ -113,6 +111,30 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, float anisotropicThreshold, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + return compute(ravv, M, clusterCount, globallyCenter, anisotropicThreshold, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * As {@link #compute(RandomAccessVectorValues, int, int, boolean, float, ForkJoinPool, ForkJoinPool)}, + * but the executors may be {@link ParallelExecutor#callerRuns()} to train synchronously on the + * calling thread with no worker threads. + *

+ * Codebook training draws random k-means++ seeds from {@code ThreadLocalRandom}, so the trained + * codebooks are already non-deterministic run-to-run and are not byte-identical between the + * caller-runs and pool-backed paths; both produce validly-trained, recall-equivalent codebooks. + * + * @param simdExecutor executor for SIMD/codebook operations + * @param parallelExecutor executor for training-vector extraction/centering + */ + public static ProductQuantization compute(RandomAccessVectorValues ravv, + int M, + int clusterCount, + boolean globallyCenter, + float anisotropicThreshold, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { checkClusterCount(clusterCount); @@ -123,9 +145,12 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, VectorFloat globalCentroid; if (globallyCenter) { globalCentroid = KMeansPlusPlusClusterer.centroidOf(vectors); - // subtract the centroid from each vector + // subtract the centroid from each vector (distinct-index array fill; the executor's + // completion barrier safely publishes the writes back to this thread) List> finalVectors = vectors; - vectors = simdExecutor.submit(() -> finalVectors.stream().parallel().map(v -> VectorUtil.sub(v, globalCentroid)).collect(Collectors.>toList())).join(); + VectorFloat[] centered = new VectorFloat[finalVectors.size()]; + simdExecutor.forEachInt(finalVectors.size(), i -> centered[i] = VectorUtil.sub(finalVectors.get(i), globalCentroid)); + vectors = Arrays.asList(centered); } else { globalCentroid = null; } @@ -135,11 +160,14 @@ public static ProductQuantization compute(RandomAccessVectorValues ravv, return new ProductQuantization(codebooks, clusterCount, subvectorSizesAndOffsets, globalCentroid, anisotropicThreshold); } - static List> extractTrainingVectors(RandomAccessVectorValues ravv, ForkJoinPool parallelExecutor) { - final IntStream ordinalStream; + static List> extractTrainingVectors(RandomAccessVectorValues ravv, ParallelExecutor parallelExecutor) { + final int[] ords; if (ravv.size() <= MAX_PQ_TRAINING_SET_SIZE) { - ordinalStream = IntStream.range(0, ravv.size()); + ords = new int[ravv.size()]; + for (int i = 0; i < ords.length; i++) { + ords[i] = i; + } } else { // Uses Floyd’s sampling algorithm to select MAX_PQ_TRAINING_SET_SIZE random ordinals from 0 to ravv.size() // while only iterating MAX_PQ_TRAINING_SET_SIZE times. @@ -154,25 +182,25 @@ static List> extractTrainingVectors(RandomAccessVectorValues ravv ordinals.add(t); } } - int[] ordinalArray = new int[ordinals.size()]; + ords = new int[ordinals.size()]; IntHashSet.IntIterator it = ordinals.iterator(); - for (int i = 0; i < ordinals.size(); i++) { + for (int i = 0; i < ords.length; i++) { assert it.hasNext(); - ordinalArray[i] = it.next(); + ords[i] = it.next(); } assert !it.hasNext(); - ordinalStream = IntStream.of(ordinalArray); } + // Distinct-index array fill in the target order; the executor's completion barrier safely + // publishes the writes (identical result under forkJoin and callerRuns). var ravvCopy = ravv.threadLocalSupplier(); - return parallelExecutor.submit(() -> ordinalStream.parallel() - .mapToObj(targetOrd -> { - var localRavv = ravvCopy.get(); - VectorFloat v = localRavv.getVector(targetOrd); - return localRavv.isValueShared() ? v.copy() : v; - }) - .collect(Collectors.toList())) - .join(); + VectorFloat[] out = new VectorFloat[ords.length]; + parallelExecutor.forEachInt(ords.length, i -> { + var localRavv = ravvCopy.get(); + VectorFloat v = localRavv.getVector(ords[i]); + out[i] = localRavv.isValueShared() ? v.copy() : v; + }); + return Arrays.asList(out); } /** @@ -193,6 +221,22 @@ public ProductQuantization refine(RandomAccessVectorValues ravv, float anisotropicThreshold, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) + { + return refine(ravv, lloydsRounds, anisotropicThreshold, + ParallelExecutor.forkJoin(simdExecutor), ParallelExecutor.forkJoin(parallelExecutor)); + } + + /** + * As {@link #refine(RandomAccessVectorValues, int, float, ForkJoinPool, ForkJoinPool)}, but the + * executors may be {@link ParallelExecutor#callerRuns()} to refine synchronously on the calling + * thread. See {@link #compute(RandomAccessVectorValues, int, int, boolean, float, ParallelExecutor, ParallelExecutor)} + * for the determinism note (codebooks are recall-equivalent, not byte-identical, across paths). + */ + public ProductQuantization refine(RandomAccessVectorValues ravv, + int lloydsRounds, + float anisotropicThreshold, + ParallelExecutor simdExecutor, + ParallelExecutor parallelExecutor) { if (lloydsRounds < 0) { throw new IllegalArgumentException("lloydsRounds must be non-negative"); @@ -201,18 +245,21 @@ public ProductQuantization refine(RandomAccessVectorValues ravv, var subvectorSizesAndOffsets = getSubvectorSizesAndOffsets(ravv.dimension(), M); var vectorsMutable = extractTrainingVectors(ravv, parallelExecutor); if (globalCentroid != null) { - var vectors = vectorsMutable; - vectorsMutable = simdExecutor.submit(() -> vectors.stream().parallel().map(v -> VectorUtil.sub(v, globalCentroid)).collect(Collectors.>toList())).join(); + List> src = vectorsMutable; + VectorFloat[] centered = new VectorFloat[src.size()]; + simdExecutor.forEachInt(src.size(), i -> centered[i] = VectorUtil.sub(src.get(i), globalCentroid)); + vectorsMutable = Arrays.asList(centered); } var vectors = vectorsMutable; // "effectively final" to make the closure happy - Callable[]> callable = () -> IntStream.range(0, M).parallel().mapToObj(m -> { + // Per-subquantizer (independent) codebook fill; completion barrier publishes the writes. + VectorFloat[] refinedCodebooks = new VectorFloat[M]; + simdExecutor.forEachInt(M, m -> { VectorFloat[] subvectors = extractSubvectors(vectors, m, subvectorSizesAndOffsets); var clusterer = new KMeansPlusPlusClusterer(subvectors, codebooks[m], anisotropicThreshold); - return clusterer.cluster(anisotropicThreshold == UNWEIGHTED ? lloydsRounds : 0, - anisotropicThreshold == UNWEIGHTED ? 0 : lloydsRounds); - }).toArray(VectorFloat[]::new); - var refinedCodebooks = simdExecutor.submit(callable).join(); + refinedCodebooks[m] = clusterer.cluster(anisotropicThreshold == UNWEIGHTED ? lloydsRounds : 0, + anisotropicThreshold == UNWEIGHTED ? 0 : lloydsRounds); + }); return new ProductQuantization(refinedCodebooks, clusterCount, subvectorSizesAndOffsets, globalCentroid, anisotropicThreshold); } @@ -255,10 +302,15 @@ public ImmutablePQVectors createCompressedVectors(Object[] compressedVectors) { * as a zero vector. */ @Override - public PQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + public PQVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { return PQVectors.encodeAndBuild(this, ravv.size(), ravv, simdExecutor); } + @Override + public PQVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + /** * Encodes the input vector using the PQ codebooks, weighing parallel loss more than orthogonal loss, into * the given ByteSequence. @@ -481,14 +533,16 @@ public int getClusterCount() { return clusterCount; } - static VectorFloat[] createCodebooks(List> vectors, int[][] subvectorSizeAndOffset, int clusters, float anisotropicThreshold, ForkJoinPool simdExecutor) { + static VectorFloat[] createCodebooks(List> vectors, int[][] subvectorSizeAndOffset, int clusters, float anisotropicThreshold, ParallelExecutor simdExecutor) { int M = subvectorSizeAndOffset.length; - Callable[]> callable = () -> IntStream.range(0, M).parallel().mapToObj(m -> { + // Per-subquantizer (independent) codebook fill; completion barrier publishes the writes. + VectorFloat[] codebooks = new VectorFloat[M]; + simdExecutor.forEachInt(M, m -> { VectorFloat[] subvectors = extractSubvectors(vectors, m, subvectorSizeAndOffset); var clusterer = new KMeansPlusPlusClusterer(subvectors, clusters, anisotropicThreshold); - return clusterer.cluster(K_MEANS_ITERATIONS, anisotropicThreshold == UNWEIGHTED ? 0 : K_MEANS_ITERATIONS); - }).toArray(VectorFloat[]::new); - return simdExecutor.submit(callable).join(); + codebooks[m] = clusterer.cluster(K_MEANS_ITERATIONS, anisotropicThreshold == UNWEIGHTED ? 0 : K_MEANS_ITERATIONS); + }); + return codebooks; } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java index c5708716c..f253acb16 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/VectorCompressor.java @@ -18,6 +18,7 @@ import io.github.jbellis.jvector.disk.IndexWriter; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.RandomAccessVectorValues; import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; @@ -44,11 +45,24 @@ default CompressedVectors encodeAll(RandomAccessVectorValues ravv) { /** * Encode all vectors in the RandomAccessVectorValues. If the RandomAccessVectorValues * has a missing vector for a given ordinal, the value will be encoded as a zero vector. + *

+ * Encoding is per-vector independent, so with {@link ParallelExecutor#callerRuns()} it runs + * synchronously on the calling thread and produces output identical to the pool-backed path; + * only wall-clock and thread usage differ. * @param ravv RandomAccessVectorValues to encode - * @param simdExecutor ForkJoinPool to use for SIMD operations + * @param simdExecutor executor hosting the parallel encode; {@link ParallelExecutor#callerRuns()} + * runs it on the calling thread * @return CompressedVectors containing the encoded vectors */ - CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor); + CompressedVectors encodeAll(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor); + + /** + * As {@link #encodeAll(RandomAccessVectorValues, ParallelExecutor)}, hosting the parallel encode + * on {@code simdExecutor}. Retained for callers holding a {@link ForkJoinPool}. + */ + default CompressedVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { + return encodeAll(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } T encode(VectorFloat v); @@ -108,8 +122,16 @@ default double[] reconstructionErrors(RandomAccessVectorValues ravv) { * @return the reconstruction error for each vector */ default double[] reconstructionErrors(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) { - return simdExecutor.submit(() -> - IntStream.range(0, ravv.size()).mapToDouble(i -> reconstructionError(ravv.getVector(i))).toArray() - ).join(); + return reconstructionErrors(ravv, ParallelExecutor.forkJoin(simdExecutor)); + } + + /** + * As {@link #reconstructionErrors(RandomAccessVectorValues, ForkJoinPool)}, hosting the parallel + * computation on the given {@link ParallelExecutor}. + */ + default double[] reconstructionErrors(RandomAccessVectorValues ravv, ParallelExecutor simdExecutor) { + double[] out = new double[ravv.size()]; + simdExecutor.forEachInt(ravv.size(), i -> out[i] = reconstructionError(ravv.getVector(i))); + return out; } } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java new file mode 100644 index 000000000..f109e6edd --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/QuantizationCallerRunsTest.java @@ -0,0 +1,202 @@ +/* + * 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.quantization; + +import com.carrotsearch.randomizedtesting.RandomizedTest; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; +import io.github.jbellis.jvector.TestUtil; +import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.ParallelExecutor; +import io.github.jbellis.jvector.graph.RandomAccessVectorValues; +import io.github.jbellis.jvector.vector.types.VectorFloat; +import org.junit.Test; + +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ForkJoinPool; + +import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/// Tests that the quantization path accepts {@link ParallelExecutor} and works caller-runs. +/// +/// The contract is deliberately split (see {@link ProductQuantization#compute} javadoc): +/// - **Encoding** is per-vector independent and deterministic given a fixed codebook, so +/// {@code encodeAll} produces byte-identical output under {@code callerRuns()} and a +/// {@code ForkJoinPool}. +/// - **Training** ({@code compute}/{@code refine}) draws k-means++ seeds from +/// {@code ThreadLocalRandom} and is therefore already non-deterministic run-to-run; it is +/// not byte-identical across executors, only recall/quality-equivalent. +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) +public class QuantizationCallerRunsTest extends RandomizedTest { + + private static final int DIM = 64; + private static final int SIZE = 2_000; + private static final int M = 8; + private static final int CLUSTERS = 256; + + private List> vectors; + private ListRandomAccessVectorValues ravv; + + private void makeData() { + vectors = TestUtil.createRandomVectors(SIZE, DIM); + ravv = new ListRandomAccessVectorValues(vectors, DIM); + } + + /// Encoding a fixed PQ codebook must be byte-identical whether run on a pool or caller-runs. + @Test + public void pqEncodeIsByteIdenticalAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization pq = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + + PQVectors onPool = pq.encodeAll(ravv, ParallelExecutor.forkJoin(pool)); + PQVectors callerRuns = pq.encodeAll(ravv, ParallelExecutor.callerRuns()); + + assertEquals("PQ encoding must be byte-identical across executors", onPool, callerRuns); + } finally { + pool.shutdown(); + } + } + + /// Binary quantization encoding must also be byte-identical across executors. + @Test + public void bqEncodeIsByteIdenticalAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + BinaryQuantization bq = new BinaryQuantization(DIM); + CompressedVectors onPool = bq.encodeAll(ravv, ParallelExecutor.forkJoin(pool)); + CompressedVectors callerRuns = bq.encodeAll(ravv, ParallelExecutor.callerRuns()); + assertEquals("BQ encoding must be byte-identical across executors", onPool, callerRuns); + } finally { + pool.shutdown(); + } + } + + /// Training on a pool vs caller-runs is NOT byte-identical (ThreadLocalRandom seeds), but must + /// produce codebooks of equivalent quality — asserted via average reconstruction error. + @Test + public void pqTrainingIsQualityEquivalentAcrossExecutors() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization onPool = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + ProductQuantization callerRuns = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + + double msePool = meanReconstructionError(onPool, ParallelExecutor.forkJoin(pool)); + double mseCaller = meanReconstructionError(callerRuns, ParallelExecutor.callerRuns()); + + assertTrue("pool-trained codebook should reconstruct well: " + msePool, msePool < 0.05); + assertTrue("caller-runs-trained codebook should reconstruct well: " + mseCaller, mseCaller < 0.05); + assertTrue("training quality must be equivalent across executors (" + msePool + " vs " + mseCaller + ")", + Math.abs(msePool - mseCaller) < 0.01); + } finally { + pool.shutdown(); + } + } + + /// refine() must also run caller-runs and produce a quality-equivalent codebook. + @Test + public void pqRefineRunsCallerRuns() { + makeData(); + ForkJoinPool pool = new ForkJoinPool(4); + try { + ProductQuantization base = ProductQuantization.compute(ravv, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.forkJoin(pool), ParallelExecutor.forkJoin(pool)); + ProductQuantization refined = base.refine(ravv, 1, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + assertTrue("refined codebook should reconstruct well", + meanReconstructionError(refined, ParallelExecutor.callerRuns()) < 0.05); + } finally { + pool.shutdown(); + } + } + + /// With callerRuns(), train + encode must execute only on the calling thread — no worker threads + /// and the common pool untouched. + @Test + public void callerRunsStaysOnCallingThread() { + makeData(); + Set observed = ConcurrentHashMap.newKeySet(); + RecordingRAVV recording = new RecordingRAVV(ravv, observed); + String mainName = Thread.currentThread().getName(); + + ProductQuantization pq = ProductQuantization.compute(recording, M, CLUSTERS, true, UNWEIGHTED, + ParallelExecutor.callerRuns(), ParallelExecutor.callerRuns()); + pq.encodeAll(recording, ParallelExecutor.callerRuns()); + + assertFalse("some vector reads must have occurred", observed.isEmpty()); + for (String name : observed) { + assertTrue("caller-runs must read only on the calling thread, but saw: " + name + " (all=" + observed + ")", + name.equals(mainName)); + } + } + + private double meanReconstructionError(ProductQuantization pq, ParallelExecutor ex) { + double[] errs = pq.reconstructionErrors(ravv, ex); + double s = 0; + for (double e : errs) { + s += e; + } + return s / errs.length; + } + + /// Records the thread of every getVector call, funnelling thread-local copies through one set. + private static final class RecordingRAVV implements RandomAccessVectorValues { + private final RandomAccessVectorValues delegate; + private final Set observed; + + RecordingRAVV(RandomAccessVectorValues delegate, Set observed) { + this.delegate = delegate; + this.observed = observed; + } + + @Override + public int size() { + return delegate.size(); + } + + @Override + public int dimension() { + return delegate.dimension(); + } + + @Override + public VectorFloat getVector(int nodeId) { + observed.add(Thread.currentThread().getName()); + return delegate.getVector(nodeId); + } + + @Override + public boolean isValueShared() { + return delegate.isValueShared(); + } + + @Override + public RandomAccessVectorValues copy() { + return new RecordingRAVV(delegate.copy(), observed); + } + } +} From c5c1b0ce7074fc72041845b91dfbad4765da23d6 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 8 Jul 2026 21:30:15 +0000 Subject: [PATCH 12/13] EmbeddedExecutionContext: carry a ParallelExecutor and support callerRuns() Now that quantization accepts ParallelExecutor, the facade carries one (plus a merge Executor and IO ExecutorService) instead of a ForkJoinPool, so a memtable flush can run build + PQ/NVQ entirely on its own thread via callerRuns(). PQRetrainer and CompactionContext.computeExecutor move to ParallelExecutor too, which also keeps retrain off the all-core pool when the merge executor isn't a ForkJoinPool. --- .../graph/EmbeddedExecutionContext.java | 139 ++++++++++++------ .../jvector/graph/disk/CompactionContext.java | 20 +-- .../graph/disk/OnDiskGraphIndexCompactor.java | 15 +- .../jvector/graph/disk/PQRetrainer.java | 21 +-- .../jvector/graph/disk/feature/FusedPQ.java | 2 +- .../jvector/quantization/PQVectors.java | 2 +- .../graph/EmbeddedExecutionContextTest.java | 29 ++++ 7 files changed, 153 insertions(+), 75 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java index 63c314769..b4dab30ac 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContext.java @@ -32,10 +32,14 @@ import java.io.IOException; import java.nio.file.Path; +import java.util.Collections; import java.util.List; import java.util.Objects; +import java.util.concurrent.AbstractExecutorService; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.TimeUnit; import static io.github.jbellis.jvector.quantization.KMeansPlusPlusClusterer.UNWEIGHTED; @@ -43,24 +47,21 @@ * A single carrier for the execution resources an embedder supplies to jvector, so the pool is * passed once and the "no parallel work escapes to {@link io.github.jbellis.jvector.util.PhysicalCoreExecutor#pool()} * or {@link ForkJoinPool#commonPool()}" guarantee lives in one place. Construct the context - * with a bounded pool, then obtain builders/compactors/writers and run quantization through it; - * every operation routes only through the carried executors and never through a default - * overload. + * with a bounded pool (or {@link #callerRuns()}), then obtain builders/compactors/writers and run + * quantization through it; every operation routes only through the carried executors and never + * through a default overload. * - *

Two executors. Graph and PQ/NVQ work is CPU-bound and runs on {@code compute}; the - * optional parallel graph write is IO-bound and runs on {@code io}, which defaults to - * {@code compute}. Passing one pool via {@link #of(ForkJoinPool)} is the simple, correct default; - * an embedder that enables parallel writes can supply a separate IO pool without touching any - * other call site. + *

Two modes. {@link #of(ForkJoinPool)} bounds all work to a supplied pool (the compaction + * use case). {@link #callerRuns()} runs everything synchronously on the calling thread — no worker + * threads, no pool — which lets a memtable flush run graph build and PQ/NVQ train/encode entirely on + * its own flush-writer thread. * - *

Why a {@code ForkJoinPool}, not a {@link ParallelExecutor}. Product/NV quantization is - * expressed with the parallel-stream-on-a-pool idiom, which only confines work when the pool is a - * {@code ForkJoinPool}. The context therefore carries a concrete compute {@code ForkJoinPool} so it - * can serve PQ (needs a {@code ForkJoinPool}), the graph builder (via - * {@link ParallelExecutor#forkJoin}), and the compactor (via {@link java.util.concurrent.Executor}) - * uniformly. A caller-runs, single-threaded mode is not offered here — PQ requires a real - * pool; for a build-only caller-runs graph, use {@link ParallelExecutor#callerRuns()} on - * {@link GraphIndexBuilder} directly. + *

Three executor roles. {@code compute} (a {@link ParallelExecutor}) drives graph build / + * cleanup and all PQ/NVQ quantization; {@code merge} (an {@link Executor}) drives the compaction + * merge; {@code io} (an {@link ExecutorService}) drives the optional IO-bound parallel graph writer. + * The factories derive all three consistently from one input, so "pass one thing" stays the simple, + * correct default; the general constructor accepts them separately for embedders that size an IO pool + * distinctly. * *

Lifecycle. The context neither owns nor shuts down the embedder's pool(s); the embedder * supplies and disposes them. @@ -72,42 +73,53 @@ */ @Experimental public final class EmbeddedExecutionContext { - private final ForkJoinPool compute; + private final ParallelExecutor compute; + private final Executor merge; private final ExecutorService io; - private final ParallelExecutor computeParallel; /** - * @param compute the pool for all CPU-bound build / cleanup / merge / PQ / NVQ work - * @param io the pool for the IO-bound parallel graph writer; {@code null} defaults to - * {@code compute} + * @param compute drives graph build/cleanup and all PQ/NVQ work + * @param merge drives the compaction merge (a {@code ForkJoinPool} or {@code Runnable::run} for + * caller-runs) + * @param io drives the IO-bound parallel graph writer */ - public EmbeddedExecutionContext(ForkJoinPool compute, ExecutorService io) { + public EmbeddedExecutionContext(ParallelExecutor compute, Executor merge, ExecutorService io) { this.compute = Objects.requireNonNull(compute, "compute"); - this.io = (io != null) ? io : compute; - this.computeParallel = ParallelExecutor.forkJoin(compute); + this.merge = Objects.requireNonNull(merge, "merge"); + this.io = Objects.requireNonNull(io, "io"); } - /** One pool for everything (compute and IO). */ + /** One pool for everything (compute, merge, and IO). */ public static EmbeddedExecutionContext of(ForkJoinPool pool) { - return new EmbeddedExecutionContext(pool, null); + return new EmbeddedExecutionContext(ParallelExecutor.forkJoin(pool), pool, pool); } - /** Separate compute and IO pools. */ - public static EmbeddedExecutionContext of(ForkJoinPool compute, ExecutorService io) { - return new EmbeddedExecutionContext(compute, io); + /** A compute pool plus a distinct IO pool; the compute pool also drives the merge. */ + public static EmbeddedExecutionContext of(ForkJoinPool computePool, ExecutorService io) { + return new EmbeddedExecutionContext(ParallelExecutor.forkJoin(computePool), computePool, io); } - /** The compute pool, as a {@link ParallelExecutor} (both simd and parallel roles). */ - public ParallelExecutor parallelExecutor() { - return computeParallel; + /** + * Runs every operation synchronously on the calling thread — no worker threads, no pool, common + * pool untouched. Suitable for a memtable flush that wants graph build + PQ/NVQ encode on its own + * thread. Encoding is byte-identical to the pool-backed path; PQ training is quality-equivalent + * (its k-means seeds are drawn from {@code ThreadLocalRandom} either way). + */ + public static EmbeddedExecutionContext callerRuns() { + return new EmbeddedExecutionContext(ParallelExecutor.callerRuns(), Runnable::run, directExecutorService()); } - /** The compute pool. */ - public ForkJoinPool computePool() { + /** The compute executor (graph build + PQ/NVQ). */ + public ParallelExecutor parallelExecutor() { return compute; } - /** The IO pool (equal to the compute pool unless a separate one was supplied). */ + /** The merge executor (compaction). */ + public Executor mergeExecutor() { + return merge; + } + + /** The IO executor (parallel graph writer). */ public ExecutorService ioExecutor() { return io; } @@ -115,8 +127,7 @@ public ExecutorService ioExecutor() { // ---- graph construction ---- /** - * A {@link GraphIndexBuilder} wired to the compute pool for both its executor roles. Equivalent - * to the pool-taking builder constructor with {@code compute} passed for both executors. + * A {@link GraphIndexBuilder} wired to the compute executor for both its executor roles. */ public GraphIndexBuilder newBuilder(BuildScoreProvider scoreProvider, int dimension, @@ -127,13 +138,13 @@ public GraphIndexBuilder newBuilder(BuildScoreProvider scoreProvider, boolean addHierarchy, boolean refineFinalGraph) { return new GraphIndexBuilder(scoreProvider, dimension, M, beamWidth, neighborOverflow, alpha, - addHierarchy, refineFinalGraph, computeParallel, computeParallel); + addHierarchy, refineFinalGraph, compute, compute); } // ---- compaction ---- /** - * An {@link OnDiskGraphIndexCompactor} wired to the compute pool. The caller sets a + * An {@link OnDiskGraphIndexCompactor} wired to the merge executor. The caller sets a * {@link io.github.jbellis.jvector.util.work.ProgressLimiter} on the returned compactor per * operation if throttling/progress is wanted (see the class note). */ @@ -142,13 +153,13 @@ public OnDiskGraphIndexCompactor newCompactor(List sources, List remappers, VectorSimilarityFunction similarityFunction, int taskWindowSize) { - return new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, similarityFunction, compute, taskWindowSize); + return new OnDiskGraphIndexCompactor(sources, liveNodes, remappers, similarityFunction, merge, taskWindowSize); } // ---- parallel (IO-bound) graph writer ---- /** - * A parallel graph-writer builder wired to the IO pool. Chain feature/offset configuration on + * A parallel graph-writer builder wired to the IO executor. Chain feature/offset configuration on * the returned builder and call {@code build()}. */ public OnDiskParallelGraphIndexWriter.Builder newParallelWriter(ImmutableGraphIndex graph, Path outputPath) throws IOException { @@ -157,23 +168,23 @@ public OnDiskParallelGraphIndexWriter.Builder newParallelWriter(ImmutableGraphIn // ---- product quantization ---- - /** Trains PQ on the compute pool (isotropic / unweighted). */ + /** Trains PQ on the compute executor (isotropic / unweighted). */ public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter) { return trainPQ(ravv, M, clusterCount, globallyCenter, UNWEIGHTED); } - /** Trains PQ on the compute pool with an explicit anisotropic threshold. */ + /** Trains PQ on the compute executor with an explicit anisotropic threshold. */ public ProductQuantization trainPQ(RandomAccessVectorValues ravv, int M, int clusterCount, boolean globallyCenter, float anisotropicThreshold) { return ProductQuantization.compute(ravv, M, clusterCount, globallyCenter, anisotropicThreshold, compute, compute); } - /** Refines an existing PQ codebook on the compute pool (one Lloyd's round, unweighted). */ + /** Refines an existing PQ codebook on the compute executor (one Lloyd's round, unweighted). */ public ProductQuantization refinePQ(ProductQuantization base, RandomAccessVectorValues ravv) { return base.refine(ravv, 1, UNWEIGHTED, compute, compute); } /** - * Retrains PQ across compaction sources on the compute pool — the direct-call counterpart to + * Retrains PQ across compaction sources on the compute executor — the direct-call counterpart to * the internal compaction retrain, and the entry point that closes the historical retrain leak. */ public ProductQuantization retrainPQ(PQRetrainer retrainer, VectorSimilarityFunction similarityFunction) { @@ -185,15 +196,49 @@ public ProductQuantization retrainPQ(PQRetrainer retrainer, VectorSimilarityFunc return retrainer.retrain(similarityFunction, basePQ, compute, compute); } - /** Encodes all vectors with PQ on the compute pool. */ + /** Encodes all vectors with PQ on the compute executor. */ public PQVectors encodePQ(ProductQuantization pq, RandomAccessVectorValues ravv) { return pq.encodeAll(ravv, compute); } // ---- non-uniform vector quantization ---- - /** Encodes all vectors with NVQ on the compute pool. */ + /** Encodes all vectors with NVQ on the compute executor. */ public NVQVectors encodeNVQ(NVQuantization nvq, RandomAccessVectorValues ravv) { return nvq.encodeAll(ravv, compute); } + + /** An {@link ExecutorService} that runs every submitted task synchronously on the calling thread. */ + private static ExecutorService directExecutorService() { + return new AbstractExecutorService() { + @Override + public void execute(Runnable command) { + command.run(); + } + + @Override + public void shutdown() { + } + + @Override + public List shutdownNow() { + return Collections.emptyList(); + } + + @Override + public boolean isShutdown() { + return false; + } + + @Override + public boolean isTerminated() { + return false; + } + + @Override + public boolean awaitTermination(long timeout, TimeUnit unit) { + return true; + } + }; + } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java index c7cf2e7d5..f5dafc70b 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionContext.java @@ -16,6 +16,7 @@ package io.github.jbellis.jvector.graph.disk; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.quantization.CompressedVectors; import io.github.jbellis.jvector.util.FixedBitSet; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; @@ -23,7 +24,6 @@ import java.util.Collections; import java.util.List; import java.util.concurrent.ExecutorService; -import java.util.concurrent.ForkJoinPool; /** * Bundle of inputs that {@link QuantizationCompactionStrategy} implementations need to do their work @@ -42,15 +42,15 @@ public final class CompactionContext { public final int maxOrdinal; public final ExecutorService executor; /** - * A concrete {@link ForkJoinPool} for PQ operations (retrain/refine), which require one and - * cannot use the abstract {@link #executor}. Set to the compaction's own pool when it supplied - * a {@code ForkJoinPool}, so PQ retrain stays on that pool rather than leaking to an all-core - * default; otherwise the shared {@link PhysicalCoreExecutor#pool()}. + * The {@link ParallelExecutor} for PQ operations (retrain/refine). Set to the compaction's own + * pool ({@code forkJoin}) when it supplied a {@code ForkJoinPool}, so PQ retrain stays there + * rather than leaking to an all-core default; {@code callerRuns()} when the compaction runs on a + * non-{@code ForkJoinPool} executor (so retrain stays on the caller's thread instead of leaking). */ - public final ForkJoinPool computePool; + public final ParallelExecutor computeExecutor; public final int taskWindowSize; - /** Back-compat: {@code computePool} defaults to the shared physical-core pool. */ + /** Back-compat: {@code computeExecutor} defaults to the shared physical-core pool. */ public CompactionContext( List sources, List sourceCompressed, @@ -61,7 +61,7 @@ public CompactionContext( ExecutorService executor, int taskWindowSize) { this(sources, sourceCompressed, liveNodes, remappers, dimension, maxOrdinal, executor, - PhysicalCoreExecutor.pool(), taskWindowSize); + ParallelExecutor.forkJoin(PhysicalCoreExecutor.pool()), taskWindowSize); } public CompactionContext( @@ -72,7 +72,7 @@ public CompactionContext( int dimension, int maxOrdinal, ExecutorService executor, - ForkJoinPool computePool, + ParallelExecutor computeExecutor, int taskWindowSize) { this.sources = Collections.unmodifiableList(sources); this.sourceCompressed = sourceCompressed == null ? null : Collections.unmodifiableList(sourceCompressed); @@ -81,7 +81,7 @@ public CompactionContext( this.dimension = dimension; this.maxOrdinal = maxOrdinal; this.executor = executor; - this.computePool = computePool; + this.computeExecutor = computeExecutor; this.taskWindowSize = taskWindowSize; } } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index 3eada3b49..bf0ad62a9 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -691,14 +691,15 @@ private QuantizationCompactionStrategy detectSidecarStrategy() { /** Snapshot the compactor's state into a {@link CompactionContext} for strategies to consume. */ private CompactionContext buildContext() { - // PQ retrain needs a concrete ForkJoinPool; hand it the compaction's own pool when we have - // one (so retrain stays pool-bounded), else the shared physical-core pool. The abstract - // executor is still wrapped for the strategies' drain-safe invokeAll. - ForkJoinPool computePool = (executor instanceof ForkJoinPool) - ? (ForkJoinPool) executor - : PhysicalCoreExecutor.pool(); + // PQ retrain runs on the compaction's own executor: forkJoin when it's a ForkJoinPool, else + // caller-runs — so a non-pool / caller-runs merge keeps retrain on the caller's thread rather + // than leaking to an all-core pool. The abstract executor is still wrapped for the strategies' + // drain-safe invokeAll. + ParallelExecutor computeExecutor = (executor instanceof ForkJoinPool) + ? ParallelExecutor.forkJoin((ForkJoinPool) executor) + : ParallelExecutor.callerRuns(); return new CompactionContext(sources, sourceCompressed, liveNodes, remappers, - dimension, maxOrdinal, asExecutorService(executor), computePool, taskWindowSize); + dimension, maxOrdinal, asExecutorService(executor), computeExecutor, taskWindowSize); } /** diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java index 83ed8b5d6..0a5b62846 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PQRetrainer.java @@ -17,6 +17,7 @@ package io.github.jbellis.jvector.graph.disk; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; +import io.github.jbellis.jvector.graph.ParallelExecutor; import io.github.jbellis.jvector.graph.disk.feature.FeatureId; import io.github.jbellis.jvector.graph.disk.feature.FusedPQ; import io.github.jbellis.jvector.quantization.ProductQuantization; @@ -78,16 +79,17 @@ public PQRetrainer(List sources, List liveNodes, * performs no I/O. */ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction) { - return retrain(similarityFunction, PhysicalCoreExecutor.pool(), ForkJoinPool.commonPool()); + return retrain(similarityFunction, + ParallelExecutor.forkJoin(PhysicalCoreExecutor.pool()), ParallelExecutor.forkJoin(ForkJoinPool.commonPool())); } /** * As {@link #retrain(VectorSimilarityFunction)}, but runs the PQ refinement on the supplied - * pools instead of the default {@link PhysicalCoreExecutor#pool()} / {@code commonPool()} — so a - * compaction that supplies its own bounded pool keeps all retrain work on it rather than leaking - * to an all-core pool. + * executors instead of the default {@link PhysicalCoreExecutor#pool()} / {@code commonPool()} — so + * a compaction that supplies its own bounded pool (or {@link ParallelExecutor#callerRuns()}) keeps + * all retrain work there rather than leaking to an all-core pool. */ - public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { + public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ParallelExecutor simdExecutor, ParallelExecutor parallelExecutor) { FusedPQ fpq = (FusedPQ) sources.get(0).getFeatures().get(FeatureId.FUSED_PQ); return retrain(similarityFunction, fpq.getPQ(), simdExecutor, parallelExecutor); } @@ -98,15 +100,16 @@ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, * non-fused source (e.g. a sidecar {@code CompressedVectors}) rather than the FUSED_PQ feature. */ public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ) { - return retrain(similarityFunction, basePQ, PhysicalCoreExecutor.pool(), ForkJoinPool.commonPool()); + return retrain(similarityFunction, basePQ, + ParallelExecutor.forkJoin(PhysicalCoreExecutor.pool()), ParallelExecutor.forkJoin(ForkJoinPool.commonPool())); } /** * As {@link #retrain(VectorSimilarityFunction, ProductQuantization)}, but runs the PQ refinement - * on the supplied pools. This is the overload that keeps a pool-bounded compaction from leaking - * PQ-retrain work to {@link PhysicalCoreExecutor#pool()} / {@code commonPool()}. + * on the supplied executors. This is the overload that keeps a pool-bounded (or caller-runs) + * compaction from leaking PQ-retrain work to {@link PhysicalCoreExecutor#pool()} / {@code commonPool()}. */ - public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ, ForkJoinPool simdExecutor, ForkJoinPool parallelExecutor) { + public ProductQuantization retrain(VectorSimilarityFunction similarityFunction, ProductQuantization basePQ, ParallelExecutor simdExecutor, ParallelExecutor parallelExecutor) { log.info("Training PQ using balanced sampling across sources"); List samples = sampleBalanced(ProductQuantization.MAX_PQ_TRAINING_SET_SIZE); diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java index 6b00a1213..ccfee7e62 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/feature/FusedPQ.java @@ -131,7 +131,7 @@ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ProductQuantization basePQ = this.pq; VectorCompressorRetrainer retrainer = vsf -> new PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension) - .retrain(vsf, basePQ, ctx.computePool, ctx.computePool); + .retrain(vsf, basePQ, ctx.computeExecutor, ctx.computeExecutor); return new FusedCompactionStrategy(ctx, this, retrainer); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java index 3ca8e60aa..3f072fd5a 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java @@ -436,7 +436,7 @@ public QuantizationCompactionStrategy createCompactionStrategy(CompactionContext ProductQuantization basePQ = this.pq; io.github.jbellis.jvector.graph.disk.VectorCompressorRetrainer retrainer = vsf -> new io.github.jbellis.jvector.graph.disk.PQRetrainer(ctx.sources, ctx.liveNodes, ctx.dimension) - .retrain(vsf, basePQ, ctx.computePool, ctx.computePool); + .retrain(vsf, basePQ, ctx.computeExecutor, ctx.computeExecutor); return new io.github.jbellis.jvector.graph.disk.SidecarCompactionStrategy(ctx, this, retrainer); } diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java index 2ff9e7b7d..becb0a32f 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/EmbeddedExecutionContextTest.java @@ -129,6 +129,35 @@ public void noParallelWorkEscapesTheSuppliedPool() throws Exception { assertTrue("no work actually ran on the supplied pool worker (observed=" + observed + ")", sawPoolWorker); } + /// A {@link EmbeddedExecutionContext#callerRuns()} context runs graph build and all PQ/NVQ work + /// synchronously on the calling thread — the memtable-flush mode. A recording view must observe + /// only the calling thread, never a worker or the common pool. + @Test + public void callerRunsExecutesEntirelyOnCallingThread() throws Exception { + List> vectors = TestUtil.createRandomVectors(SIZE, DIM); + Set observed = ConcurrentHashMap.newKeySet(); + RecordingRAVV ravv = new RecordingRAVV(new ListRandomAccessVectorValues(vectors, DIM), observed); + ListRandomAccessVectorValues plain = new ListRandomAccessVectorValues(vectors, DIM); + String mainName = Thread.currentThread().getName(); + + var ctx = EmbeddedExecutionContext.callerRuns(); + try (var builder = ctx.newBuilder(BuildScoreProvider.randomAccessScoreProvider(ravv, VSF), + DIM, 16, 100, 1.2f, 1.2f, true, true)) { + builder.build(ravv); + } + ProductQuantization pq = ctx.trainPQ(ravv, 8, 256, true); + ctx.refinePQ(pq, ravv); + ctx.encodePQ(pq, ravv); + NVQuantization nvq = NVQuantization.compute(plain, 2); + ctx.encodeNVQ(nvq, ravv); + + assertFalse("expected some vector reads", observed.isEmpty()); + for (String name : observed) { + assertTrue("caller-runs context must execute only on the calling thread, but saw: " + name + + " (observed=" + observed + ")", name.equals(mainName)); + } + } + /// End-to-end through the context: build two FusedPQ source graphs, then merge them with a /// context-created compactor (which retrains PQ internally on the context's pool). Asserts the /// merged graph is complete and searchable — exercising newBuilder, trainPQ, encodePQ, From ae5d1684db366b8b959063f9a26b6ee46ed9c14b Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Fri, 24 Jul 2026 19:48:35 +0000 Subject: [PATCH 13/13] Compaction: chunk the fused-PQ pre-encode code cache Replace the single 2 GiB-capped MappedByteBuffer cache with a chunked PqCodeCache of uniform power-of-two chunks (shift/mask ordinal addressing), so it is no longer silently skipped above ~21M nodes and its degree-x encode fallback avoided. Adds a compose-in/leave-off PqCodeCacheConfig and an internal dynamic bypass on the cache itself (PqCodeCache.NONE null-object, paralleling QuantizationCompactionStrategy.NONE). --- .../jvector/graph/disk/CompactWriter.java | 37 ++- .../graph/disk/FusedCompactionStrategy.java | 90 +++---- .../graph/disk/OnDiskGraphIndexCompactor.java | 44 +++- .../jvector/graph/disk/PqCodeCache.java | 247 ++++++++++++++++++ .../jvector/graph/disk/PqCodeCacheConfig.java | 75 ++++++ .../disk/QuantizationCompactionStrategy.java | 17 +- .../disk/TestOnDiskGraphIndexCompactor.java | 116 ++++++++ .../jvector/graph/disk/TestPqCodeCache.java | 176 +++++++++++++ 8 files changed, 706 insertions(+), 96 deletions(-) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCache.java create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCacheConfig.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPqCodeCache.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java index 182648735..06bdf789e 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactWriter.java @@ -19,7 +19,6 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.ByteOrder; -import java.nio.MappedByteBuffer; import java.nio.file.Path; import java.util.ArrayList; import java.util.LinkedHashMap; @@ -71,12 +70,12 @@ final class CompactWriter implements AutoCloseable { // Mirrors what AbstractGraphIndexWriter.writeSparseLevels writes in the no-hierarchy branch. private ByteSequence entryNodePqCode; - // Optional pre-computed PQ codes by new ordinal. When set, writeInlineNodeRecord copies - // codes from this buffer instead of calling pq.encodeTo per neighbor. Each worker thread - // gets its own duplicated view via pqCacheViewPerThread so positions don't race. - private MappedByteBuffer pqCodeCache; + // Optional pre-computed PQ codes by new ordinal. When set and active, writeInlineNodeRecord + // copies codes from this cache instead of calling pq.encodeTo per neighbor. Each worker thread + // gets its own per-chunk duplicate views via pqCacheViewsPerThread so positions don't race. + private PqCodeCache pqCodeCache; private int pqCodeSize; - private ThreadLocal pqCacheViewPerThread; + private ThreadLocal pqCacheViewsPerThread; private ThreadLocal pqCodeBufPerThread; CompactWriter(Path outputPath, @@ -136,18 +135,17 @@ final class CompactWriter implements AutoCloseable { /** * Enables the pre-computed PQ code cache. Must be called before any writeInlineNodeRecord - * call. Once enabled, neighbor PQ codes are copied from {@code cache} instead of being - * re-encoded per write. + * call. Once enabled and while {@link PqCodeCache#isActive() active}, neighbor PQ codes are + * copied from {@code cache} (indexed by new ordinal) instead of being re-encoded per write. * - * @param cache a buffer holding pqCodeSize bytes per new ordinal (length must be at - * least {@code (maxOrdinal + 1) * pqCodeSize}) - * @param pqCodeSize bytes per code (== FusedFeature.codeSize() of the source's feature) + * @param cache the chunked pre-encode cache holding one code per new ordinal; carries its own + * {@link PqCodeCache#codeSize() code size} */ - public void enablePqCodeCache(MappedByteBuffer cache, int pqCodeSize) { + public void enablePqCodeCache(PqCodeCache cache) { this.pqCodeCache = cache; - this.pqCodeSize = pqCodeSize; - // Each worker thread gets its own ByteBuffer view so absolute-position seeks don't race. - this.pqCacheViewPerThread = ThreadLocal.withInitial(() -> cache.duplicate()); + this.pqCodeSize = cache.codeSize(); + // Each worker thread gets its own per-chunk views so relative-position reads don't race. + this.pqCacheViewsPerThread = ThreadLocal.withInitial(cache::newViews); this.pqCodeBufPerThread = ThreadLocal.withInitial(() -> new byte[pqCodeSize]); } @@ -270,15 +268,14 @@ public WriteResult writeInlineNodeRecord(int ordinal, VectorFloat vec, Select // we cannot use fusedPQfeature.writeInline if (fusedPQEnabled) { int k = 0; - if (pqCodeCache != null) { + // Decide cache-vs-encode once per record (a single volatile read of the dynamic bypass). + if (pqCodeCache != null && pqCodeCache.isActive()) { // Look up neighbors' codes from the pre-encoded mmap'd cache instead of re-encoding. - ByteBuffer cacheView = pqCacheViewPerThread.get(); + ByteBuffer[] cacheViews = pqCacheViewsPerThread.get(); byte[] codeBuf = pqCodeBufPerThread.get(); for (; k < selectedCache.size; k++) { int newOrd = selectedCache.nodes[k]; // already remapped before this call - int offset = newOrd * pqCodeSize; - cacheView.position(offset); - cacheView.get(codeBuf, 0, pqCodeSize); + pqCodeCache.copyCode(cacheViews, newOrd, codeBuf); bwriter.write(codeBuf, 0, pqCodeSize); } } else { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java index 597ab2211..6bc4a6e51 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedCompactionStrategy.java @@ -25,12 +25,8 @@ import io.github.jbellis.jvector.vector.types.VectorTypeSupport; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import sun.misc.Unsafe; import java.io.IOException; -import java.lang.reflect.Field; -import java.nio.ByteBuffer; -import java.nio.MappedByteBuffer; import java.nio.channels.FileChannel; import java.nio.file.Path; import java.nio.file.StandardOpenOption; @@ -56,18 +52,6 @@ public final class FusedCompactionStrategy extends QuantizationCompactionStrategy { private static final Logger log = LoggerFactory.getLogger(FusedCompactionStrategy.class); private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport(); - private static final Unsafe UNSAFE = getUnsafe(); - - private static Unsafe getUnsafe() { - try { - Field f = Unsafe.class.getDeclaredField("theUnsafe"); - f.setAccessible(true); - return (Unsafe) f.get(null); - } catch (Exception e) { - log.warn("FusedCompactionStrategy can't acquire needed Unsafe access; code-cache will not be explicitly unmapped"); - return null; - } - } // Non-final: nulled by releaseSources() after compactGraphImpl so the source graphs reachable // through ctx.sources can be GC'd before refinement. onAfterClose must not touch ctx. @@ -77,13 +61,20 @@ private static Unsafe getUnsafe() { private VectorCompressor> retrainedCompressor; - // Transient pre-encode cache: lives in a memory-mapped section appended past the projected - // end of the output graph file. Truncated away in onAfterClose. Off-heap; single-mapping - // limit (2 GB) caps this at ~10M nodes for typical codeSize. - private MappedByteBuffer codeCache; - private int cacheCodeSize; + // Transient pre-encode cache: lives in a memory-mapped section appended past the projected end + // of the output graph file, chunked into <= 1 GiB mappings so it has no single-mapping size cap. + // Truncated away in onAfterClose. May be PqCodeCache.NONE when the cache is configured off. + private PqCodeCache codeCache; private long cacheTruncateAt; + // Compose-in / leave-off + chunk-size config; default composes the cache in at 1 GiB chunks. + private PqCodeCacheConfig pqCodeCacheConfig = PqCodeCacheConfig.DEFAULT; + + /** For compaction use. Sets the {@link PqCodeCache} configuration for this run. */ + public void setPqCodeCacheConfig(PqCodeCacheConfig pqCodeCacheConfig) { + this.pqCodeCacheConfig = pqCodeCacheConfig; + } + public FusedCompactionStrategy(CompactionContext ctx, FusedFeature sourceFusedFeature, VectorCompressorRetrainer retrainer) { @@ -105,15 +96,10 @@ public VectorCompressor compressor() { } @Override - public MappedByteBuffer getCodeCache() { + public PqCodeCache getCodeCache() { return codeCache; } - @Override - public int getCacheCodeSize() { - return cacheCodeSize; - } - @Override public boolean writesCodesInline() { return true; @@ -146,8 +132,8 @@ public void onAfterHeader(CompactWriter writer) throws IOException { } try { precomputeCodes(writer); - if (codeCache != null) { - writer.enablePqCodeCache(codeCache, cacheCodeSize); + if (codeCache != null && codeCache.isActive()) { + writer.enablePqCodeCache(codeCache); } } catch (IOException e) { // The fallback exists for environmental failures (mapping limits, transient IO) — @@ -186,12 +172,8 @@ public void onAfterLevels(CompactWriter writer, int[] entryNodeSource, List 0) { - if (codeCache != null && UNSAFE != null) { - try { - UNSAFE.invokeCleaner(codeCache); - } catch (IllegalArgumentException ignored) { - // duplicated/indirect buffer; not cleanable - } + if (codeCache != null) { + codeCache.unmap(); } codeCache = null; try (FileChannel fc = FileChannel.open(graphPath, StandardOpenOption.WRITE)) { @@ -205,27 +187,36 @@ public void onAfterClose(Path graphPath) { } } - /** Pre-encode every live node's code into a memory-mapped section past the projected output end. */ + /// Pre-encode every live node's code into a memory-mapped, chunked cache appended past the + /// projected output end. Composed in or left off per [PqCodeCacheConfig]; when off, the cache is + /// [PqCodeCache#NONE] and consumers encode per neighbor write. The chunked mapping has no + /// single-mapping (2 GiB) size cap, so — unlike the previous single-buffer cache — it is never + /// silently skipped above ~21M nodes. private void precomputeCodes(CompactWriter writer) throws IOException { - cacheCodeSize = retrainedCompressor.compressedVectorSize(); - long tempSize = (long) (ctx.maxOrdinal + 1) * cacheCodeSize; - if (tempSize <= 0 || tempSize > Integer.MAX_VALUE) { - log.info("Code pre-encode skipped: required cache size {} bytes exceeds single-mapping limit", tempSize); + if (!pqCodeCacheConfig.enabled()) { + log.info("Code pre-encode disabled by configuration; neighbor codes will be encoded per write"); + codeCache = PqCodeCache.NONE; + return; + } + + final int cs = retrainedCompressor.compressedVectorSize(); + long numCodes = (long) (ctx.maxOrdinal + 1); + long tempSize = numCodes * cs; + if (tempSize <= 0) { + log.info("Code pre-encode skipped: non-positive cache size {} bytes", tempSize); + codeCache = PqCodeCache.NONE; return; } long tempOffset = writer.projectedOutputSize(); cacheTruncateAt = tempOffset; - long totalSize = tempOffset + tempSize; try (FileChannel fc = FileChannel.open(writer.getOutputPath(), StandardOpenOption.READ, StandardOpenOption.WRITE)) { - ByteBuffer pad = ByteBuffer.wrap(new byte[]{0}); - fc.write(pad, totalSize - 1); - codeCache = fc.map(FileChannel.MapMode.READ_WRITE, tempOffset, tempSize); + codeCache = PqCodeCache.map(fc, tempOffset, cs, numCodes, pqCodeCacheConfig.maxChunkBytes()); } - final int cs = cacheCodeSize; + final PqCodeCache cache = codeCache; final VectorCompressor> compressor = retrainedCompressor; List> tasks = new ArrayList<>(); int targetTasks = Math.max(ctx.taskWindowSize * 4, 16); @@ -252,10 +243,7 @@ private void precomputeCodes(CompactWriter writer) throws IOException { code.zero(); compressor.encodeTo(vec, code); int newOrd = ctx.remappers.get(sIdx).oldToNew(oldOrd); - int offset = newOrd * cs; - for (int i = 0; i < cs; i++) { - codeCache.put(offset + i, code.get(i)); - } + cache.putCode(newOrd, code); count++; } } @@ -268,8 +256,8 @@ private void precomputeCodes(CompactWriter writer) throws IOException { for (Future f : ctx.executor.invokeAll(tasks)) { total += f.get(); } - log.info("Code pre-encode: {} nodes encoded into {} MB in-output cache (offset {})", - total, tempSize / (1024 * 1024), tempOffset); + log.info("Code pre-encode: {} nodes encoded into {} MB in-output cache across {} chunk(s) (offset {})", + total, tempSize / (1024 * 1024), cache.chunkCount(), tempOffset); } catch (InterruptedException | ExecutionException e) { throw new IOException("Code pre-encode failed", e); } diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java index bf0ad62a9..961e45282 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexCompactor.java @@ -104,6 +104,7 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; private boolean refineAfterCompaction = false; + private PqCodeCacheConfig pqCodeCacheConfig = PqCodeCacheConfig.DEFAULT; // Embedder progress + work-admission control surface (see io.github.jbellis.jvector.util.work). // Default UNLIMITED = no observation and no throttling → byte-identical output and equivalent @@ -162,6 +163,18 @@ public void setRefineAfterCompaction(boolean refineAfterCompaction) { this.refineAfterCompaction = refineAfterCompaction; } + /** + * Configures the fused-PQ pre-encode {@link PqCodeCache} for this run: whether it is composed in + * or left off, and its chunk sizing. Defaults to {@link PqCodeCacheConfig#DEFAULT} (composed in, + * 1 GiB chunks). Passing {@code null} restores the default. Only affects fused-PQ compactions. + * Returns {@code this} for chaining. + */ + @Experimental + public OnDiskGraphIndexCompactor setPqCodeCacheConfig(PqCodeCacheConfig config) { + this.pqCodeCacheConfig = (config == null) ? PqCodeCacheConfig.DEFAULT : config; + return this; + } + /** * Primary constructor: merges multiple graph indexes using any {@link Executor} with an * explicit in-flight window. @@ -670,7 +683,11 @@ private void releaseSourcesBeforeRefine(QuantizationCompactionStrategy strategy) private QuantizationCompactionStrategy detectInlineStrategy() { for (var feature : sources.get(0).getFeatures().values()) { if (feature instanceof FusedFeature) { - return ((FusedFeature) feature).createCompactionStrategy(buildContext()); + QuantizationCompactionStrategy strategy = ((FusedFeature) feature).createCompactionStrategy(buildContext()); + if (strategy instanceof FusedCompactionStrategy) { + ((FusedCompactionStrategy) strategy).setPqCodeCacheConfig(pqCodeCacheConfig); + } + return strategy; } } return QuantizationCompactionStrategy.NONE; @@ -783,11 +800,11 @@ private void refineCompactedGraph(Path outputPath, long startOffset, Quantizatio baseDegree * SEARCH_TOP_K_MULTIPLIER); final int beamWidth = Math.max(baseDegree, searchTopK) * BEAM_WIDTH_MULTIPLIER; - // Code cache may or may not be present; capture once so refineOneNode can take the fast path. - // The cache is shared across threads; refineOneNode duplicates per call (cheap; no per-thread - // state to track and the duplicates are tiny GC-friendly ByteBuffer wrappers). - final java.nio.MappedByteBuffer codeCache = hasFusedPQ ? strategy.getCodeCache() : null; - final int cacheCodeSize = hasFusedPQ ? strategy.getCacheCodeSize() : 0; + // Code cache may or may not be present/active; capture once so refineOneNode can take the + // fast path. The cache is shared across threads; refineOneNode takes its own per-chunk views + // per call (cheap; the duplicates are tiny GC-friendly ByteBuffer wrappers). + final PqCodeCache codeCache = hasFusedPQ ? strategy.getCodeCache() : null; + final int cacheCodeSize = codeCache != null ? codeCache.codeSize() : 0; try (var supplier = new SimpleReader.Supplier(outputPath); FileChannel fc = FileChannel.open(outputPath, StandardOpenOption.WRITE, StandardOpenOption.READ)) { @@ -824,7 +841,7 @@ private void refineCompactedGraph(Path outputPath, long startOffset, Quantizatio final int codeSize = pqCodeSize; final VectorCompressor> cmp = compressor; final int bw = beamWidth; - final java.nio.MappedByteBuffer cache = codeCache; + final PqCodeCache cache = codeCache; final int cacheSz = cacheCodeSize; final OnDiskGraphIndex graphRef = mergedGraph; @@ -914,7 +931,7 @@ private void refineOneNode(int node, VectorCompressor> compressor, int beamWidth, OnDiskGraphIndex mergedGraph, - java.nio.MappedByteBuffer codeCache, + PqCodeCache codeCache, int cacheCodeSize) { OnDiskGraphIndex.View view = scratch.view; view.getVectorInto(node, scratch.queryVec, 0); @@ -1006,16 +1023,15 @@ private void refineOneNode(int node, if (hasFusedPQ) { // PQ codes block sits between the inline vector and the neighbor count. writeOffset = view.offsetFor(node, FeatureId.FUSED_PQ); - if (codeCache != null) { + if (codeCache != null && codeCache.isActive()) { // Memcpy from the pre-encoded cache (indexed by new ordinal). Avoids one FP - // vector read AND one PQ encode per selected neighbor. duplicate() gives this - // call its own position cursor without racing other workers. - ByteBuffer cacheView = codeCache.duplicate(); + // vector read AND one PQ encode per selected neighbor. newViews() gives this + // call its own per-chunk cursors without racing other workers. + ByteBuffer[] cacheViews = codeCache.newViews(); byte[] codeBuf = scratch.pqCodeBytes; for (int k = 0; k < selectedSize; k++) { int newOrd = scratch.selectedNodes[k]; - cacheView.position(newOrd * cacheCodeSize); - cacheView.get(codeBuf, 0, cacheCodeSize); + codeCache.copyCode(cacheViews, newOrd, codeBuf); rec.put(codeBuf, 0, cacheCodeSize); } } else { diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCache.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCache.java new file mode 100644 index 000000000..0cf9a36de --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCache.java @@ -0,0 +1,247 @@ +/* + * 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.disk; + +import io.github.jbellis.jvector.vector.types.ByteSequence; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import sun.misc.Unsafe; + +import java.io.IOException; +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.nio.MappedByteBuffer; +import java.nio.channels.FileChannel; + +/// Chunked, memory-mapped pre-encode cache of fused-PQ codes, indexed by new (post-remap) ordinal. +/// +/// Replaces the single-[MappedByteBuffer] cache, whose `Integer.MAX_VALUE` (2 GiB) mapping limit +/// silently disabled the cache — forcing roughly `degree ×` encode amplification — once +/// `(maxOrdinal + 1) * codeSize` crossed `2^31`. Codes are split across an array of mappings, each +/// holding a whole number of codes and at most `maxChunkBytes` (≤ 1 GiB), so total capacity is +/// bounded only by disk. A code never straddles a chunk boundary, so every read/write is a single +/// contiguous `codeSize`-byte access. +/// +/// **Uniform, power-of-two chunks.** Every chunk spans exactly `codesPerChunk` ordinals, and +/// `codesPerChunk` is a power of two ([#codesPerChunkFor]). The global-ordinal address therefore +/// reduces to register shift/mask rather than integer division: +/// ``` +/// chunk = ord >>> codesPerChunkShift // ord / codesPerChunk +/// within = (ord & codesPerChunkMask) * codeSize // (ord % codesPerChunk) * codeSize +/// ``` +/// Because `within < codesPerChunk * codeSize <= maxChunkBytes <= 1 GiB`, it always fits an `int` — +/// the old `int offset = ord * codeSize` overflow is structurally impossible. +/// +/// Threading mirrors the previous cache: writes use absolute [MappedByteBuffer#put(int, byte)] (the +/// pre-encode tasks cover disjoint ordinals, so concurrent absolute puts never race), and readers +/// each take their own [#newViews()] duplicates so relative-position reads don't share a cursor. +/// jvector-base targets `--release 11`, so this uses NIO mappings rather than a +/// `java.lang.foreign` `MemorySegment`. +public final class PqCodeCache { + private static final Logger log = LoggerFactory.getLogger(PqCodeCache.class); + private static final Unsafe UNSAFE = getUnsafe(); + + /// Maximum bytes per chunk mapping: 1 GiB. Also the hard ceiling enforced by + /// [#codesPerChunkFor], which keeps `codesPerChunk * codeSize` (and thus every within-chunk + /// offset) below `2^31`, comfortably under the 2 GiB single-mapping limit. + public static final long DEFAULT_MAX_CHUNK_BYTES = 1L << 30; + + /// Null-object sentinel for "no active cache", paralleling [QuantizationCompactionStrategy#NONE]. + /// It holds no mappings, so [#isActive()] is permanently `false` and every consumer takes its + /// encode fallback. Returned instead of `null` when the cache is composed in but configured off + /// (see [PqCodeCacheConfig]), keeping the object graph uniform. Safe to share as a singleton: a + /// chunk-less cache is inactive regardless of the bypass flag, so it cannot be toggled into a + /// state that would read from absent mappings. + public static final PqCodeCache NONE = new PqCodeCache(new MappedByteBuffer[0], 0, 1, false); + + private final MappedByteBuffer[] chunks; + private final int codeSize; + private final int codesPerChunk; // power of two: uniform stride across chunks + private final int codesPerChunkShift; // log2(codesPerChunk): ord -> chunk index + private final int codesPerChunkMask; // codesPerChunk - 1: ord -> in-chunk code index + + /// Internal bypass. When {@code false}, consumers must encode rather than read this cache; a + /// mapped cache starts active and can be flipped off at runtime via [#setActive] / [#bypass]. + /// Volatile so a mid-run toggle is observed by reader threads. Bypass is always correctness-safe + /// because a memcpy from the cache and a fresh `encodeTo` produce identical bytes for the same + /// compressor, so consumers may switch paths per record. + private volatile boolean active; + + private PqCodeCache(MappedByteBuffer[] chunks, int codeSize, int codesPerChunk, boolean active) { + assert Integer.bitCount(codesPerChunk) == 1 : "codesPerChunk must be a power of two, got " + codesPerChunk; + this.chunks = chunks; + this.codeSize = codeSize; + this.codesPerChunk = codesPerChunk; + this.codesPerChunkShift = Integer.numberOfTrailingZeros(codesPerChunk); + this.codesPerChunkMask = codesPerChunk - 1; + this.active = active; + } + + /// Maps a code cache over `[baseOffset, baseOffset + numCodes * codeSize)` of `fc`, growing the + /// file if needed. The region is split into uniform chunks of `codesPerChunk` codes each (the + /// last chunk maps only the remaining codes), every chunk a separate [MappedByteBuffer]. + /// + /// @param fc the output channel, opened for read and write + /// @param baseOffset byte offset at which the cache region begins + /// @param codeSize bytes per code (`> 0`) + /// @param numCodes number of codes to hold, one per new ordinal (`> 0`) + /// @param maxChunkBytes soft cap on bytes per chunk; values above 1 GiB are clamped down + public static PqCodeCache map(FileChannel fc, long baseOffset, int codeSize, long numCodes, long maxChunkBytes) throws IOException { + if (codeSize <= 0) { + throw new IllegalArgumentException("codeSize must be > 0, got " + codeSize); + } + if (numCodes <= 0) { + throw new IllegalArgumentException("numCodes must be > 0, got " + numCodes); + } + int codesPerChunk = codesPerChunkFor(codeSize, maxChunkBytes); + long totalBytes = numCodes * codeSize; + + // Ensure the backing file spans the whole cache region before mapping any chunk. + long end = baseOffset + totalBytes; + if (fc.size() < end) { + ByteBuffer pad = ByteBuffer.wrap(new byte[]{0}); + fc.write(pad, end - 1); + } + + int numChunks = (int) ((numCodes + codesPerChunk - 1) / codesPerChunk); + MappedByteBuffer[] chunks = new MappedByteBuffer[numChunks]; + for (int c = 0; c < numChunks; c++) { + long firstCode = (long) c * codesPerChunk; + long chunkOffset = baseOffset + firstCode * codeSize; + long thisChunkCodes = Math.min(codesPerChunk, numCodes - firstCode); + long thisChunkBytes = thisChunkCodes * codeSize; + chunks[c] = fc.map(FileChannel.MapMode.READ_WRITE, chunkOffset, thisChunkBytes); + } + return new PqCodeCache(chunks, codeSize, codesPerChunk, true); + } + + /// Largest power-of-two number of codes that fits in `maxChunkBytes`, at least 1. Power-of-two + /// sizing makes the per-chunk stride uniform so ordinal addressing is shift/mask (see class + /// doc). The budget is capped at [#DEFAULT_MAX_CHUNK_BYTES] so that `codesPerChunk * codeSize` — + /// the largest within-chunk byte offset — stays below `2^31` and cannot overflow `int`. + public static int codesPerChunkFor(int codeSize, long maxChunkBytes) { + long budget = Math.min(maxChunkBytes, DEFAULT_MAX_CHUNK_BYTES); + long perChunk = budget / codeSize; + if (perChunk < 1) { + return 1; // 2^0: a single code larger than the budget still gets its own chunk + } + // Largest power of two <= perChunk. perChunk <= 1 GiB / codeSize <= 2^30, so this fits int. + return Integer.highestOneBit((int) perChunk); + } + + /// Bytes per code held in this cache. + public int codeSize() { + return codeSize; + } + + /// Number of codes stored per chunk (a power of two). The last chunk may hold fewer. + public int codesPerChunk() { + return codesPerChunk; + } + + /// `log2(codesPerChunk)`: the shift that maps a new ordinal to its chunk index. + public int codesPerChunkShift() { + return codesPerChunkShift; + } + + /// `codesPerChunk - 1`: the mask that maps a new ordinal to its in-chunk code index. + public int codesPerChunkMask() { + return codesPerChunkMask; + } + + /// Number of chunk mappings backing this cache. + public int chunkCount() { + return chunks.length; + } + + /// Whether consumers should read codes from this cache ({@code true}) or take their encode + /// fallback ({@code false}). Read once per record on the write/refine hot paths. A chunk-less + /// cache (the [#NONE] sentinel) is never active regardless of the bypass flag, which keeps the + /// shared sentinel a safe, immutable-in-effect null-object. + public boolean isActive() { + return active && chunks.length > 0; + } + + /// Dynamically enables or disables this cache. Bypass ({@code false}) is correctness-safe: a + /// consumer that encodes instead of reading produces identical bytes for the same compressor. + public void setActive(boolean active) { + this.active = active; + } + + /// Convenience for {@code setActive(false)} — dynamically bypasses this cache. + public void bypass() { + this.active = false; + } + + /// Per-reader duplicate views (one per chunk) so the relative-position reads in [#copyCode] + /// don't share a cursor with other threads. The duplicates are cheap wrappers over the same + /// mapped memory; hand a distinct array to each reader thread. + public ByteBuffer[] newViews() { + ByteBuffer[] views = new ByteBuffer[chunks.length]; + for (int c = 0; c < chunks.length; c++) { + views[c] = chunks[c].duplicate(); + } + return views; + } + + /// Writes `code` (length `codeSize`) at new ordinal `ord` using absolute puts. Thread-safe for + /// concurrent callers as long as they target disjoint ordinals (as the pre-encode tasks do). + public void putCode(int ord, ByteSequence code) { + MappedByteBuffer chunk = chunks[ord >>> codesPerChunkShift]; + int within = (ord & codesPerChunkMask) * codeSize; + for (int i = 0; i < codeSize; i++) { + chunk.put(within + i, code.get(i)); + } + } + + /// Copies the `codeSize` bytes stored at new ordinal `ord` into `out`, using this reader's own + /// `views` (from [#newViews()]). `out` must be at least `codeSize` long. + public void copyCode(ByteBuffer[] views, int ord, byte[] out) { + ByteBuffer view = views[ord >>> codesPerChunkShift]; + view.position((ord & codesPerChunkMask) * codeSize); + view.get(out, 0, codeSize); + } + + /// Explicitly unmaps every chunk mapping. Best-effort (no-op if `Unsafe` is unavailable); call + /// once, before truncating the cache section off the output file. + public void unmap() { + if (UNSAFE == null) { + return; + } + for (MappedByteBuffer chunk : chunks) { + if (chunk == null) { + continue; + } + try { + UNSAFE.invokeCleaner(chunk); + } catch (IllegalArgumentException ignored) { + // duplicated/indirect buffer; not cleanable + } + } + } + + private static Unsafe getUnsafe() { + try { + Field f = Unsafe.class.getDeclaredField("theUnsafe"); + f.setAccessible(true); + return (Unsafe) f.get(null); + } catch (Exception e) { + log.warn("PqCodeCache can't acquire needed Unsafe access; mapped chunks will not be explicitly unmapped"); + return null; + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCacheConfig.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCacheConfig.java new file mode 100644 index 000000000..755969450 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/PqCodeCacheConfig.java @@ -0,0 +1,75 @@ +/* + * 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.disk; + +/// Immutable configuration for the fused-PQ pre-encode code cache ([PqCodeCache]), composed into a +/// compaction run via {@code OnDiskGraphIndexCompactor.setPqCodeCacheConfig(...)}. +/// +/// Two independent controls: +/// - **compose in or leave off** — [#enabled]. When off, no cache is built and every consumer +/// encodes per neighbor write (the pre-fix behavior); the strategy hands consumers +/// [PqCodeCache#NONE] so the object graph stays uniform. +/// - **chunk sizing** — [#maxChunkBytes], the soft cap on bytes per mmap chunk (clamped to +/// [PqCodeCache#DEFAULT_MAX_CHUNK_BYTES]). Smaller values force more chunks; the default keeps a +/// graph in as few ≤ 1 GiB mappings as possible. +/// +/// Distinct from the cache's *internal* [PqCodeCache#isActive() dynamic bypass]: this config is the +/// build-time, compose-in/leave-off decision, whereas the bypass can flip a built cache off at +/// runtime. +public final class PqCodeCacheConfig { + /// Cache composed in, chunks sized to the 1 GiB default. The compactor default. + public static final PqCodeCacheConfig DEFAULT = new PqCodeCacheConfig(true, PqCodeCache.DEFAULT_MAX_CHUNK_BYTES); + + /// Cache left off: consumers encode per neighbor write (no pre-encode cache built). + public static final PqCodeCacheConfig DISABLED = new PqCodeCacheConfig(false, PqCodeCache.DEFAULT_MAX_CHUNK_BYTES); + + private final boolean enabled; + private final long maxChunkBytes; + + public PqCodeCacheConfig(boolean enabled, long maxChunkBytes) { + if (maxChunkBytes <= 0) { + throw new IllegalArgumentException("maxChunkBytes must be > 0, got " + maxChunkBytes); + } + this.enabled = enabled; + this.maxChunkBytes = maxChunkBytes; + } + + /// Whether to build the pre-encode cache at all. + public boolean enabled() { + return enabled; + } + + /// Soft cap on bytes per mmap chunk (clamped to 1 GiB by [PqCodeCache#codesPerChunkFor]). + public long maxChunkBytes() { + return maxChunkBytes; + } + + /// Returns a copy with {@link #enabled()} set as given. + public PqCodeCacheConfig withEnabled(boolean enabled) { + return new PqCodeCacheConfig(enabled, maxChunkBytes); + } + + /// Returns a copy with {@link #maxChunkBytes()} set as given. + public PqCodeCacheConfig withMaxChunkBytes(long maxChunkBytes) { + return new PqCodeCacheConfig(enabled, maxChunkBytes); + } + + @Override + public String toString() { + return "PqCodeCacheConfig{enabled=" + enabled + ", maxChunkBytes=" + maxChunkBytes + '}'; + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java index 5d9c7311c..1c258d906 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/QuantizationCompactionStrategy.java @@ -22,7 +22,6 @@ import io.github.jbellis.jvector.vector.VectorSimilarityFunction; import java.io.IOException; -import java.nio.MappedByteBuffer; import java.nio.file.Path; import java.util.List; @@ -148,20 +147,16 @@ public FusedFeature outputFusedFeature(int maxDegree) { /** * For compaction use. Returns the precomputed code cache built by {@link #onAfterHeader}, - * indexed by new ordinal so refinement can memcpy neighbor codes instead of re-encoding them. - * Returns {@code null} when no cache is held (non-fused strategy, NONE, or graph too large for - * a single mapping). The returned buffer is shared; callers must {@code .duplicate()} per - * thread before using. + * indexed by new ordinal so both the base-layer writer and refinement can memcpy neighbor codes + * instead of re-encoding them. Returns {@code null} when no cache is held (non-fused strategy, + * NONE); a fused strategy that was configured off returns {@link PqCodeCache#NONE}. Callers + * must gate reads on {@link PqCodeCache#isActive()} and take their own {@link PqCodeCache#newViews()} + * per thread. The cache carries its own {@link PqCodeCache#codeSize() code size}. */ - public MappedByteBuffer getCodeCache() { + public PqCodeCache getCodeCache() { return null; } - /** For compaction use. Bytes per code in {@link #getCodeCache()}, or {@code 0} when no cache. */ - public int getCacheCodeSize() { - return 0; - } - /** * For compaction use. Drops the strategy's hold on the {@link CompactionContext} (and thus the * source graphs) once {@code compactGraphImpl} no longer needs it, so the source graphs become diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java index e9bc30d49..a698f60f3 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestOnDiskGraphIndexCompactor.java @@ -676,6 +676,122 @@ public void testCallerRunsProducesValidGraph() throws Exception { } } + /** Search the compacted graph with exact scoring and return recall@10 vs brute-force GT. */ + private double recallOnCompacted(OnDiskGraphIndex g) throws IOException { + List> queries = new ArrayList<>(); + for (int i = 0; i < numQueries; ++i) queries.add(allVecs.get(randomIntBetween(0, allVecs.size() - 1))); + List> gt = buildGT(queries, 10); + List results = new ArrayList<>(); + try (GraphSearcher searcher = new GraphSearcher(g)) { + for (VectorFloat q : queries) { + SearchScoreProvider ssp = DefaultSearchScoreProvider.exact(q, similarityFunction, allravv); + results.add(searcher.search(ssp, 10, Bits.ALL)); + } + } + return AccuracyMetrics.recallFromSearchResults(gt, results, 10, 10); + } + + @Test + public void testFusedCompactionAcrossMultipleCacheChunks() throws Exception { + // 3 sources x 256 vectors = 768 merged nodes; PQ has 8 subspaces => 8-byte codes. A 256-byte + // chunk budget => 32 codes/chunk => 24 chunks, so the pre-encode cache genuinely spans many + // chunks. Refinement is enabled so BOTH cache consumers (CompactWriter and refineOneNode) + // read neighbor codes across chunk boundaries in a real compaction. + assertEquals("test intends a multi-chunk setup", 32, PqCodeCache.codesPerChunkFor(8, 256)); + + List rss = new ArrayList<>(); + try { + var path = testDirectory.resolve("multichunk"); + var compactor = newAllLiveCompactor(rss); + compactor.setPqCodeCacheConfig(PqCodeCacheConfig.DEFAULT.withMaxChunkBytes(256)); + compactor.setRefineAfterCompaction(true); + compactor.compact(path); + + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + var g = OnDiskGraphIndex.load(out); + assertEquals(numSources * numVectorsPerGraph, g.size(0)); + double recall = recallOnCompacted(g); + assertTrue("multi-chunk compacted graph recall should be reasonable, got " + recall, recall >= 0.2); + } + } finally { + for (var rs : rss) rs.close(); + } + } + + @Test + public void testFusedCompactionWithCodeCacheDisabled() throws Exception { + // Leaving the cache off must still produce a valid graph — every neighbor code is encoded + // per write (the pre-fix fallback path), exercised here as a first-class config, not a cliff. + List rss = new ArrayList<>(); + try { + var path = testDirectory.resolve("cache_disabled"); + var compactor = newAllLiveCompactor(rss); + compactor.setPqCodeCacheConfig(PqCodeCacheConfig.DISABLED); + compactor.setRefineAfterCompaction(true); + compactor.compact(path); + + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + var g = OnDiskGraphIndex.load(out); + assertEquals(numSources * numVectorsPerGraph, g.size(0)); + double recall = recallOnCompacted(g); + assertTrue("cache-disabled compacted graph recall should be reasonable, got " + recall, recall >= 0.2); + } + } finally { + for (var rs : rss) rs.close(); + } + } + + /** + * Search the compacted graph with the FUSED approximate score function — the decoder reads the + * on-disk neighbor PQ codes during traversal — plus an exact reranker, returning recall@10 vs + * brute-force GT. Unlike {@link #recallOnCompacted}, this path actually decodes the codes the + * pre-encode cache wrote. + */ + private double fusedApproxRecall(OnDiskGraphIndex g, int rerankK) throws IOException { + List> queries = new ArrayList<>(); + for (int i = 0; i < numQueries; ++i) queries.add(allVecs.get(randomIntBetween(0, allVecs.size() - 1))); + List> gt = buildGT(queries, 10); + List results = new ArrayList<>(); + try (GraphSearcher searcher = new GraphSearcher(g)) { + for (VectorFloat q : queries) { + var scoringView = (ImmutableGraphIndex.ScoringView) searcher.getView(); + var asf = scoringView.approximateScoreFunctionFor(q, similarityFunction); + var rr = scoringView.rerankerFor(q, similarityFunction); + SearchScoreProvider ssp = new DefaultSearchScoreProvider(asf, rr); + results.add(searcher.search(ssp, 10, rerankK, 0f, 0f, Bits.ALL)); + } + } + return AccuracyMetrics.recallFromSearchResults(gt, results, 10, 10); + } + + @Test + public void testFusedApproximateSearchDecodesMultiChunkCodes() throws Exception { + // Compact across many cache chunks, then search with the FUSED approximate score function so + // the FusedPQDecoder reads the multi-chunk-written neighbor codes back off disk during + // traversal. A chunk-addressing bug in the code cache would corrupt those codes and collapse + // recall; the exact-scoring integration tests above never read the on-disk codes, this does. + assertEquals("test intends a multi-chunk setup", 32, PqCodeCache.codesPerChunkFor(8, 256)); + + List rss = new ArrayList<>(); + try { + var path = testDirectory.resolve("multichunk_fused_search"); + var compactor = newAllLiveCompactor(rss); + compactor.setPqCodeCacheConfig(PqCodeCacheConfig.DEFAULT.withMaxChunkBytes(256)); + compactor.setRefineAfterCompaction(true); + compactor.compact(path); + + try (ReaderSupplier out = ReaderSupplierFactory.open(path)) { + var g = OnDiskGraphIndex.load(out); + assertEquals(numSources * numVectorsPerGraph, g.size(0)); + double recall = fusedApproxRecall(g, 32); + assertTrue("fused-approximate recall over multi-chunk codes should be high, got " + recall, + recall >= 0.6); + } + } finally { + for (var rs : rss) rs.close(); + } + } + @Test public void testCallerRunsRunsOnCallingThreadOnly() throws Exception { List rss = new ArrayList<>(); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPqCodeCache.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPqCodeCache.java new file mode 100644 index 000000000..5a79733c8 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/disk/TestPqCodeCache.java @@ -0,0 +1,176 @@ +/* + * 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.disk; + +import io.github.jbellis.jvector.vector.VectorizationProvider; +import io.github.jbellis.jvector.vector.types.ByteSequence; +import io.github.jbellis.jvector.vector.types.VectorTypeSupport; +import org.junit.Test; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; + +import static java.nio.file.StandardOpenOption.CREATE; +import static java.nio.file.StandardOpenOption.READ; +import static java.nio.file.StandardOpenOption.TRUNCATE_EXISTING; +import static java.nio.file.StandardOpenOption.WRITE; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +/// Proves the chunked {@link PqCodeCache} primitive that replaces the single-mapping fused-PQ code +/// cache: multi-chunk round-trip, power-of-two chunk sizing, and int-overflow-safe addressing for +/// ordinals whose old {@code ord * codeSize} offset would have overflowed {@code int}. +public class TestPqCodeCache { + private static final VectorTypeSupport vts = VectorizationProvider.getInstance().getVectorTypeSupport(); + + private static Path targetFile(String name) throws IOException { + Path dir = Paths.get("target", "pqcodecache"); + Files.createDirectories(dir); + return dir.resolve(name); + } + + private static ByteSequence codeFor(int ord, int codeSize) { + ByteSequence code = vts.createByteSequence(codeSize); + for (int i = 0; i < codeSize; i++) { + // distinct, ordinal- and position-dependent so a mis-addressed read is detected + code.set(i, (byte) (ord * 31 + i * 7 + 1)); + } + return code; + } + + /// Write a distinct code per ordinal across many chunks, then read every one back through the + /// reader views and assert exact identity. A wrong chunk index or within-chunk offset — the + /// failure mode of the old single-mapping cache above ~21M nodes — would corrupt this. + @Test + public void roundTripAcrossManyChunks() throws IOException { + int codeSize = 7; + long maxChunkBytes = 20; // budget/codeSize = 2, largest pow2 <= 2 is 2 -> 2 codes/chunk + int numCodes = 11; // -> ceil(11/2) = 6 chunks, last chunk partial (1 code) + + Path file = targetFile("roundtrip.bin"); + try (FileChannel fc = FileChannel.open(file, CREATE, READ, WRITE, TRUNCATE_EXISTING)) { + PqCodeCache cache = PqCodeCache.map(fc, 0, codeSize, numCodes, maxChunkBytes); + assertEquals("codes per chunk (power of two)", 2, cache.codesPerChunk()); + assertEquals("chunk count", 6, cache.chunkCount()); + assertEquals("shift = log2(codesPerChunk)", 1, cache.codesPerChunkShift()); + assertEquals("mask = codesPerChunk - 1", 1, cache.codesPerChunkMask()); + assertTrue("a mapped cache is active", cache.isActive()); + + for (int ord = 0; ord < numCodes; ord++) { + cache.putCode(ord, codeFor(ord, codeSize)); + } + + ByteBuffer[] views = cache.newViews(); + byte[] out = new byte[codeSize]; + for (int ord = 0; ord < numCodes; ord++) { + cache.copyCode(views, ord, out); + ByteSequence expected = codeFor(ord, codeSize); + for (int i = 0; i < codeSize; i++) { + assertEquals("ord " + ord + " byte " + i, expected.get(i), out[i]); + } + } + cache.unmap(); + } + } + + /// Chunk sizing is the largest power of two of codes that fits the (1 GiB-capped) budget. + @Test + public void codesPerChunkForIsPowerOfTwoWithinBudget() { + assertEquals(8, PqCodeCache.codesPerChunkFor(100, 1000)); // 1000/100=10 -> 8 + assertEquals(1, PqCodeCache.codesPerChunkFor(100, 100)); // 1 + assertEquals(1, PqCodeCache.codesPerChunkFor(100, 50)); // code larger than budget -> 1 + assertEquals(8388608, PqCodeCache.codesPerChunkFor(100, 1L << 30)); // 2^30/100 -> 2^23 + + for (long budget : new long[]{1, 63, 256, 999, 1L << 20, 1L << 30, (1L << 30) + 1}) { + for (int codeSize : new int[]{1, 3, 8, 100, 999}) { + int perChunk = PqCodeCache.codesPerChunkFor(codeSize, budget); + assertTrue("power of two: " + perChunk, Integer.bitCount(perChunk) == 1); + assertTrue("at least one code", perChunk >= 1); + long effectiveBudget = Math.min(budget, PqCodeCache.DEFAULT_MAX_CHUNK_BYTES); + assertTrue("fits budget unless a single code exceeds it", + (long) perChunk * codeSize <= effectiveBudget || perChunk == 1); + // within-chunk offsets stay < 2^31 (the whole point of the chunked addressing) + assertTrue("chunk span fits int", (long) perChunk * codeSize < (1L << 31)); + } + } + } + + /// The addressing (shift/mask) is int-overflow-safe for ordinals whose old + /// {@code int offset = ord * codeSize} would have wrapped past {@code Integer.MAX_VALUE}. This is + /// the exact cliff that silently disabled the single-mapping cache; here it is structurally gone. + @Test + public void addressingIsIntOverflowSafeForLargeOrdinals() throws IOException { + int codeSize = 100; + // Map a tiny cache purely to obtain the production shift/mask for this codeSize + 1 GiB budget. + Path file = targetFile("addressing.bin"); + try (FileChannel fc = FileChannel.open(file, CREATE, READ, WRITE, TRUNCATE_EXISTING)) { + PqCodeCache cache = PqCodeCache.map(fc, 0, codeSize, 4, PqCodeCache.DEFAULT_MAX_CHUNK_BYTES); + int shift = cache.codesPerChunkShift(); + int mask = cache.codesPerChunkMask(); + int codesPerChunk = cache.codesPerChunk(); + + int ord = 30_000_000; // ord * codeSize = 3.0e9, which overflows a signed int + assertTrue("precondition: old int offset would overflow", + (long) ord * codeSize > Integer.MAX_VALUE); + + int chunkIndex = ord >>> shift; + int within = (ord & mask) * codeSize; + assertEquals("chunk index matches ord / codesPerChunk", ord / codesPerChunk, chunkIndex); + assertTrue("within-chunk offset is non-negative (no overflow)", within >= 0); + assertTrue("within-chunk offset stays inside the chunk", + within < (long) codesPerChunk * codeSize); + cache.unmap(); + } + } + + /// Internal dynamic bypass: a mapped cache is active and can be flipped off/on at runtime; the + /// NONE sentinel is a composed-in-but-off cache holding no mappings. + @Test + public void dynamicBypassAndNoneSentinel() throws IOException { + assertFalse("NONE reports inactive", PqCodeCache.NONE.isActive()); + assertEquals("NONE holds no chunks", 0, PqCodeCache.NONE.chunkCount()); + + Path file = targetFile("bypass.bin"); + try (FileChannel fc = FileChannel.open(file, CREATE, READ, WRITE, TRUNCATE_EXISTING)) { + PqCodeCache cache = PqCodeCache.map(fc, 0, 4, 8, PqCodeCache.DEFAULT_MAX_CHUNK_BYTES); + assertTrue(cache.isActive()); + cache.bypass(); + assertFalse("bypass() disables", cache.isActive()); + cache.setActive(true); + assertTrue("setActive(true) re-enables", cache.isActive()); + cache.unmap(); + } + } + + /// Config: compose-in/leave-off and chunk sizing are independent, immutable settings. + @Test + public void configComposition() { + assertTrue(PqCodeCacheConfig.DEFAULT.enabled()); + assertEquals(PqCodeCache.DEFAULT_MAX_CHUNK_BYTES, PqCodeCacheConfig.DEFAULT.maxChunkBytes()); + assertFalse(PqCodeCacheConfig.DISABLED.enabled()); + + assertFalse(PqCodeCacheConfig.DEFAULT.withEnabled(false).enabled()); + assertEquals(256, PqCodeCacheConfig.DEFAULT.withMaxChunkBytes(256).maxChunkBytes()); + // withX preserves the other field + assertTrue(PqCodeCacheConfig.DEFAULT.withMaxChunkBytes(256).enabled()); + } +}