From 7c9375d17a6f7700d527ede4c1c891bc770751ff Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 16:13:01 +0000 Subject: [PATCH 1/5] util/work: cooperative work-limiting and progress interface Introduce an opt-in control surface a host can install on long-running jvector operations, without touching any existing call site: - WorkStage / WorkLimiter / ProgressTracker: the admission + reporting primitives (throttle down, progress up). - ProgressLimiter: the single combined facet a host implements. - LeakyBucketLimiter: a default rate-limited WorkLimiter. Purely additive. Nothing in the library calls these yet; they establish the seam that host-driven operations (e.g. compaction) can later be rewritten against. --- .../jvector/util/work/LeakyBucketLimiter.java | 67 ++++ .../jvector/util/work/ProgressLimiter.java | 110 ++++++ .../jvector/util/work/ProgressTracker.java | 61 ++++ .../jvector/util/work/WorkLimiter.java | 59 ++++ .../jbellis/jvector/util/work/WorkStage.java | 33 ++ .../util/work/TestProgressLimiter.java | 314 ++++++++++++++++++ 6 files changed, 644 insertions(+) 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/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..1ae328ee6 --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressLimiter.java @@ -0,0 +1,110 @@ +/* + * 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 PhaseScope startPhase(WorkStage stage) { + sink.accept("phase[" + stage.name() + "] started"); + PhaseScope scope = d.startPhase(stage); + return () -> { + 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..4e596c42f --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/work/ProgressTracker.java @@ -0,0 +1,61 @@ +/* + * 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); + + /** + * Starts observing the lifetime of {@code stage}. The returned scope must be closed exactly + * once, normally with try-with-resources. Implementations may use this for a long-task timer; + * the default keeps instrumentation optional for embedders that only consume progress. + */ + default PhaseScope startPhase(WorkStage stage) { + return PhaseScope.NOOP; + } + + /** Closeable lifetime token returned by {@link #startPhase}. */ + @FunctionalInterface + interface PhaseScope extends AutoCloseable { + @Override + void close(); + + PhaseScope NOOP = () -> { }; + } + + /** 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/util/work/TestProgressLimiter.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.java new file mode 100644 index 000000000..e446ffc61 --- /dev/null +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/util/work/TestProgressLimiter.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.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 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"))); + 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 + ProgressLimiter.UNLIMITED.startPhase(STAGE).close(); + + // 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(); + final AtomicInteger phaseStarts = new AtomicInteger(); + final AtomicInteger phaseCloses = 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; + } + + @Override + public PhaseScope startPhase(WorkStage stage) { + phaseStarts.incrementAndGet(); + return 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; + } + } +} From 15579ba301aff8415da4c01146dfa0928264aeaa Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 16:13:01 +0000 Subject: [PATCH 2/5] disk: pluggable output sink for compaction writes Add a host-supplied destination abstraction so compaction output can be written through a caller-owned channel instead of always allocating its own file: - SeekableSink / FileChannelSeekableSink: a minimal seekable byte sink and its file-backed implementation. - CompactionDestination / FileCompactionDestination: the compaction output target, resolvable to a SeekableSink. Purely additive interface types; no existing code is wired to them on this branch. --- .../jvector/disk/FileChannelSeekableSink.java | 70 +++++++++++++++ .../jbellis/jvector/disk/SeekableSink.java | 65 ++++++++++++++ .../graph/disk/CompactionDestination.java | 83 ++++++++++++++++++ .../graph/disk/FileCompactionDestination.java | 66 ++++++++++++++ .../jvector/disk/TestSeekableSink.java | 85 +++++++++++++++++++ 5 files changed, 369 insertions(+) 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/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/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-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); + } +} From 914f83e5c91223e91c413aca5ea455197f696c72 Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 16:30:25 +0000 Subject: [PATCH 3/5] graph: ParallelExecutor abstraction for host-provided execution Add ParallelExecutor, a minimal parallel-for abstraction that lets an embedding host supply its own execution strategy -- its own pool, or caller-runs on the calling thread -- instead of jvector reaching for ForkJoinPool.commonPool(): - forkJoin(pool): run on a caller-supplied ForkJoinPool - callerRuns(): run inline, no work escaping to a shared pool - forEachInt / forEach: the parallel-for entry points Additive: a standalone interface with no consumers on this branch. It establishes the execution seam that build, quantization, and compaction paths can be rewritten against later. --- .../jvector/graph/ParallelExecutor.java | 118 ++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java new file mode 100644 index 000000000..e4ac29d6d --- /dev/null +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ParallelExecutor.java @@ -0,0 +1,118 @@ +/* + * Copyright DataStax, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.github.jbellis.jvector.graph; + +import java.util.concurrent.ForkJoinPool; +import java.util.function.Consumer; +import java.util.function.IntConsumer; +import java.util.stream.IntStream; +import java.util.stream.Stream; + +/** + * Runs a {@link GraphIndexBuilder}'s internal build/finalize iterations to completion, blocking the + * calling thread until every element has been processed. The implementation decides how the + * iteration is distributed: {@link #forkJoin(ForkJoinPool)} hosts a parallel stream on a dedicated + * pool (the historical behavior), while {@link #callerRuns()} runs everything sequentially on the + * calling thread with no worker threads and no pool. + *

+ * This is the seam that lets an embedder bound vector-graph construction to its own thread budget — + * e.g. one thread per compaction — instead of a jvector-owned all-core pool. It is the build/finalize + * counterpart to the caller-runs executor injection already available on the compaction merge path. + */ +public interface ParallelExecutor { + /** + * Runs {@code body} for each {@code i} in {@code [0, upperBound)}, blocking until all complete. + * + * @param upperBound the exclusive upper bound of the index range (may be {@code 0}) + * @param body the action to apply to each index + */ + void forEachInt(int upperBound, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream of primitive ints to iterate + * @param body the action to apply to each element + */ + void forEach(IntStream source, IntConsumer body); + + /** + * Runs {@code body} for each element produced by {@code source}, blocking until all complete. + * Callers pass a sequential stream; the implementation decides whether to parallelize it. + * + * @param source the (sequential) stream to iterate + * @param body the action to apply to each element + * @param the stream element type + */ + void forEach(Stream source, Consumer body); + + /** + * Returns an executor backed by {@code pool}: each iteration is hosted as a parallel stream on + * that pool and the calling thread blocks on the result. This reproduces the behavior of the + * {@code ForkJoinPool}-based {@link GraphIndexBuilder} constructors. + * + * @param pool the pool that hosts the parallel iterations + * @return a pool-backed {@code ParallelExecutor} + */ + static ParallelExecutor forkJoin(ForkJoinPool pool) { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + pool.submit(() -> IntStream.range(0, upperBound).parallel().forEach(body)).join(); + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + + @Override + public void forEach(Stream source, Consumer body) { + pool.submit(() -> source.parallel().forEach(body)).join(); + } + }; + } + + /** + * Returns an executor that runs every iteration sequentially on the calling thread — no worker + * threads, no pool, and the common pool is left untouched. Graph structure and recall are + * equivalent to the {@link #forkJoin(ForkJoinPool)} path; only wall-clock and thread usage differ. + * + * @return a caller-runs {@code ParallelExecutor} + */ + static ParallelExecutor callerRuns() { + return new ParallelExecutor() { + @Override + public void forEachInt(int upperBound, IntConsumer body) { + for (int i = 0; i < upperBound; i++) { + body.accept(i); + } + } + + @Override + public void forEach(IntStream source, IntConsumer body) { + source.forEach(body); + } + + @Override + public void forEach(Stream source, Consumer body) { + source.forEach(body); + } + }; + } +} From 6b528cbc8876c316fb0af9d4f848a585beaa7fab Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 16:30:25 +0000 Subject: [PATCH 4/5] util: RuntimeMode gate for opt-in diagnostic work Add RuntimeMode, a process-level switch (jvector.mode) separating production from diagnostic runs. Unset is production; diagnostic-only work (e.g. verification walks) sits behind this gate and must be opted into explicitly. Additive: a standalone type with no call sites on this branch. It establishes the gate that diagnostic code paths can consult once wired. --- .../jbellis/jvector/util/RuntimeMode.java | 89 +++++++++++++++++++ .../jbellis/jvector/util/TestRuntimeMode.java | 58 ++++++++++++ 2 files changed, 147 insertions(+) create mode 100644 jvector-base/src/main/java/io/github/jbellis/jvector/util/RuntimeMode.java create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/util/TestRuntimeMode.java 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-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)); + } +} From d55d0d574c9be6700d9e3c2c84d874c304daffbc Mon Sep 17 00:00:00 2001 From: jshook Date: Wed, 12 Aug 2026 16:40:14 +0000 Subject: [PATCH 5/5] disk: memory-safety guards for host-managed reads Harden the read path for embedding hosts that manage their own mmap lifecycle, where a stale offset or a close racing an in-flight read faults the JVM (SIGSEGV) instead of throwing: - OnDiskGraphIndex: bound record reads. A node id outside the graph would become a wild offset into the mapped file (garbage, or a fault); requireValidNode() rejects it up front, and a stale/corrupt neighbor block now fails with IllegalStateException instead of consuming garbage ints. - ReaderSupplier / SimpleMappedReader: document the close() contract -- the raw-release (immediate unmap -> SIGSEGV) vs coordinated (liveness handshake -> IllegalStateException) families -- so a host knows a supplier must not be closed until every vended reader is quiescent. Non-breaking: internal bounds checks and documentation only, no signature or call-path changes. The compactor-side memory safety (drain-on-unwind, truncate-reused-outputs) stays with the compaction work, where it lives. --- .../jbellis/jvector/disk/ReaderSupplier.java | 17 +++++++ .../jvector/disk/SimpleMappedReader.java | 9 ++++ .../jvector/graph/disk/OnDiskGraphIndex.java | 45 +++++++++++++++++-- 3 files changed, 68 insertions(+), 3 deletions(-) diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java index 7ebb3f9b0..8827bcdb1 100644 --- a/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java +++ b/jvector-base/src/main/java/io/github/jbellis/jvector/disk/ReaderSupplier.java @@ -41,6 +41,23 @@ public interface ReaderSupplier extends AutoCloseable { default void prefetch(long offset, long length) { } + /** + * Releases the supplier's underlying resource. Two implementation families exist, with very + * different safety under concurrency: + *

+ * 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/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 {