feat(mem_wal): make write backpressure injectable and account index memory

[hamersaw]

Why

MemWAL's backpressure is a concrete struct whose only seam is a per-call closure, and it can see one shard's row bytes. That leaves two gaps:

1. An embedder cannot express its own budget. A process running many shards has limits lance cannot know — a process-wide memtable total, a page-cache working set. Nothing could represent them, and nothing could refuse a write: the valve only ever blocked, unboundedly.
2. The bytes it measures are not the bytes that OOM you. MemTable::estimated_size counts buffered batches + the PK bloom filter. Every in-memory index is invisible to it.

Gap 2 is not a rounding error. HnswGraph::try_new pre-allocates one node per unit of capacity (= max_memtable_rows), so a vector memtable commits its entire graph on the first insert:

max_memtable_rows	HNSW on row #1
125k	64.7 MiB
500k	258 MiB
1M	517 MiB
2M	1033 MiB

Measured, not estimated, and identical at dim=768 and dim=8 — vectors are held by reference, so the cost is ~542 B per row of capacity regardless of dimension. A memtable holding one row costs the same as a full one, while estimated_size reads ~zero.

What

Accounting (resident_bytes at each layer, all cheap enough for the write path):

• HNSW — the node arena and lookup slabs are sized from capacity at construction, so the total is computed once in the loop try_new already runs, plus an atomic for the rebuilt packed_level0.
• BTree — the skiplist arena counts chunks in its cold grow path: free per insert, exact for the nodes.
• FTS — partitions are capped at MAX_PARTITIONS and size themselves in O(1); only the mutable tail needed a running counter, kept in step with the existing walk at its single growth point. A test asserts the two agree exactly, including across a freeze.

Row bytes keep their data-only meaning — they size the flush unit, so a generation stays a function of the rows in it. They move off MemTableStats onto ShardMemory::row_bytes(), alongside index_bytes(), frozen_bytes(), grace_bytes() and retained_bytes(); InMemoryMemTableRef::resident_bytes is the per-memtable total the ceiling is built on.

The seam:

async fn maybe_apply_backpressure(&self, shard: ShardMemory) -> Result<()>

ShardWriterConfig::backpressure replaces the built-in valve rather than layering on it, so one implementation owns the whole policy. That is why the gate receives a live view of what the calling shard holds: a replacement can enforce a per-shard ceiling in the same place as its process-wide one, without a back-reference into the writer calling it. ShardMemory is deliberately live rather than a snapshot — a controller that delays is waiting for those bytes to fall, so a captured copy would never observe the drain.

The gate is deliberately not told the incoming batch's size. Batches are decoded and resident by the time it runs, so there is nothing left to reserve against — refusing does not un-allocate them, and bounding a single write's memory belongs to the ingress. This is the settled shape rather than a first cut: adding a parameter later would break every external implementor.

With nothing injected, LocalBackpressureController keeps today's per-shard behaviour, its two write modes covered by a LocalSource enum instead of the closure.

Counters. unflushed_memtable_bytes read through a try_read() that yields 0 whenever the write lock is held — and tokio's RwLock is write-preferring, so also whenever a writer is merely queued. It zeroed exactly the shards taking writes, reporting lowest under the heaviest load. Replaced with active_bytes() / frozen_bytes(): relaxed loads of counters maintained under the write lock, index memory included.

Error::Backpressure (with is_backpressure()) makes a rejection distinguishable from a real failure without matching on the message.

Reservations are validated at open(). An HNSW graph is charged from max_memtable_rows before its first insert, while only row bytes seal a memtable. A reservation with no room left under max_unflushed_memtable_bytes would therefore put a shard over budget at zero rows — nothing to seal, so nothing to flush, so every put stalls and then fails as Backpressure, which is supposed to mean "retry later" and never comes true. open() now rejects that configuration outright, requiring index_reserved + max_memtable_size <= max_unflushed_memtable_bytes, with both figures named in the error. Only the built-in valve reads that ceiling, so the check is skipped when a controller is injected.

Update: the process-wide counter is gone

Earlier revisions carried ShardWriterConfig::pod_memory_bytes, an Arc<AtomicUsize> every writer added to on insert and subtracted from on flush-commit, so an embedder could read a process-wide total with one load instead of scanning every shard.

It has been removed, because an incremental counter of that shape has no way back. Nothing releases a writer's residual bytes when it goes away without a final flush — abort() is shutdown_all() and nothing else, and there is no Drop on MemoryCounters — so every teardown that skips a flush leaks its memtable into the sum permanently. The downstream embedder evicts that way on four separate paths, and enough of them leave a process refusing every write against a memtable that reads empty.

A Drop impl would close that particular hole. But the embedder is better served recomputing the total from memtable_stats(), which is also what it exports: derived, self-healing within one interval, and identical to its own gauge by construction. So the counter is deleted rather than patched, and the seam is smaller for it — net −41 lines.

ShardMemory and the per-shard active/frozen counters are untouched. They are the live per-shard view a controller needs to enforce a per-shard ceiling, and they have no cross-shard lifetime problem.

Contract changes

• "Never errors due to backpressure" → an injected controller may reject. The built-in valve still never does.
• "Never drops data" → never drops acked data. A rejected write was never accepted.
• active/frozen now count index memory, so the built-in valve is meaningfully tighter on vector tables at a given max_unflushed_memtable_bytes. That is the correct direction for a memory valve, but it is a behaviour change worth knowing about.

