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 1/5] 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 aeb61c02f43f93ae5d75f678d2c9e7a74ecb843e Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 1 Jul 2026 08:02:17 +0000 Subject: [PATCH 2/5] 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 | 196 ++++++++++++ .../util/work/TestProgressLimiter.java | 289 ++++++++++++++++++ 8 files changed, 948 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 ddf5d3cc5..488912177 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; @@ -85,6 +88,51 @@ public final class OnDiskGraphIndexCompactor implements Accountable { private final int taskWindowSize; private final VectorSimilarityFunction similarityFunction; + // 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; + } + /** * Constructs a new OnDiskGraphIndexCompactor for graphs without a non-fused compressed sidecar. * Equivalent to calling the 6-arg constructor with {@code sourceCompressed = null}. @@ -108,6 +156,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( @@ -294,11 +349,32 @@ 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); - refineCompactedGraph(outputPath, strategy); + refineCompactedGraph(outputPath, startOffset, 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,8 +405,8 @@ public void compact(Path graphPath, Path compressedPath) throws FileNotFoundExce QuantizationCompactionStrategy sidecarStrategy = detectSidecarStrategy(); try { sidecarStrategy.retrain(similarityFunction); - compactGraphImpl(graphPath, inlineStrategy); - refineCompactedGraph(graphPath, inlineStrategy); + compactGraphImpl(graphPath, 0L, inlineStrategy); + refineCompactedGraph(graphPath, 0L, inlineStrategy); sidecarStrategy.writeSidecar(compressedPath); } catch (IOException e) { throw new RuntimeException("Sidecar compaction failed", e); @@ -401,7 +477,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(); @@ -417,7 +493,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(); @@ -457,7 +533,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(); @@ -484,7 +560,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 @@ -520,33 +596,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 @@ -868,6 +962,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); @@ -908,7 +1007,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 ); } @@ -945,7 +1051,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 ); } } @@ -1291,7 +1405,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(); @@ -1307,11 +1424,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..e5a3540ab 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,194 @@ 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); + + 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 22d1de3784707e0921891f38782dacd3f43fe6cd Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Wed, 1 Jul 2026 17:41:57 +0000 Subject: [PATCH 3/5] 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 | 134 +++++++++---- .../jvector/example/CompactionBench.java | 2 +- .../jvector/disk/TestSeekableSink.java | 85 +++++++++ .../disk/TestOnDiskGraphIndexCompactor.java | 176 +++++++++++++++--- 11 files changed, 720 insertions(+), 67 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 488912177..8e2219ec4 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; @@ -134,35 +134,32 @@ public int getTaskWindowSize() { } /** - * 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( @@ -171,22 +168,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; @@ -210,6 +201,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 @@ -382,6 +416,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 @@ -464,7 +520,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 e5a3540ab..ce32b8d43 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(); } } @@ -732,6 +791,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. @@ -775,7 +909,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); @@ -842,7 +976,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); @@ -920,7 +1054,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); @@ -1026,7 +1160,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"); @@ -1109,7 +1243,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", @@ -1152,7 +1286,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", @@ -1188,7 +1322,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"); @@ -1248,7 +1382,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 a4f4c57fbb5fd783afb05dce115b014409b9d1ee Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 6 Jul 2026 22:55:30 +0000 Subject: [PATCH 4/5] 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 30431d6ca..e7fc850d6 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,23 @@ public interface ReaderSupplier extends AutoCloseable { */ RandomAccessReader get() throws IOException; + /** + * 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 f759bc702..e0673b10c 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 3fb69d967..a21e9e441 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. */ @@ -475,6 +491,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; @@ -493,6 +524,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 @@ -509,6 +541,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 @@ -569,8 +602,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 8e2219ec4..599d33fb7 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); @@ -224,24 +240,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; + } } /** @@ -380,6 +484,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 { @@ -392,8 +502,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 @@ -404,6 +527,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); @@ -423,6 +549,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} */ @@ -441,11 +570,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 { @@ -455,6 +589,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(); @@ -471,6 +611,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 @@ -684,9 +839,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); @@ -695,6 +851,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(); } @@ -1470,40 +1633,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 ad29eeaa7a70d44bc9b6f1895a3dc0d5dc55a29b Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Mon, 6 Jul 2026 22:55:30 +0000 Subject: [PATCH 5/5] add memory-safety repro and regression-guard tests --- .../example/repro/BoundsCheckBenchTest.java | 125 +++++ .../jvector/example/repro/ChildJvm.java | 157 ++++++ .../repro/CompactorStragglerReproTest.java | 497 ++++++++++++++++++ .../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, 2385 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..7b12947eb --- /dev/null +++ b/jvector-examples/src/test/java/io/github/jbellis/jvector/example/repro/CompactorStragglerReproTest.java @@ -0,0 +1,497 @@ +/* + * 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); + + 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); + } + } +}