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/ReaderSupplier.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java index 7ebb3f9b0..8827bcdb1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java @@ -41,6 +41,23 @@ public interface ReaderSupplier extends AutoCloseable { default void prefetch(long offset, long length) { } + /** + * Releases the supplier's underlying resource. Two implementation families exist, with very + * different safety under concurrency: + * + * 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/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/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/ChunkingParallelExecutor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ChunkingParallelExecutor.java new file mode 100644 index 000000000..3ce220950 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ChunkingParallelExecutor.java @@ -0,0 +1,248 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.PrimitiveIterator; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.RejectedExecutionException; +import java.util.function.Consumer; +import java.util.function.IntConsumer; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +/** + * {@link ParallelExecutor} over a caller-supplied {@link ExecutorService}: the calling thread + * traverses the source and submits fixed-size chunks to the executor, keeping a bounded number in + * flight. Backs {@link ParallelExecutor#over(ExecutorService, int)}; see the {@link ParallelExecutor} + * class javadoc for how this implementation's caveats compare with the other factories. + */ +final class ChunkingParallelExecutor implements ParallelExecutor { + /** Elements per submitted chunk on the stream paths ({@code forEachInt} splits its range evenly instead). */ + private static final int BATCH_SIZE = 32; + /** Chunks per worker on the {@code forEachInt} path: more than one smooths skewed per-element cost. */ + private static final int CHUNKS_PER_WORKER = 4; + + // Nested use from a body degrades to inline execution: a chunk that blocked on sub-chunks of + // the same bounded executor could starve it into deadlock (every worker waiting on chunks that + // can never be scheduled). + private static final ThreadLocal IN_BODY = ThreadLocal.withInitial(() -> Boolean.FALSE); + + private final ExecutorService executor; + private final int parallelism; + + ChunkingParallelExecutor(ExecutorService executor, int parallelism) { + this.executor = executor; + this.parallelism = parallelism; + } + + @Override + public void forEachInt(int upperBound, IntConsumer body) { + if (upperBound <= 0) { + return; + } + if (IN_BODY.get()) { + for (int i = 0; i < upperBound; i++) { + body.accept(i); + } + return; + } + int chunks = Math.min(upperBound, CHUNKS_PER_WORKER * parallelism); + Drain drain = new Drain(Integer.MAX_VALUE); // chunk count is fixed and small: no window needed + for (int c = 0; c < chunks && drain.healthy(); c++) { + int start = (int) ((long) upperBound * c / chunks); + int end = (int) ((long) upperBound * (c + 1) / chunks); + drain.submit(() -> { + for (int i = start; i < end; i++) { + body.accept(i); + } + }); + } + drain.finish(); + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + if (IN_BODY.get()) { + source.forEach(body); + return; + } + Drain drain = new Drain(2 * parallelism); + try { + PrimitiveIterator.OfInt it = source.iterator(); + int[] batch = new int[BATCH_SIZE]; + int n = 0; + while (drain.healthy() && it.hasNext()) { + batch[n++] = it.nextInt(); + if (n == BATCH_SIZE) { + int[] b = batch; + int len = n; + batch = new int[BATCH_SIZE]; + n = 0; + drain.submit(() -> { + for (int i = 0; i < len; i++) { + body.accept(b[i]); + } + }); + } + } + if (drain.healthy() && n > 0) { + int[] b = batch; + int len = n; + drain.submit(() -> { + for (int i = 0; i < len; i++) { + body.accept(b[i]); + } + }); + } + } catch (Throwable t) { + drain.fail(t); // traversal failure: recorded, then drained below like any other + } + drain.finish(); + } + + @Override + public void forEach(Stream source, Consumer body) { + if (IN_BODY.get()) { + source.forEach(body); + return; + } + Drain drain = new Drain(2 * parallelism); + try { + Iterator it = source.iterator(); + List batch = new ArrayList<>(BATCH_SIZE); + while (drain.healthy() && it.hasNext()) { + batch.add(it.next()); + if (batch.size() == BATCH_SIZE) { + List b = batch; + batch = new ArrayList<>(BATCH_SIZE); + drain.submit(() -> b.forEach(body)); + } + } + if (drain.healthy() && !batch.isEmpty()) { + List b = batch; + drain.submit(() -> b.forEach(body)); + } + } catch (Throwable t) { + drain.fail(t); + } + drain.finish(); + } + + /** + * Chunk bookkeeping with the drain-before-unwind discipline used elsewhere in jvector: once any + * chunk (or the traversal) fails, chunks that have not begun skip themselves via {@code aborted}, + * but every started chunk is waited out — the caller never unwinds (and so never releases + * resources) beneath a still-running body. {@link Future#cancel} is deliberately never used: + * {@code cancel(false)} succeeds on a running {@code FutureTask} (the flag only governs + * interruption), making {@code get()} return while the body still executes — the exact unwind + * this class exists to prevent. Interruption is noted, honored after the drain, and never used + * to abort a running chunk. + */ + private final class Drain { + private final int window; + private final ArrayDeque> inFlight = new ArrayDeque<>(); + // Written by the orchestrating thread, read by workers: a failing iteration stops issuing + // chunks here and not-yet-started chunks become no-ops. + private volatile boolean aborted; + private Throwable failure; + private boolean interrupted; + + Drain(int window) { + this.window = window; + } + + boolean healthy() { + return failure == null && !interrupted; + } + + void fail(Throwable t) { + if (failure == null) { + failure = t; + } + aborted = true; + } + + void submit(Runnable chunk) { + if (!healthy()) { + return; + } + try { + inFlight.add(executor.submit(() -> { + if (aborted) { + return; // iteration is failing: skip a chunk that has not begun + } + IN_BODY.set(Boolean.TRUE); + try { + chunk.run(); + } finally { + IN_BODY.remove(); + } + })); + } catch (RejectedExecutionException e) { + fail(e); + return; + } + if (inFlight.size() >= window) { + settle(inFlight.poll()); + } + } + + /** Waits {@code f} out (across interrupts), recording any failure. */ + private void settle(Future f) { + while (true) { + try { + f.get(); + return; + } catch (ExecutionException e) { + fail(e.getCause()); + return; + } catch (InterruptedException e) { + interrupted = true; // note it and keep waiting: no unwind under running work + aborted = true; // unstarted chunks need not run + } + } + } + + /** Awaits every outstanding chunk, then rethrows the first failure and restores the interrupt flag. */ + void finish() { + while (!inFlight.isEmpty()) { + settle(inFlight.poll()); + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + if (failure instanceof Error) { + throw (Error) failure; + } + if (failure != null) { + throw new RuntimeException(failure); + } + if (interrupted) { + throw new RuntimeException(new InterruptedException("interrupted while awaiting parallel iteration")); + } + } + } +} diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java new file mode 100644 index 000000000..727457437 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java @@ -0,0 +1,178 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import java.util.Objects; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ForkJoinPool; +import java.util.function.Consumer; +import java.util.function.IntConsumer; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +/** + * Runs a {@link GraphIndexBuilder}'s internal build/finalize iterations to completion, blocking the + * calling thread until every element has been processed. The implementation decides how the + * iteration is distributed. + *

+ * This is the seam that lets an embedder bound vector-graph construction to its own thread budget — + * e.g. one thread per compaction — instead of a jvector-owned all-core pool. It is the build/finalize + * counterpart to the caller-runs executor injection already available on the compaction merge path. + * + *

Choosing a factory

+ * The provided implementations trade off differently; the caveats below are relative to each other: + * + * Common to all three: the calling thread blocks until the iteration settles, and the first + * observed body failure is rethrown to the caller; whether other elements still run after a + * failure differs per the caveats above. + */ +public interface ParallelExecutor { + /** + * Runs {@code body} for each {@code i} in {@code [0, upperBound)}, blocking until all complete. + * + * @param upperBound the exclusive upper bound of the index range (may be {@code 0}) + * @param body the action to apply to each index + */ + void forEachInt(int upperBound, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream of primitive ints to iterate + * @param body the action to apply to each element + */ + void forEach(IntStream source, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream to iterate + * @param body the action to apply to each element + * @param the stream element type + */ + void forEach(Stream source, Consumer body); + + /** + * Returns an executor backed by {@code pool}: each iteration is hosted as a parallel stream on + * that pool and the calling thread blocks on the result. This reproduces the behavior of the + * {@code ForkJoinPool}-based {@link GraphIndexBuilder} constructors. See the class javadoc for + * how its caveats compare with the other factories. + * + * @param pool the pool that hosts the parallel iterations + * @return a pool-backed {@code ParallelExecutor} + */ + static ParallelExecutor forkJoin(ForkJoinPool pool) { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + pool.submit(() -> IntStream.range(0, upperBound).parallel().forEach(body)).join(); + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + + @Override + public void forEach(Stream source, Consumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + }; + } + + /** + * Returns an executor that runs every iteration sequentially on the calling thread — no worker + * threads, no pool, and the common pool is left untouched. Graph structure and recall are + * equivalent to the {@link #forkJoin(ForkJoinPool)} path; only wall-clock and thread usage differ. + * See the class javadoc for how its caveats compare with the other factories. + * + * @return a caller-runs {@code ParallelExecutor} + */ + static ParallelExecutor callerRuns() { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + for (int i = 0; i < upperBound; i++) { + body.accept(i); + } + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + source.forEach(body); + } + + @Override + public void forEach(Stream source, Consumer body) { + source.forEach(body); + } + }; + } + + /** + * Returns an executor backed by a caller-supplied {@code ExecutorService}: iterations are split + * into chunks submitted to {@code executor} while the calling thread traverses the source and + * bounds the in-flight window. Parallel streams cannot be hosted on a generic + * {@code ExecutorService} (inside its workers they would silently run on the common pool), so + * this adapter distributes only the body; see the class javadoc for the full caveat + * comparison with {@link #forkJoin(ForkJoinPool)} and {@link #callerRuns()}. + *

+ * {@code parallelism} must be stated explicitly because {@code ExecutorService} does not expose + * its width: state the executor's actual thread count — a larger value merely queues, a smaller + * one under-uses it. Passing a {@link ForkJoinPool} delegates to {@link #forkJoin(ForkJoinPool)} + * (whole-pipeline stream decomposition is strictly better there) and {@code parallelism} is + * ignored in that case. + * + * @param executor runs the chunked iterations; its lifecycle remains the caller's (never shut down here) + * @param parallelism the intended number of chunks in flight, typically the executor's thread + * count; must be {@code >= 1} + * @return a chunk-submitting {@code ParallelExecutor} + * @throws IllegalArgumentException if {@code parallelism < 1} + */ + static ParallelExecutor over(ExecutorService executor, int parallelism) { + Objects.requireNonNull(executor, "executor"); + if (parallelism < 1) { + throw new IllegalArgumentException("parallelism must be >= 1, got " + parallelism); + } + if (executor instanceof ForkJoinPool) { + return forkJoin((ForkJoinPool) executor); + } + return new ChunkingParallelExecutor(executor, parallelism); + } +} 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..2dbaf5cd1 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/CompactionDestination.java @@ -0,0 +1,91 @@ +/* + * 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. The destination itself is stateless + * configuration an embedder can build once and hold; each {@code compact(...)} call + * {@link #reserve() reserves} its own {@link OutputReservation}, commits it on success, and always + * closes it. + * + * 

{@code
+ *   try (CompactionDestination.OutputReservation r = destination.reserve()) {
+ *       // ...compactor writes the graph into r.file() at r.startOffset()...
+ *       r.commit(bodyLength);   // success: body written & durable; embedder finalizes its footer
+ *   }                           // close() always runs; no commit() => reservation released, partial output discarded
+ * }
+ * + *

The compactor needs a real file (it uses a memory-mapped read-back during refinement and a + * random-access writer), so an {@link OutputReservation} 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, r.startOffset())}. + */ +@FunctionalInterface +@Experimental +public interface CompactionDestination { + + /** Reserve a fresh output region for one compaction. */ + OutputReservation reserve() throws IOException; + + /** + * One compaction's reserved output region plus its commit/abort lifecycle: live, single-use + * state made against the destination, fulfilled by {@link #commit} or released by a + * {@link #close} without one. + */ + interface OutputReservation 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(); + + /** + * Fulfills the reservation. 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: the reservation is released, the embedder discards the partial + * output, and embedder resources are freed. + */ + @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..b9a972fd6 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FileCompactionDestination.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.graph.disk; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Objects; + +/** + * {@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) { + Objects.requireNonNull(path, "path"); + this.path = path; + } + + @Override + public OutputReservation reserve() { + return new OutputReservation() { + 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/OnDiskGraphIndex.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java index ba34e2cb0..915eb08c1 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 @@ -227,7 +227,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. @@ -239,6 +242,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. @@ -270,6 +283,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. */ @@ -494,6 +510,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; @@ -512,6 +543,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 @@ -528,6 +560,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 @@ -588,8 +621,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/util/RuntimeMode.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/RuntimeMode.java new file mode 100644 index 000000000..50758893b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/RuntimeMode.java @@ -0,0 +1,89 @@ +/* + * 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; + +import java.util.Locale; +import java.util.logging.Logger; + +/** + * Library-wide runtime mode, from the {@code jvector.mode} system property. + * + *

Two modes: {@code prod} (or {@code production}) and {@code dev} (or + * {@code development}), case-insensitive. The mode gates diagnostic + * walks — computations whose value is re-derived from first principles by + * traversing a structure, such as {@link + * io.github.jbellis.jvector.quantization.PQVectors#ramBytesUsed()} summing + * every compressed chunk. In production mode such walks are categorically + * replaced by incrementally-maintained values; in development mode they run in + * full, so a drifted cache or accounting bug is observable. + * + *

Default is production. An embedding host once called the chunk + * walk once per inserted vector, turning index-build accounting into + * O(n²) — a single compaction burned two CPU-hours inside + * {@code ramBytesUsed} while build workers starved. A diagnostic full-walk is + * something a developer opts into, not something a production host should + * have to know to opt out of. + * + *

Unrecognized values log a warning and resolve to production, not + * development: falling back to the diagnostic mode on a typo would silently + * reintroduce exactly the pathology above. + */ +public final class RuntimeMode { + private static final Logger LOG = Logger.getLogger(RuntimeMode.class.getName()); + + public static final String PROPERTY = "jvector.mode"; + + private static final boolean DEVELOPMENT = + parseIsDevelopment(System.getProperty(PROPERTY), LOG); + + private RuntimeMode() { + } + + /** True when diagnostic walks should run in full. */ + public static boolean isDevelopment() { + return DEVELOPMENT; + } + + /** True when diagnostic walks are replaced by maintained values. */ + public static boolean isProduction() { + return !DEVELOPMENT; + } + + /** + * Pure parse, exposed for tests (the static mode is fixed at class load). + */ + static boolean parseIsDevelopment(String raw, Logger log) { + if (raw == null) { + return false; + } + switch (raw.trim().toLowerCase(Locale.ROOT)) { + case "dev": + case "development": + return true; + case "prod": + case "production": + case "": + return false; + default: + log.warning(() -> PROPERTY + "=" + raw + + " is not recognized (expected prod|production|dev|development); " + + "using production. Development mode re-enables full diagnostic " + + "walks and must be asked for exactly."); + return false; + } + } +} 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..4f31f0431 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/LeakyBucketLimiter.java @@ -0,0 +1,68 @@ +/* + * 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 #startPhase} is inherited as a no-op scope factory: 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..481f88406 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java @@ -0,0 +1,114 @@ +/* + * 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 PhaseScope startPhase(WorkStage stage) { return PhaseScope.NOOP; } + + @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 #startPhase} + * returns a no-op scope 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 phase start, + * each {@link PhaseScope#onProgress} report, each phase completion, and 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 PhaseScope startPhase(WorkStage stage) { + sink.accept("phase[" + stage.name() + "] started"); + PhaseScope scope = d.startPhase(stage); + return new PhaseScope() { + @Override + public void onProgress(long completed, long total) { + sink.accept("progress[" + stage.name() + "] " + completed + "/" + (total < 0 ? "?" : Long.toString(total))); + scope.onProgress(completed, total); + } + + @Override + public void close() { + try { + scope.close(); + } finally { + sink.accept("phase[" + stage.name() + "] completed"); + } + } + }; + } + + @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..90a60011b --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java @@ -0,0 +1,72 @@ +/* + * 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: observes the phases of a long-running operation. {@link #startPhase} is + * the single entry point; progress is reported through the returned {@link PhaseScope}, so the + * scope itself is the capability to report — progress for a phase that was never started is + * unrepresentable, and per-phase implementation state (a progress bar, a long-task timer) lives + * in the scope instance rather than in a map keyed by {@link WorkStage}. Concurrent phases, even + * of the same stage, are distinguished by scope identity. + * + *

Best-effort and cheap: implementations must not throw from {@code startPhase}, + * {@link PhaseScope#onProgress}, or {@link PhaseScope#close} (the caller invokes these on its + * orchestrating thread and treats them as fire-and-forget). The minimal consumer is a single + * expression — {@code stage -> (completed, total) -> bar.update(stage, completed, total)} — since + * {@code close()} defaults to a no-op. See {@link ProgressLimiter} for the melded progress + + * throttle surface that most consumers accept. + */ +@Experimental +@FunctionalInterface +public interface ProgressTracker { + /** + * Starts one phase of the operation. The returned scope receives that phase's progress and + * must be closed exactly once, normally with try-with-resources. + * + * @param stage identifies the phase's stage, in consumer-defined terms + */ + PhaseScope startPhase(WorkStage stage); + + /** + * One started phase: receives its progress reports and, on {@link #close}, its completion. + * {@code close()} defaults to a no-op so a progress-only implementation stays a single lambda. + */ + @FunctionalInterface + interface PhaseScope extends AutoCloseable { + /** + * Reports progress for this phase. + * + * @param completed work done so far, in stage-defined units; monotonically non-decreasing + * within the phase + * @param total total work for the phase, or {@code -1} if not yet known + */ + void onProgress(long completed, long total); + + /** Marks the end of the phase. Called exactly once; defaults to a no-op. */ + @Override + default void close() { } + + /** A scope that discards every update. */ + PhaseScope NOOP = (completed, total) -> { }; + } + + /** A tracker that discards every phase. */ + ProgressTracker NOOP = stage -> PhaseScope.NOOP; +} 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/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/TestParallelExecutor.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestParallelExecutor.java new file mode 100644 index 000000000..744034ddd --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/TestParallelExecutor.java @@ -0,0 +1,286 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import org.junit.Test; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentSkipListSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ForkJoinPool; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicIntegerArray; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestParallelExecutor { + + // ---- over(...) argument validation ---- + + @Test(expected = NullPointerException.class) + public void overRejectsNullExecutor() { + ParallelExecutor.over(null, 1); + } + + @Test + public void overRejectsNonPositiveParallelism() { + ExecutorService es = Executors.newSingleThreadExecutor(); + try { + for (int bad : new int[]{0, -1}) { + try { + ParallelExecutor.over(es, bad); + fail("expected IllegalArgumentException for parallelism " + bad); + } catch (IllegalArgumentException expected) { + // ok + } + } + } finally { + es.shutdown(); + } + } + + // ---- completeness: every element exactly once ---- + + @Test + public void chunkingForEachIntCoversRangeExactlyOnce() { + ExecutorService es = Executors.newFixedThreadPool(3); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 3); + AtomicIntegerArray hits = new AtomicIntegerArray(10_000); + pe.forEachInt(10_000, hits::incrementAndGet); + for (int i = 0; i < hits.length(); i++) { + assertEquals("index " + i, 1, hits.get(i)); + } + // Zero-size range: no calls, no error. + AtomicInteger calls = new AtomicInteger(); + pe.forEachInt(0, i -> calls.incrementAndGet()); + assertEquals(0, calls.get()); + } finally { + es.shutdown(); + } + } + + @Test + public void chunkingForEachIntStreamCoversAllElements() { + ExecutorService es = Executors.newFixedThreadPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 2); + // A filtered (non-SIZED) stream exercises traversal-side batching. + Set expected = IntStream.range(0, 5_000).filter(i -> i % 3 == 0) + .boxed().collect(Collectors.toSet()); + Set seen = new ConcurrentSkipListSet<>(); + pe.forEach(IntStream.range(0, 5_000).filter(i -> i % 3 == 0), seen::add); + assertEquals(expected, seen); + } finally { + es.shutdown(); + } + } + + @Test + public void chunkingForEachGenericStreamCoversAllElements() { + ExecutorService es = Executors.newFixedThreadPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 2); + Set expected = IntStream.range(0, 1_000).mapToObj(Integer::toString) + .collect(Collectors.toSet()); + Set seen = new ConcurrentSkipListSet<>(); + pe.forEach(IntStream.range(0, 1_000).mapToObj(Integer::toString), seen::add); + assertEquals(expected, seen); + } finally { + es.shutdown(); + } + } + + // ---- thread placement: what is (and is not) distributed ---- + + @Test + public void chunkingTraversesSourceOnCallingThreadAndRunsBodyOnWorkers() { + ExecutorService es = Executors.newFixedThreadPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 2); + Thread caller = Thread.currentThread(); + Set traversalThreads = ConcurrentHashMap.newKeySet(); + Set bodyThreads = ConcurrentHashMap.newKeySet(); + IntStream source = IntStream.range(0, 2_000).map(i -> { + traversalThreads.add(Thread.currentThread()); + return i; + }); + pe.forEach(source, i -> bodyThreads.add(Thread.currentThread())); + assertEquals("upstream stages run on the calling thread only", + Set.of(caller), traversalThreads); + assertFalse("the body must not run on the calling thread", bodyThreads.contains(caller)); + } finally { + es.shutdown(); + } + } + + @Test + public void overForkJoinPoolDelegatesToStreamDecomposition() { + ForkJoinPool pool = new ForkJoinPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(pool, 5); + Thread caller = Thread.currentThread(); + AtomicBoolean traversedOnCaller = new AtomicBoolean(false); + IntStream source = IntStream.range(0, 2_000).map(i -> { + if (Thread.currentThread() == caller) { + traversedOnCaller.set(true); + } + return i; + }); + AtomicInteger count = new AtomicInteger(); + pe.forEach(source, i -> count.incrementAndGet()); + assertEquals(2_000, count.get()); + assertFalse("a ForkJoinPool must delegate to forkJoin(): upstream stages run inside the pool", + traversedOnCaller.get()); + } finally { + pool.shutdown(); + } + } + + @Test + public void callerRunsExecutesInOrderOnCallingThread() { + ParallelExecutor pe = ParallelExecutor.callerRuns(); + Thread caller = Thread.currentThread(); + List order = new ArrayList<>(); + pe.forEachInt(100, i -> { + assertSame(caller, Thread.currentThread()); + order.add(i); + }); + assertEquals(IntStream.range(0, 100).boxed().collect(Collectors.toList()), order); + } + + // ---- failure and interrupt semantics ---- + + @Test + public void chunkingBodyFailurePropagatesWithNothingLeftRunning() { + ExecutorService es = Executors.newFixedThreadPool(4); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 4); + RuntimeException marker = new RuntimeException("marker"); + AtomicInteger active = new AtomicInteger(); + try { + pe.forEachInt(100, i -> { + active.incrementAndGet(); + try { + if (i == 41) { + throw marker; + } + Thread.sleep(5); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + active.decrementAndGet(); + } + }); + fail("expected the body failure to propagate"); + } catch (RuntimeException e) { + assertSame("the first observed failure must propagate as-is", marker, e); + } + assertEquals("no body may still be running once the failure unwinds", 0, active.get()); + } finally { + es.shutdown(); + } + } + + @Test(timeout = 10_000) + public void chunkingInterruptDrainsStartedWorkAndRestoresFlag() throws Exception { + ExecutorService es = Executors.newFixedThreadPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 2); + AtomicInteger active = new AtomicInteger(); + AtomicReference caught = new AtomicReference<>(); + AtomicBoolean flagRestored = new AtomicBoolean(); + AtomicInteger activeAtUnwind = new AtomicInteger(-1); + Thread runner = new Thread(() -> { + try { + pe.forEachInt(64, i -> { + active.incrementAndGet(); + try { + Thread.sleep(20); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + active.decrementAndGet(); + } + }); + } catch (Throwable t) { + caught.set(t); + flagRestored.set(Thread.currentThread().isInterrupted()); + activeAtUnwind.set(active.get()); + } + }, "interrupted-orchestrator"); + runner.start(); + Thread.sleep(100); + runner.interrupt(); + runner.join(8_000); + + assertFalse("orchestrator must not hang", runner.isAlive()); + assertTrue("expected a RuntimeException, got " + caught.get(), + caught.get() instanceof RuntimeException); + assertTrue("cause must be the InterruptedException, got " + caught.get().getCause(), + caught.get().getCause() instanceof InterruptedException); + assertTrue("interrupt flag must be restored before unwinding", flagRestored.get()); + assertEquals("no body may still be running once the interrupt unwinds", 0, activeAtUnwind.get()); + } finally { + es.shutdown(); + } + } + + @Test + public void chunkingRejectedExecutionPropagates() { + ExecutorService es = Executors.newSingleThreadExecutor(); + es.shutdown(); + ParallelExecutor pe = ParallelExecutor.over(es, 1); + try { + pe.forEachInt(10, i -> { }); + fail("expected RejectedExecutionException"); + } catch (RuntimeException e) { + assertTrue("expected RejectedExecutionException, got " + e, + e instanceof java.util.concurrent.RejectedExecutionException); + } + } + + // ---- reentrancy ---- + + @Test(timeout = 10_000) + public void chunkingNestedUseRunsInlineWithoutDeadlock() { + ExecutorService es = Executors.newFixedThreadPool(2); + try { + ParallelExecutor pe = ParallelExecutor.over(es, 2); + AtomicInteger inner = new AtomicInteger(); + // Without the inline-degrade guard, every worker blocks awaiting sub-chunks that can + // never be scheduled on the bounded pool; the timeout turns that hang into a failure. + pe.forEachInt(8, i -> pe.forEachInt(10, j -> inner.incrementAndGet())); + assertEquals(80, inner.get()); + } finally { + es.shutdown(); + } + } +} diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/util/TestRuntimeMode.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/TestRuntimeMode.java new file mode 100644 index 000000000..d05e6f50d --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/TestRuntimeMode.java @@ -0,0 +1,58 @@ +/* + * 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; + +import java.util.logging.Logger; + +import org.junit.Test; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class TestRuntimeMode { + private static final Logger LOG = Logger.getLogger(TestRuntimeMode.class.getName()); + + @Test + public void terseAndVerboseSpellingsBothParse() { + assertTrue(RuntimeMode.parseIsDevelopment("dev", LOG)); + assertTrue(RuntimeMode.parseIsDevelopment("development", LOG)); + assertTrue(RuntimeMode.parseIsDevelopment("DEVELOPMENT", LOG)); + assertTrue(RuntimeMode.parseIsDevelopment(" Dev ", LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("prod", LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("production", LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("PROD", LOG)); + } + + @Test + public void defaultIsProduction() { + assertFalse("unset must be production — the diagnostic walk is opt-in", + RuntimeMode.parseIsDevelopment(null, LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("", LOG)); + } + + /** + * A typo must NOT fall back to development: that would silently + * reintroduce the per-insert full-walk pathology the gate exists to + * prevent. + */ + @Test + public void unknownValuesResolveToProduction() { + assertFalse(RuntimeMode.parseIsDevelopment("porduction", LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("debug", LOG)); + assertFalse(RuntimeMode.parseIsDevelopment("true", LOG)); + } +} 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..ab2c85ac8 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java @@ -0,0 +1,356 @@ +/* + * 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 + try (ProgressTracker.PhaseScope scope = limiter.startPhase(STAGE)) { + scope.onProgress(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); + + try (ProgressTracker.PhaseScope scope = limiter.startPhase(STAGE)) { + scope.onProgress(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 loggingDelegatesPhaseLifetime() { + RecordingLimiter delegate = new RecordingLimiter(); + List log = Collections.synchronizedList(new ArrayList<>()); + ProgressLimiter limiter = ProgressLimiter.logging(delegate, log::add); + + try (ProgressTracker.PhaseScope ignored = limiter.startPhase(STAGE)) { + assertEquals(1, delegate.phaseStarts.get()); + assertEquals(0, delegate.phaseCloses.get()); + } + + assertEquals(1, delegate.phaseCloses.get()); + assertTrue(log.stream().anyMatch(s -> s.contains("phase[TEST] started"))); + assertTrue(log.stream().anyMatch(s -> s.contains("phase[TEST] completed"))); + } + + @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"))); + try (ProgressTracker.PhaseScope scope = limiter.startPhase(STAGE)) { + scope.onProgress(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); + try (ProgressTracker.PhaseScope scope = ProgressLimiter.UNLIMITED.startPhase(STAGE)) { + scope.onProgress(7, -1); // no-op, must not throw + } + + // Facet no-op constants exist and are safe. + WorkLimiter.Grant.NOOP.close(); + ProgressTracker.PhaseScope.NOOP.onProgress(1, 1); + ProgressTracker.NOOP.startPhase(STAGE).onProgress(1, 1); + try (Grant g = WorkLimiter.UNLIMITED.acquire(99)) { + assertNotNull(g); + } + } + + @Test + public void facetsAreIndependentlyOverridable() throws Exception { + // Tracker-only: overrides startPhase, inherits no-op acquire. + AtomicInteger progressSeen = new AtomicInteger(); + ProgressLimiter trackerOnly = new ProgressLimiter() { + @Override public PhaseScope startPhase(WorkStage stage) { + return (completed, total) -> progressSeen.incrementAndGet(); + } + }; + try (Grant g = trackerOnly.acquire(1_000_000)) { // inherited no-op: must not block + assertNotNull(g); + } + try (ProgressTracker.PhaseScope scope = trackerOnly.startPhase(STAGE)) { + scope.onProgress(1, 1); + } + assertEquals(1, progressSeen.get()); + + // Throttle-only: overrides acquire, inherits no-op startPhase. + AtomicInteger acquireSeen = new AtomicInteger(); + ProgressLimiter throttleOnly = new ProgressLimiter() { + @Override public Grant acquire(long amount) { + acquireSeen.incrementAndGet(); + return Grant.NOOP; + } + }; + throttleOnly.startPhase(STAGE).onProgress(1, 1); // inherited no-op scope: must not throw + throttleOnly.acquire(5).close(); + assertEquals(1, acquireSeen.get()); + } + + @Test + public void progressOnlyTrackerIsASingleLambda() { + List seen = new ArrayList<>(); + // The minimal consumer shape the SPI promises: one expression, default no-op close(). + ProgressTracker tracker = stage -> (completed, total) -> seen.add(stage.name() + ":" + completed + "/" + total); + try (ProgressTracker.PhaseScope scope = tracker.startPhase(STAGE)) { + scope.onProgress(1, 4); + scope.onProgress(4, 4); + } + assertEquals(List.of("TEST:1/4", "TEST:4/4"), seen); + } + + @Test + public void concurrentPhasesOfSameStageGetIndependentScopes() { + RecordingLimiter delegate = new RecordingLimiter(); + ProgressTracker.PhaseScope first = delegate.startPhase(STAGE); + ProgressTracker.PhaseScope second = delegate.startPhase(STAGE); + assertEquals(2, delegate.phaseStarts.get()); + + first.onProgress(1, 10); + second.onProgress(9, 10); + first.close(); + assertEquals("closing one scope must not close the other", 1, delegate.phaseCloses.get()); + second.close(); + assertEquals(2, delegate.phaseCloses.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(); + final AtomicInteger phaseStarts = new AtomicInteger(); + final AtomicInteger phaseCloses = new AtomicInteger(); + volatile long lastCompleted, lastTotal, lastAmount; + + @Override + public Grant acquire(long amount) { + acquireCalls.incrementAndGet(); + lastAmount = amount; + return grantCloses::incrementAndGet; + } + + @Override + public PhaseScope startPhase(WorkStage stage) { + phaseStarts.incrementAndGet(); + return new PhaseScope() { + @Override + public void onProgress(long completed, long total) { + progressCalls.incrementAndGet(); + lastCompleted = completed; + lastTotal = total; + } + + @Override + public void close() { + phaseCloses.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; + } + } +}