Follow-up (not this PR)

The graph costs ~5x the adjacency it stores: level-0 links are held three times (ranked with distances, published, packed_level0), and it is ~2-3 allocations per node. The reference implementation for the fix already lives next door in mem_wal/index/arena_skiplist.rs, whose BTree gets 4 MiB where HNSW spends 64.7 MiB. Filing separately — this PR only makes the cost visible.

Test

cargo test -p lance --lib -- dataset::mem_wal. New: HNSW pre-allocation behaviour, FTS counter-vs-walk equality, injected-replaces-default, default-when-none-injected, reject-surfaces-as-Backpressure, and ShardMemory liveness across polls (it hangs if the view ever regresses to a copy).

cargo fmt --all + cargo clippy -p lance --lib --tests --benches clean in touched files.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 7678afb9-dcbe-4cf8-8cb0-c80593802c01

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 86.68407% with 51 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance/src/dataset/mem_wal/write.rs	80.16%	44 Missing and 3 partials ⚠️
rust/lance/src/dataset/mem_wal/index/btree.rs	86.66%	2 Missing ⚠️
rust/lance-core/src/error.rs	90.00%	0 Missing and 1 partial ⚠️
rust/lance/src/dataset/mem_wal/index.rs	88.88%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[wjones127]

issue (blocking): The ceiling now counts index bytes but the seal trigger still counts row bytes (write.rs:1489, memtable.rs:511), so a shard configured with a large-capacity HNSW is over max_unflushed_memtable_bytes at open() with zero rows, can never seal, and — via this arm — refuses every put after STALL_GRACE; index.rs:1919 confirms the charge lands from configuration. Acceptance: either the seal trigger consults resident bytes so a stalled shard can actually drain, or open() validates that the configured indexes' reserved bytes leave headroom under the ceiling and fails fast there — plus a test that a writer with an HNSW index sized well above max_unflushed_memtable_bytes still accepts writes (or is rejected at open, not at put #1).

[hamersaw]

Updated this so that open() validates headroom. This feels like the right call to make right now, but the failing if there is no room is probably not the right long-term call. We should then proactively create room by flushing + compacting table data.

[wjones127]

question (blocking): The PR body advertises maybe_apply_backpressure(&self, incoming_bytes: usize, shard: ShardMemory) but the trait ships without incoming_bytes — is the no-size shape final? Adding a parameter later breaks every external implementor, so I'd like this settled now rather than discovered by an embedder; the "batches are already resident" rationale is reasonable, I just want it to be the decided one.

[hamersaw]

Updated the PR body to be correct with the code.

[wjones127]

issue (non-blocking): row_bytes is not active_bytes() - index_bytes(): active_bytes is built on retained_row_bytes (full pinned parent capacity) while this is the window sum, and the gap is unbounded for zero-copy slices — your own test_memory_handle_tracks_the_live_memtable asserts resident == retained_row_bytes + index_bytes, not row_bytes + index_bytes.

[hamersaw]

I think this can be solved with just a doc update, the functionality seems correct to me.

[wjones127]

question (non-blocking): The "same parent charged once" claim holds because ArrayData::slice advances offset rather than the buffer pointer — but a batch whose buffers came from a kernel that re-slices (Buffer::slice_with_length, concat/take output) gets a distinct data_ptr for the same allocation and is charged again in full; worth softening the doc to "once per distinct buffer view" since it over-counts rather than under-counts. Separately: retained_buffers is never pruned and to_data() + a mutex + a hash insert per column now run on every append — fine at current batch rates, but worth a note that it grows with accepted batches.

[hamersaw]

Both fair. Softened to "once per distinct buffer view", spelling out that ArrayData::slice advances the offset and leaves the pointer alone so ordinary slices dedupe, while a buffer that came back re-sliced from a kernel presents a different data_ptr for the same allocation and is charged again in full — over-counting, which is the safe direction for a ceiling.

Also documented the growth: retained_buffers is never pruned, bounded only by the store being dropped at flush, and the walk plus to_data, the mutex and a hash insert run per column per append — fine at current batch rates, and the first thing to look at if that changes.

[wjones127]

issue (non-blocking): null_bytes takes a blocking lock() on null_positions, and this path is now reached from ShardMemory::active_bytes() on every put and every 10ms DRAIN_POLL_INTERVAL tick — so a memory poll can block behind an in-progress insert holding the same mutex; a try_lock with a fallback (it's a size estimate) would keep the read path non-blocking.

[hamersaw]

Updated this to use an atomic instead to mirror the existing approaches.

[wjones127]

That sounds like the right call to me.

[wjones127]

question (non-blocking): Python gained index_bytes, grace_bytes, and retained_bytes in this PR while Java still exposes only the one repurposed row_bytes slot — is the JNI surface intentionally left behind, or should it get the same breakdown so cross-language embedders can build the same budget?

[hamersaw]

These should all be exposed through the Java SDK as well now.

[lance-gatekeeper]

✅ Gate recommendation: approve.

3 fixed / 0 remain. The built-in controller now rejects zero row headroom before claiming the epoch, and the pre-valve seal recheck lets both post-index growth and replay tails rotate before a subsequent write enters backpressure. The solution preserves row-window fragment sizing while closing every verified permanent no-drain path.

Please mark this PR with the breaking-change label.

[wjones127]

Looks good now. Thanks for addressing my comments.