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 (memory_size 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.

estimated_size keeps its data-only meaning — it sizes the flush unit, so a generation stays a function of the rows in it. MemTable::memory_size is the new resident-bytes accessor.

The seam:

async fn maybe_apply_backpressure(&self, incoming_bytes: usize, 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 both the size of the batch about to be admitted and 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.

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. pod_memory_bytes optionally sums them across every shard in the process, so a global gate reads one atomic instead of scanning.

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

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 — 507 pass. 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!