From dd019c416dc65f88534743499f438c83759c60cc Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 16 Jul 2026 09:26:19 -0500 Subject: [PATCH 1/2] feat(mem_wal): account in-memory index bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MemTable::estimated_size` counts row data and the PK bloom filter only, so every in-memory index is invisible to callers sizing memtable memory. That gap is not a rounding error: a configured HNSW index pre-allocates its whole graph from `capacity` on the first insert, so at the WAL's defaults (125k rows, m=16) a vector memtable costs 64.7 MiB the moment row #1 lands — while `estimated_size` still reads near zero. Many small vector tables therefore sail past any budget built on `estimated_size` and OOM the process. Add `memory_size` at each layer, 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 `HnswGraph::try_new` (which already walks the nodes) plus an atomic for the rebuilt `packed_level0`. Vectors are held by reference, so this is independent of `dim`. - BTree: the skiplist arena counts its chunks in the cold `grow` path — free per insert, and 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 (`append_batch`). `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. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance/src/dataset/mem_wal/hnsw/graph.rs | 50 +++++++++++++ .../lance/src/dataset/mem_wal/hnsw/storage.rs | 11 +++ rust/lance/src/dataset/mem_wal/index.rs | 20 +++++ .../dataset/mem_wal/index/arena_skiplist.rs | 25 ++++++- rust/lance/src/dataset/mem_wal/index/btree.rs | 22 ++++++ rust/lance/src/dataset/mem_wal/index/fts.rs | 73 ++++++++++++++++++- rust/lance/src/dataset/mem_wal/index/hnsw.rs | 58 +++++++++++++++ rust/lance/src/dataset/mem_wal/memtable.rs | 16 +++- 8 files changed, 269 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/hnsw/graph.rs b/rust/lance/src/dataset/mem_wal/hnsw/graph.rs index 4009557702c..f82cea9468c 100644 --- a/rust/lance/src/dataset/mem_wal/hnsw/graph.rs +++ b/rust/lance/src/dataset/mem_wal/hnsw/graph.rs @@ -231,6 +231,17 @@ impl LevelLinks { } } + /// Heap bytes for one level, counting `published` at its full width even + /// while empty — the build fills it, and the caller budgets against a + /// ceiling where over-counting is the safe direction. + fn allocated_bytes(max_neighbors: usize) -> usize { + // Arc>: strong + weak refcounts, the Vec header, then the ids. + let published = 2 * std::mem::size_of::() + + std::mem::size_of::>() + + max_neighbors * std::mem::size_of::(); + published + max_neighbors * std::mem::size_of::() + } + fn publish_from_ranked(&self, ranked: &[ScoredPoint]) { self.published.store(Arc::new( ranked.iter().map(|point| point.id).collect::>(), @@ -259,6 +270,16 @@ impl Node { } } + /// Heap bytes held by a node of `target_level`, excluding the `Node` itself + /// (which lives inline in the graph's node arena). + fn allocated_bytes(target_level: u16, m: usize) -> usize { + let levels = target_level as usize + 1; + levels * std::mem::size_of::() + + (0..=target_level) + .map(|level| LevelLinks::allocated_bytes(max_neighbors(m, level))) + .sum::() + } + fn has_level(&self, level: u16) -> bool { (level as usize) < self.levels.len() } @@ -292,6 +313,12 @@ pub struct HnswGraph { visible_len: AtomicUsize, visited_pool: ArrayQueue, packed_level0: ArcSwap, + /// Heap bytes of the node arena and visited pool. Fixed at construction: + /// both are sized from `capacity`, not from `len()`. + base_bytes: usize, + /// Heap bytes of the current `packed_level0` snapshot, which is rebuilt + /// wholesale on each level-0 publish rather than grown. + packed_bytes: AtomicUsize, } impl HnswGraph { @@ -309,12 +336,14 @@ impl HnswGraph { let mut rng = SmallRng::seed_from_u64(params.seed); let mut nodes = Vec::with_capacity(capacity); + let mut node_bytes = 0; for id in 0..capacity { let target_level = if id == 0 { 0 } else { random_level(¶ms, &mut rng) }; + node_bytes += Node::allocated_bytes(target_level, params.m); nodes.push(Node::new(target_level, params.m)); } @@ -323,9 +352,15 @@ impl HnswGraph { for _ in 0..pool_size { let _ = visited_pool.push(VisitedList::new(0)); } + // Pooled lists start empty but `VisitedList::reset` resizes each to one + // bit per node on first use. + let visited_bytes = pool_size + * (std::mem::size_of::() + + capacity.div_ceil(WORD_BITS) * std::mem::size_of::()); Ok(Self { params, + base_bytes: capacity * std::mem::size_of::() + node_bytes + visited_bytes, nodes, build_entry_point: AtomicU32::new(0), build_max_level: AtomicU16::new(0), @@ -335,9 +370,20 @@ impl HnswGraph { visible_len: AtomicUsize::new(0), visited_pool, packed_level0: ArcSwap::from_pointee(PackedLevel::empty()), + packed_bytes: AtomicUsize::new(0), }) } + /// Upper bound on the graph's dominant heap allocations. + /// + /// Near-constant from the first insert rather than proportional to `len()`: + /// the node arena is allocated in full at construction, sized by `capacity`. + /// Callers budgeting memtable memory must account for this the moment a + /// vector memtable takes its first row. + pub fn memory_size(&self) -> usize { + self.base_bytes + self.packed_bytes.load(Ordering::Relaxed) + } + /// Number of nodes visible to readers. pub fn len(&self) -> usize { self.visible_len.load(Ordering::Acquire) @@ -1051,9 +1097,13 @@ impl HnswGraph { offsets.push(neighbors.len()); } + let packed_bytes = offsets.capacity() * std::mem::size_of::() + + neighbors.capacity() * std::mem::size_of::(); + // ArcSwap reclaims the prior snapshot once no reader guard holds it. self.packed_level0 .store(Arc::new(PackedLevel { offsets, neighbors })); + self.packed_bytes.store(packed_bytes, Ordering::Relaxed); Ok(()) } } diff --git a/rust/lance/src/dataset/mem_wal/hnsw/storage.rs b/rust/lance/src/dataset/mem_wal/hnsw/storage.rs index bbeb57a5fe2..93b1947fa27 100644 --- a/rust/lance/src/dataset/mem_wal/hnsw/storage.rs +++ b/rust/lance/src/dataset/mem_wal/hnsw/storage.rs @@ -196,6 +196,17 @@ impl ArrowFixedSizeListVectorStore { }) } + /// Heap bytes of the store's own slabs, all sized from `capacity` and + /// `max_batches` at construction. + /// + /// Excludes the vectors themselves: batches are held by reference, so their + /// bytes belong to the MemTable's batch store and counting them here would + /// double-count. That also makes this independent of `dim`. + pub fn memory_size(&self) -> usize { + self.max_batches * std::mem::size_of::() + + self.capacity * (std::mem::size_of::() + std::mem::size_of::()) + } + /// Number of committed vectors. pub fn committed_len(&self) -> usize { self.committed_len.load(Ordering::Acquire) diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 8077f06149e..f6f0d2666b8 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -985,6 +985,26 @@ impl IndexStore { self.btree_indexes.len() + self.hnsw_indexes.len() + self.fts_indexes.len() } + /// Heap bytes held by every index in the registry. + /// + /// `MemTable::estimated_size` deliberately omits this — it sizes the flush + /// unit, which is row data. Callers budgeting *resident* memory must add it: + /// a configured HNSW index pre-allocates its whole graph on the first insert + /// (see [`HnswMemIndex::memory_size`]), so it can dwarf a memtable's row + /// bytes while `estimated_size` still reads near zero. + pub fn memory_size(&self) -> usize { + let btrees: usize = self.btree_indexes.values().map(|b| b.memory_size()).sum(); + let hnsw: usize = self.hnsw_indexes.values().map(|h| h.memory_size()).sum(); + let fts: usize = self.fts_indexes.values().map(|f| f.memory_size()).sum(); + // A `Single` PK aliases a `btree_indexes` entry, already counted above. + // A composite PK's index is held only here. + let pk = match &self.pk_index { + Some(PkIndex::Composite { index, .. }) => index.memory_size(), + Some(PkIndex::Single(_)) | None => 0, + }; + btrees + hnsw + fts + pk + } + /// Get the visibility watermark (max batch position safe to read). /// /// Returns the highest batch position whose data is durable in the WAL diff --git a/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs b/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs index 6b7361e9f1b..0120e981b13 100644 --- a/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs +++ b/rust/lance/src/dataset/mem_wal/index/arena_skiplist.rs @@ -103,11 +103,12 @@ impl Arena { } /// Bump-allocate `layout`. Caller must have exclusive access (single writer). - unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 { + /// `allocated` accumulates chunk bytes; only the cold `grow` path touches it. + unsafe fn alloc(&mut self, layout: Layout, allocated: &AtomicUsize) -> *mut u8 { let align = layout.align(); let mut aligned = (self.cursor as usize).wrapping_add(align - 1) & !(align - 1); if self.cursor.is_null() || aligned + layout.size() > self.end as usize { - self.grow(layout); + self.grow(layout, allocated); aligned = (self.cursor as usize + align - 1) & !(align - 1); } self.cursor = (aligned + layout.size()) as *mut u8; @@ -116,7 +117,7 @@ impl Arena { /// Allocate a fresh chunk large enough for `layout` and make it current. #[cold] - unsafe fn grow(&mut self, layout: Layout) { + unsafe fn grow(&mut self, layout: Layout, allocated: &AtomicUsize) { let align = layout.align().max(64); let size = CHUNK_SIZE.max(layout.size().next_power_of_two()); let chunk_layout = Layout::from_size_align(size, align).expect("valid chunk layout"); @@ -126,6 +127,7 @@ impl Arena { } self.chunks .push((NonNull::new_unchecked(ptr), chunk_layout)); + allocated.fetch_add(size, Ordering::Relaxed); self.cursor = ptr; self.end = ptr.add(size); } @@ -154,6 +156,10 @@ struct SkipListCore { height: AtomicUsize, /// Number of entries. len: AtomicUsize, + /// Bytes of arena chunks allocated so far. Maintained here rather than read + /// off `arena.chunks` because the arena is writer-only; this is readable by + /// anyone. Only `Arena::grow` touches it, so it costs nothing per insert. + arena_bytes: AtomicUsize, } // SAFETY: `arena` (the only non-Sync field) is mutated exclusively by the single @@ -174,6 +180,7 @@ impl SkipListCore { arena: UnsafeCell::new(Arena::new()), height: AtomicUsize::new(1), len: AtomicUsize::new(0), + arena_bytes: AtomicUsize::new(0), } } @@ -303,7 +310,8 @@ impl SkipListWriter { // only mutator, so no link changes between read and publish. let layout = node_layout::(height); // SAFETY: single-writer exclusive access to the arena. - let node = unsafe { (*self.core.arena.get()).alloc(layout) } as *mut Node; + let node = unsafe { (*self.core.arena.get()).alloc(layout, &self.core.arena_bytes) } + as *mut Node; // SAFETY: `node` points to a fresh, uninitialized, correctly-sized and // -aligned block; we write the key then `height` tower slots. unsafe { @@ -342,6 +350,15 @@ pub struct SkipListReader { } impl SkipListReader { + /// Bytes of arena chunks backing this skiplist's nodes. + /// + /// Counts chunks, not entries, so it steps by `CHUNK_SIZE` and overshoots + /// the live nodes by at most one partly-filled chunk. Excludes any bytes a + /// key owns outside its node (e.g. a long `Box<[u8]>` key). + pub fn memory_size(&self) -> usize { + self.core.arena_bytes.load(Ordering::Relaxed) + } + /// Greatest node with `key <= target`, mapped through `f` while it is alive. /// Equivalent to crossbeam's `upper_bound(Included(target))`. `None` if no /// such node. The closure avoids cloning the key on the hot path. diff --git a/rust/lance/src/dataset/mem_wal/index/btree.rs b/rust/lance/src/dataset/mem_wal/index/btree.rs index ca4dd178548..fa25e0c07cb 100644 --- a/rust/lance/src/dataset/mem_wal/index/btree.rs +++ b/rust/lance/src/dataset/mem_wal/index/btree.rs @@ -833,6 +833,20 @@ impl Backend { } } + fn memory_size(&self) -> usize { + fn null_bytes(nulls: &Mutex>) -> usize { + nulls + .lock() + .map(|n| n.capacity() * std::mem::size_of::()) + .unwrap_or(0) + } + match self { + Self::FixedInt(b) => b.reader.memory_size() + null_bytes(&b.null_positions), + Self::Bytes(b) => b.reader.memory_size() + null_bytes(&b.null_positions), + Self::Scalar(b) => b.reader.memory_size(), + } + } + fn data_type(&self) -> Option { match self { Self::FixedInt(b) => Some(b.data_type()), @@ -937,6 +951,14 @@ impl BTreeMemIndex { self.backend.get().map(|b| b.len()).unwrap_or(0) } + /// Heap bytes held by this index; zero before the first insert. + /// + /// Grows with rows (unlike the pre-allocated HNSW index). Arena-chunk + /// granular, so it steps rather than climbs smoothly. + pub fn memory_size(&self) -> usize { + self.backend.get().map(|b| b.memory_size()).unwrap_or(0) + } + /// Check if the index is empty. pub fn is_empty(&self) -> bool { self.len() == 0 diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 6e0d12a276e..46b982396ca 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -561,6 +561,10 @@ impl BatchMeta { } } +/// Per-entry overhead charged for a `SkipMap` node (tower + links) when sizing +/// the tail's term map. An estimate — the node layout is crossbeam-internal. +const SKIPMAP_ENTRY_OVERHEAD: usize = 32; + /// Size of a sealed batch block. Small enough that copying the partial tail /// block on append stays cheap; large enough that sealing (which clones the /// block-pointer vec) is rare. @@ -770,6 +774,11 @@ struct TailIndex { /// hash probe instead of a skiplist search. Reset implicitly when the tail /// is replaced on freeze. Uncontended — the single writer holds it briefly. writer_term_cache: Mutex, Arc>>>, + /// Running total mirroring [`Self::memory_size`], maintained by + /// `append_batch`. Exists so the write path can budget memtable memory + /// without the O(terms) walk. `test_tail_bytes_tracks_memory_size` pins the + /// two together. + bytes: AtomicUsize, } impl TailIndex { @@ -779,9 +788,15 @@ impl TailIndex { snapshot: ArcSwap::from(Snapshot::empty()), next_batch_position: AtomicUsize::new(0), writer_term_cache: Mutex::new(FxHashMap::default()), + bytes: AtomicUsize::new(std::mem::size_of::()), }) } + /// [`Self::memory_size`] without the walk. See [`Self::bytes`]. + fn memory_size_cached(&self) -> usize { + self.bytes.load(Ordering::Relaxed) + } + fn snapshot(&self) -> Arc { self.snapshot.load_full() } @@ -815,8 +830,19 @@ impl TailIndex { .writer_term_cache .lock() .expect("writer term cache poisoned — single-writer invariant violated"); + // Mirrors `memory_size`'s per-term arithmetic; keep the two in step. + let mut added = 0; for (term, builder) in term_builders { let chunk = builder.build(batch_position, with_position); + added += std::mem::size_of::() + chunk.memory_size(); + if !cache.contains_key(&term) { + // First sight this generation: the SkipMap entry plus the + // slice's empty root node. + added += std::mem::size_of::>() + + term.len() + + SKIPMAP_ENTRY_OVERHEAD + + std::mem::size_of::(); + } // First sight of the term this generation populates the SkipMap // (so readers can find it) and caches the slot; later batches hit // only the cache. @@ -836,6 +862,8 @@ impl TailIndex { doc_lengths, rows, }); + added += new_meta.memory_size(); + self.bytes.fetch_add(added, Ordering::Relaxed); let cur = self.snapshot.load(); self.snapshot.store(Arc::new(Snapshot { visible_count: cur.visible_count + 1, @@ -849,7 +877,7 @@ impl TailIndex { let mut total = std::mem::size_of::(); for entry in self.terms.iter() { let term: &Arc = entry.key(); - total += std::mem::size_of::>() + term.len() + 32; + total += std::mem::size_of::>() + term.len() + SKIPMAP_ENTRY_OVERHEAD; total += entry.value().load().memory_size(); } total += self @@ -1040,6 +1068,8 @@ impl FtsMemIndex { } /// Estimated bytes of heap memory held by this index. + /// + /// Walks every tail term. Prefer [`Self::memory_size`] on the write path. pub fn memory_usage(&self) -> usize { let st = self.state.load_full(); let mut total = std::mem::size_of::(); @@ -1048,6 +1078,16 @@ impl FtsMemIndex { total } + /// [`Self::memory_usage`] without the per-term walk: partitions are capped + /// at `MAX_PARTITIONS` and size themselves in O(1), and the tail keeps a + /// running total. Cheap enough for the write path. + pub fn memory_size(&self) -> usize { + let st = self.state.load(); + std::mem::size_of::() + + st.partitions.iter().map(|p| p.memory_size()).sum::() + + st.tail.memory_size_cached() + } + /// Component memory breakdown (bytes), for diagnostics: /// `(num_partitions, term_strings, postings_meta, block_meta, doc_freq, pos, docs, tail)`. pub fn memory_breakdown(&self) -> (usize, usize, usize, usize, usize, usize, usize, usize) { @@ -4587,6 +4627,37 @@ mod tests { ); } + /// The tail's running byte counter must stay exactly in step with the walk + /// it replaces, including across a freeze (which swaps in a fresh tail). + #[test] + fn test_tail_bytes_tracks_memory_size() { + let schema = create_test_schema(); + // Freeze partway through so the counter is checked on both a live tail + // and a post-freeze one. + let index = FtsMemIndex::new(1, "description".to_string()).with_freeze_threshold_rows(4); + + for round in 0..6 { + let batch = if round % 2 == 0 { + create_test_batch(&schema) + } else { + create_phrase_test_batch(&schema) + }; + index.insert(&batch, round * 100).unwrap(); + + let st = index.state.load(); + assert_eq!( + st.tail.memory_size_cached(), + st.tail.memory_size(), + "tail byte counter drifted from the walk at round {round}" + ); + assert_eq!( + index.memory_size(), + index.memory_usage(), + "index memory_size drifted from memory_usage at round {round}" + ); + } + } + #[test] fn test_partial_doc_never_visible_phrase() { // A phrase query inside a single document must either match fully diff --git a/rust/lance/src/dataset/mem_wal/index/hnsw.rs b/rust/lance/src/dataset/mem_wal/index/hnsw.rs index 7c7993bb298..0725d38e6df 100644 --- a/rust/lance/src/dataset/mem_wal/index/hnsw.rs +++ b/rust/lance/src/dataset/mem_wal/index/hnsw.rs @@ -156,6 +156,19 @@ impl HnswMemIndex { self.len() == 0 } + /// Upper bound on heap bytes held by this index. + /// + /// Zero until the first insert, then a near-constant sized by `capacity` + /// (the writer's `max_memtable_rows`) rather than by rows inserted — the + /// graph and lookup slabs are pre-allocated in full by [`Self::ensure_state`]. + /// So an idle-but-touched vector memtable costs the same as a full one. + pub fn memory_size(&self) -> usize { + self.state + .get() + .map(|s| s.graph.memory_size() + s.storage.memory_size()) + .unwrap_or(0) + } + fn ensure_state(&self, dim: usize) -> Result<&HnswState> { if let Some(state) = self.state.get() { if state.storage.dim() != dim { @@ -462,6 +475,51 @@ mod tests { RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids)), Arc::new(fsl)]).unwrap() } + /// The graph is pre-allocated from `capacity`, so its footprint appears in + /// full on the first insert and barely moves as rows arrive. A memory + /// budget that samples only row bytes would miss all of it. + #[test] + fn test_memory_size_is_preallocated_not_proportional_to_rows() { + let dim = 8; + let capacity = 4_000; + let index = || { + HnswMemIndex::with_capacity( + 1, + "vector".to_string(), + DistanceType::L2, + HnswBuildParams::default().num_edges(16).ef_construction(64), + capacity, + 64, + ) + }; + + let untouched = index(); + assert_eq!( + untouched.memory_size(), + 0, + "an index with no state built holds nothing" + ); + + let sparse = index(); + sparse.insert(&make_batch(0, 1, dim), 0).unwrap(); + let one_row = sparse.memory_size(); + + let full = index(); + full.insert(&make_batch(0, capacity, dim), 0).unwrap(); + let all_rows = full.memory_size(); + + // One row already pays for the whole graph: well over a KB per slot of + // capacity, and within a small factor of the fully-populated index. + assert!( + one_row > capacity * 128, + "one row should commit the pre-allocated graph, got {one_row} for capacity {capacity}" + ); + assert!( + all_rows < one_row * 2, + "a full index ({all_rows}) should not dwarf a one-row index ({one_row})" + ); + } + #[test] fn test_index_insert_and_search() { let dim = 8; diff --git a/rust/lance/src/dataset/mem_wal/memtable.rs b/rust/lance/src/dataset/mem_wal/memtable.rs index 77611fd7c47..c7d35c0d4f8 100644 --- a/rust/lance/src/dataset/mem_wal/memtable.rs +++ b/rust/lance/src/dataset/mem_wal/memtable.rs @@ -721,11 +721,25 @@ impl MemTable { self.batch_count() } - /// Get estimated size in bytes. + /// Row-data bytes: the buffered batches plus the PK bloom filter. + /// + /// This is the **flush unit**, not the memtable's footprint — it drives the + /// `max_memtable_size` freeze trigger and deliberately excludes index + /// memory, so a generation's size stays a function of the rows in it. Use + /// [`Self::memory_size`] to budget resident memory. pub fn estimated_size(&self) -> usize { self.batch_store.estimated_bytes() + self.pk_bloom_filter.estimated_memory_size() } + /// Total resident heap bytes: row data plus every in-memory index. + /// + /// Can far exceed [`Self::estimated_size`] on an indexed table — an HNSW + /// index pre-allocates its full graph on the first insert — so this, not + /// `estimated_size`, is what a memory ceiling must be built on. + pub fn memory_size(&self) -> usize { + self.estimated_size() + self.indexes().map_or(0, IndexStore::memory_size) + } + /// Get the WAL batch mapping. pub fn wal_batch_mapping(&self) -> &HashMap { &self.wal_batch_mapping From a6f19a5c8a2c128923226679f20afe4f2ce2803c Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 16 Jul 2026 10:21:39 -0500 Subject: [PATCH 2/2] feat(mem_wal): make write backpressure injectable and bound resident bytes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backpressure was a concrete struct whose only seam was a per-call closure, and it could see one shard's row bytes. An embedder running many shards in one process has budgets lance cannot: a process-wide memtable total, a page-cache working set. Nothing could express them, and nothing could refuse a write — the valve only ever blocked, unboundedly. Turn it into a trait an embedder can implement: 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 is handed both the size of the batch about to be admitted and what the calling shard already holds: a replacement can enforce a per-shard ceiling in the same place as its process-wide one, with no back-reference into the writer. With nothing injected, `LocalBackpressureController` keeps today's per-shard behaviour, its two write modes covered by a `LocalSource` enum rather than the closure. Fix what the counters measure. `unflushed_memtable_bytes` read through a `try_read()` that yielded 0 whenever the write lock was held — and tokio's `RwLock` is write-preferring, so it also yielded 0 whenever a writer was merely queued. It zeroed exactly the shards taking writes, reporting lowest under the heaviest load. Replace it with `active_bytes()`/`frozen_bytes()`: relaxed loads of counters maintained under the write lock, counting index memory too (a vector memtable's HNSW graph is ~65 MiB that `estimated_size` cannot see). `ShardWriterConfig::pod_memory_bytes` optionally sums them across every shard in the process, so a global gate reads one atomic instead of scanning shards. `Error::Backpressure` (with `is_backpressure()`) lets a rejection be told from a real failure without matching on the message. This relaxes the contract: "never errors due to backpressure" becomes "an injected controller may reject", and "never drops data" becomes "never drops *acked* data" — a rejected write was never accepted. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/lance-core/src/error.rs | 33 +- .../benches/mem_wal/write/mem_wal_write.rs | 4 + rust/lance/src/dataset/mem_wal/write.rs | 662 +++++++++++++++--- 3 files changed, 615 insertions(+), 84 deletions(-) diff --git a/rust/lance-core/src/error.rs b/rust/lance-core/src/error.rs index a85f08cb741..99b27a7ef31 100644 --- a/rust/lance-core/src/error.rs +++ b/rust/lance-core/src/error.rs @@ -373,6 +373,19 @@ pub enum Error { #[snafu(implicit)] location: Location, }, + /// A write was refused to keep the writer inside its memory budget. + /// + /// Unlike every other write error this one is *expected* under load and + /// carries no data loss: the write was never accepted, so a caller that + /// retries once the flush pipeline drains loses nothing. Callers should + /// surface it as a retryable "busy" signal (HTTP 503), not a failure. + /// Match via [`Error::is_backpressure`] rather than on the message. + #[snafu(display("Write rejected by backpressure: {message}, {location}"))] + Backpressure { + message: String, + #[snafu(implicit)] + location: Location, + }, } impl Error { @@ -423,7 +436,8 @@ impl Error { | Self::FieldNotFound { .. } | Self::Timeout { .. } | Self::DiskCapExceeded { .. } - | Self::Fenced { .. } => None, + | Self::Fenced { .. } + | Self::Backpressure { .. } => None, } } @@ -485,6 +499,23 @@ impl Error { } } + /// A write was refused because the writer is at its memory ceiling; the + /// data was never accepted. See [`Error::Backpressure`]. + #[track_caller] + pub fn backpressure(message: impl Into) -> Self { + BackpressureSnafu { + message: message.into(), + } + .build() + } + + /// Whether this is [`Error::Backpressure`] — i.e. a retryable "writer is + /// full" signal rather than a real failure. Prefer this over matching the + /// error message. + pub fn is_backpressure(&self) -> bool { + matches!(self, Self::Backpressure { .. }) + } + #[track_caller] pub fn io_source(source: BoxedError) -> Self { IOSnafu.into_error(source) diff --git a/rust/lance/benches/mem_wal/write/mem_wal_write.rs b/rust/lance/benches/mem_wal/write/mem_wal_write.rs index db28c5b33ab..733834dac5c 100644 --- a/rust/lance/benches/mem_wal/write/mem_wal_write.rs +++ b/rust/lance/benches/mem_wal/write/mem_wal_write.rs @@ -658,6 +658,10 @@ fn bench_lance_memwal_write(c: &mut Criterion) { warmer: None, store_params: default_config.store_params, session: default_config.session, + // Measure the built-in per-shard valve, not + // an injected policy. + backpressure: None, + pod_memory_bytes: None, }; // Get writer through Dataset API (index configs loaded automatically) diff --git a/rust/lance/src/dataset/mem_wal/write.rs b/rust/lance/src/dataset/mem_wal/write.rs index 4c1cccd5af8..190525161c1 100644 --- a/rust/lance/src/dataset/mem_wal/write.rs +++ b/rust/lance/src/dataset/mem_wal/write.rs @@ -12,9 +12,10 @@ //! - [`IndexStore`] - In-memory index management //! - [`MemTableFlusher`] - Flush MemTable to storage as single Lance file +use std::cmp::Ordering as CmpOrdering; use std::collections::{HashMap, VecDeque}; use std::fmt::Debug; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, RwLock as StdRwLock}; use std::time::{Duration, Instant}; @@ -261,6 +262,30 @@ pub struct ShardWriterConfig { /// Session for those opens, injected alongside `store_params`. /// Default: `None`. pub session: Option>, + + /// Shared counter this writer adds its resident memtable bytes to, so an + /// embedder running many shards in one process can read a process-wide + /// total with a single load instead of scanning every shard. + /// + /// Lance treats it as opaque: it adds on insert and subtracts on + /// flush-commit, and never reads or interprets it. Give every writer in the + /// process the same `Arc` to make the sum meaningful. Because the increment + /// lands under the write lock before the write is acked, a reader sees a + /// safe upper bound rather than a lagging one. + /// Default: `None` (no process-wide accounting). + pub pod_memory_bytes: Option>, + + /// Admission control for every `put`, **replacing** lance's built-in + /// per-shard valve ([`LocalBackpressureController`]). + /// + /// For embedders whose budgets lance cannot see: a process-wide memtable + /// total across shards, a page-cache working set. Because it replaces + /// rather than layers, the injected controller owns the whole policy — a + /// per-shard ceiling included, for which it is handed [`ShardMemory`]. + /// Unlike the built-in valve it may also reject: see + /// [`Error::Backpressure`]. + /// Default: `None` (use the built-in valve). + pub backpressure: Option>, } impl Default for ShardWriterConfig { @@ -289,6 +314,8 @@ impl Default for ShardWriterConfig { warmer: None, store_params: None, session: None, + pod_memory_bytes: None, + backpressure: None, } } } @@ -376,6 +403,20 @@ impl ShardWriterConfig { self } + /// Set the process-wide resident-bytes counter. See + /// [`Self::pod_memory_bytes`]. + pub fn with_pod_memory_bytes(mut self, counter: Arc) -> Self { + self.pod_memory_bytes = Some(counter); + self + } + + /// Replace the built-in per-shard valve with `controller`. See + /// [`Self::backpressure`]. + pub fn with_backpressure(mut self, controller: Arc) -> Self { + self.backpressure = Some(controller); + self + } + /// Set backpressure log interval. pub fn with_backpressure_log_interval(mut self, interval: Duration) -> Self { self.backpressure_log_interval = interval; @@ -691,19 +732,174 @@ pub struct BackpressureStatsSnapshot { pub total_wait_ms: u64, } -/// Backpressure controller for managing write flow. -pub struct BackpressureController { - /// Configuration. - config: ShardWriterConfig, - /// Stats for monitoring. +/// A **live** view of the calling shard's resident memtable bytes, handed to +/// the controller so it can apply a per-shard rule without a back-reference +/// into the writer that is calling it. +/// +/// Re-read it on every poll rather than reading once: a controller that delays +/// is waiting for exactly these numbers to fall, so a captured copy would never +/// observe the drain and the wait would never end. Each accessor is a relaxed +/// atomic load, so polling is cheap. +#[derive(Clone)] +pub struct ShardMemory(ShardMemorySource); + +#[derive(Clone)] +enum ShardMemorySource { + Memtable(Arc), + /// WAL-only mode has no memtable; the pending queue is the whole pool. + WalOnly(Arc), +} + +impl ShardMemory { + /// Resident bytes of the active memtable — row data plus its in-memory + /// indexes. In WAL-only mode, the pending queue's bytes. + pub fn active(&self) -> usize { + match &self.0 { + ShardMemorySource::Memtable(m) => m.active(), + ShardMemorySource::WalOnly(s) => s.estimated_size(), + } + } + + /// Resident bytes of sealed memtables whose flush has not committed. + /// Always `0` in WAL-only mode. + pub fn frozen(&self) -> usize { + match &self.0 { + ShardMemorySource::Memtable(m) => m.frozen(), + ShardMemorySource::WalOnly(_) => 0, + } + } + + /// Bytes only a flush can reclaim. Sums two independent loads, which may + /// cross a freeze — never derive one term from the other. + pub fn unflushed(&self) -> usize { + self.active() + self.frozen() + } +} + +impl Debug for ShardMemory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ShardMemory") + .field("active", &self.active()) + .field("frozen", &self.frozen()) + .finish() + } +} + +/// Admission control for [`ShardWriter::put`], consulted before each write. +/// +/// The default ([`LocalBackpressureController`]) bounds one shard's memtable +/// pool and only ever delays. An embedder running many shards in one process +/// has budgets lance cannot see — a process-wide memtable total, a page-cache +/// working set — and can **replace** the default via +/// [`ShardWriterConfig::backpressure`]. +/// +/// A replacement owns the whole policy, per-shard rules included: nothing else +/// gates the write. That is what [`ShardMemory`] is for — the controller is +/// told what this shard holds, so it can enforce a per-shard ceiling in the +/// same place it enforces its process-wide one, rather than relying on a second +/// valve underneath it. +#[async_trait::async_trait] +pub trait BackpressureController: Send + Sync + Debug { + /// Decide whether to admit a write of `incoming_bytes` into a shard + /// currently holding `shard`. + /// + /// May await to delay the writer, or return [`Error::Backpressure`] to + /// refuse it. Any other error is a real failure. `incoming_bytes` is the + /// batch's in-memory size, so an implementation can size relief against + /// what is about to be admitted, or refuse before it is buffered. + async fn maybe_apply_backpressure( + &self, + incoming_bytes: usize, + shard: ShardMemory, + ) -> Result<()>; +} + +/// In-memory size of a batch list, for sizing an admission decision against +/// what is about to be inserted. Matches how `WalOnlyState` sizes its queue. +fn batches_memory_size(batches: &[RecordBatch]) -> usize { + batches.iter().map(|b| b.get_array_memory_size()).sum() +} + +/// The controller guarding this writer: the embedder's if one was injected, +/// otherwise lance's own [`LocalBackpressureController`]. +fn resolve_backpressure( + config: &ShardWriterConfig, + default_source: impl FnOnce() -> LocalSource, +) -> Arc { + match &config.backpressure { + Some(injected) => injected.clone(), + None => Arc::new(LocalBackpressureController::new(default_source(), config)), + } +} + +/// Which pool a [`LocalBackpressureController`] watches, and how it waits. +/// +/// The two write modes differ only in where the bytes live and whether a flush +/// can be awaited — both known in-crate, so this is an enum rather than another +/// trait. +enum LocalSource { + /// Memtable mode: active + frozen bytes, with a flush to wait on. + Memtable(Arc), + /// WAL-only mode: the pending queue, with no flush watcher, so the loop + /// falls back to a short sleep. + WalOnly(Arc), + /// Test-only: a synthetic reading, re-polled each iteration, so the loop + /// can be driven without standing up a writer. + #[cfg(test)] + Fake(Box (usize, Option) + Send + Sync>), +} + +impl LocalSource { + fn read(&self) -> (usize, Option) { + match self { + Self::Memtable(s) => (s.unflushed_memtable_bytes(), s.oldest_memtable_watcher()), + Self::WalOnly(s) => (s.estimated_size(), None), + #[cfg(test)] + Self::Fake(f) => f(), + } + } +} + +/// The per-shard memtable valve: lance's default when nothing is injected. +/// +/// Soft and blocking — it stalls the producer until a flush drains the pool and +/// never refuses a write. Replaced wholesale by +/// [`ShardWriterConfig::backpressure`], so an embedder that injects a +/// controller takes on the per-shard ceiling this provides (see +/// [`ShardMemory`]). +pub struct LocalBackpressureController { + source: LocalSource, + max_unflushed_memtable_bytes: usize, + log_interval: Duration, stats: Arc, } -impl BackpressureController { - /// Create a new backpressure controller. - pub fn new(config: ShardWriterConfig) -> Self { +impl Debug for LocalBackpressureController { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LocalBackpressureController") + .field( + "mode", + match &self.source { + LocalSource::Memtable(_) => &"memtable", + LocalSource::WalOnly(_) => &"wal_only", + #[cfg(test)] + LocalSource::Fake(_) => &"fake", + }, + ) + .field( + "max_unflushed_memtable_bytes", + &self.max_unflushed_memtable_bytes, + ) + .finish() + } +} + +impl LocalBackpressureController { + fn new(source: LocalSource, config: &ShardWriterConfig) -> Self { Self { - config, + source, + max_unflushed_memtable_bytes: config.max_unflushed_memtable_bytes, + log_interval: config.backpressure_log_interval, stats: Arc::new(BackpressureStats::new()), } } @@ -712,31 +908,30 @@ impl BackpressureController { pub fn stats(&self) -> &Arc { &self.stats } +} - /// Check and apply backpressure if needed. - /// - /// This method blocks if the system is under memory pressure, waiting for - /// frozen memtables to be flushed to storage until under threshold. - /// - /// Backpressure is applied when: - /// - `unflushed_memtable_bytes` >= `max_unflushed_memtable_bytes` - /// - /// # Arguments - /// - `get_state`: Closure that returns current (unflushed_memtable_bytes, oldest_memtable_watcher) +#[async_trait::async_trait] +impl BackpressureController for LocalBackpressureController { + /// Blocks while this shard's unflushed bytes are at or above + /// `max_unflushed_memtable_bytes`, waiting on the oldest flush. Never + /// returns [`Error::Backpressure`] — this valve only ever delays. /// - /// The closure is called in a loop to get fresh state after each wait. - pub async fn maybe_apply_backpressure(&self, mut get_state: F) -> Result<()> - where - F: FnMut() -> (usize, Option), - { + /// Reads its own source rather than the passed-in `shard`: it must re-poll + /// each iteration to see the flush drain, and in WAL-only mode the pool it + /// guards is the pending queue. + async fn maybe_apply_backpressure( + &self, + _incoming_bytes: usize, + _shard: ShardMemory, + ) -> Result<()> { let start = std::time::Instant::now(); let mut iteration = 0u32; loop { - let (unflushed_memtable_bytes, oldest_watcher) = get_state(); + let (unflushed_memtable_bytes, oldest_watcher) = self.source.read(); // Check if under threshold - if unflushed_memtable_bytes < self.config.max_unflushed_memtable_bytes { + if unflushed_memtable_bytes < self.max_unflushed_memtable_bytes { if iteration > 0 { let wait_ms = start.elapsed().as_millis() as u64; self.stats.record(wait_ms); @@ -748,18 +943,18 @@ impl BackpressureController { debug!( "Backpressure triggered: unflushed_bytes={}, max={}, iteration={}", - unflushed_memtable_bytes, self.config.max_unflushed_memtable_bytes, iteration + unflushed_memtable_bytes, self.max_unflushed_memtable_bytes, iteration ); // Wait for oldest memtable to flush if let Some(mut mem_watcher) = oldest_watcher { tokio::select! { _ = mem_watcher.await_value() => {} - _ = tokio::time::sleep(self.config.backpressure_log_interval) => { + _ = tokio::time::sleep(self.log_interval) => { warn!( "Backpressure wait timeout, continuing to wait: unflushed_bytes={}, interval={}s, iteration={}", unflushed_memtable_bytes, - self.config.backpressure_log_interval.as_secs(), + self.log_interval.as_secs(), iteration ); } @@ -789,12 +984,91 @@ struct FrozenMemTable { flushed_at_ms: Option, } +/// Resident memtable bytes for one shard: what is actually held in memory and +/// reclaimable only by flushing. +/// +/// Lives outside [`WriterState`] on purpose. A caller deciding whether to admit +/// a write must read these *while* the writer holds the write lock, which is +/// exactly when a `try_read()` on that lock fails — and tokio's `RwLock` is +/// write-preferring, so it fails whenever a writer is merely queued. Reading +/// through the lock would therefore report zero precisely under load. These are +/// plain atomics instead: every mutation happens under the write lock (so the +/// per-shard counters are single-writer), and any reader gets a relaxed load. +/// +/// Counts row data *and* index memory ([`MemTable::memory_size`]), since the +/// point is to bound RSS. +struct MemoryCounters { + /// The active memtable's resident bytes. + active: AtomicUsize, + /// Resident bytes of every sealed memtable still awaiting its flush. + frozen: AtomicUsize, + /// Process-wide `Σ(active + frozen)` over every shard, when the embedder + /// injects one. Opaque to lance — it only adds and subtracts. Bumped under + /// the write lock before the write is acked, so it is a safe upper bound. + pod: Option>, +} + +impl MemoryCounters { + fn new(pod: Option>) -> Self { + Self { + active: AtomicUsize::new(0), + frozen: AtomicUsize::new(0), + pod, + } + } + + /// Mirror a `from -> to` move on one counter into the pod-wide sum. + fn shift_pod(&self, from: usize, to: usize) { + let Some(pod) = &self.pod else { return }; + match to.cmp(&from) { + CmpOrdering::Greater => pod.fetch_add(to - from, Ordering::Relaxed), + CmpOrdering::Less => pod.fetch_sub(from - to, Ordering::Relaxed), + CmpOrdering::Equal => return, + }; + } + + /// Record the active memtable's size. Call under the write lock after any + /// mutation of it. + fn set_active(&self, bytes: usize) { + let prev = self.active.swap(bytes, Ordering::Relaxed); + self.shift_pod(prev, bytes); + } + + /// A memtable of `sealed` bytes was frozen and replaced by a fresh active + /// one of `new_active` bytes. Pod-wide this is ~net-zero: the bytes are + /// reclassified, not released. + fn seal(&self, sealed: usize, new_active: usize) { + self.set_active(new_active); + let prev = self.frozen.load(Ordering::Relaxed); + self.frozen.store(prev + sealed, Ordering::Relaxed); + self.shift_pod(prev, prev + sealed); + } + + /// Release a frozen memtable's bytes once its flush has committed. + fn release_frozen(&self, bytes: usize) { + let prev = self.frozen.load(Ordering::Relaxed); + let next = prev.saturating_sub(bytes); + self.frozen.store(next, Ordering::Relaxed); + self.shift_pod(prev, next); + } + + fn active(&self) -> usize { + self.active.load(Ordering::Relaxed) + } + + fn frozen(&self) -> usize { + self.frozen.load(Ordering::Relaxed) + } + + fn unflushed(&self) -> usize { + self.active() + self.frozen() + } +} + /// ShardWriter state shared across tasks. struct WriterState { memtable: MemTable, last_flushed_wal_entry_position: u64, - /// Total size of frozen memtables (for backpressure). - frozen_memtable_bytes: usize, /// Flush watchers for frozen memtables (for backpressure). frozen_flush_watchers: VecDeque<(usize, DurabilityWatcher)>, /// Sealed memtables, kept queryable so a concurrent reader sees no hole @@ -1013,6 +1287,9 @@ fn build_tombstone_batch( /// Shared state for writer operations. struct SharedWriterState { state: Arc>, + /// Resident-bytes counters, shared with the memtable flush handler (which + /// releases frozen bytes on commit). + memory: Arc, wal_flusher: Arc, wal_flush_tx: mpsc::UnboundedSender, memtable_flush_tx: mpsc::UnboundedSender, @@ -1031,6 +1308,7 @@ impl SharedWriterState { #[allow(clippy::too_many_arguments)] fn new( state: Arc>, + memory: Arc, wal_flusher: Arc, wal_flush_tx: mpsc::UnboundedSender, memtable_flush_tx: mpsc::UnboundedSender, @@ -1044,6 +1322,7 @@ impl SharedWriterState { ) -> Self { Self { state, + memory, wal_flusher, wal_flush_tx, memtable_flush_tx, @@ -1111,8 +1390,12 @@ impl SharedWriterState { )?; } - let frozen_size = old_memtable.estimated_size(); - state.frozen_memtable_bytes += frozen_size; + // Resident bytes, not `estimated_size`: the sealed memtable keeps its + // indexes in memory until the flush takes them. `flush_memtable` + // re-reads `memory_size` on the same (now immutable) memtable to + // release exactly this much, so the two stay paired. + let frozen_size = old_memtable.memory_size(); + self.memory.seal(frozen_size, state.memtable.memory_size()); let flush_watcher = old_memtable .get_memtable_flush_watcher() @@ -1248,16 +1531,15 @@ impl SharedWriterState { } impl SharedWriterState { + /// Resident bytes of the active plus all frozen memtables. + /// + /// A relaxed load of counters maintained under the write lock. It replaced + /// a `try_read()` on that lock, which reported `0` whenever the writer held + /// it — and because tokio's `RwLock` is write-preferring, also whenever a + /// writer was merely queued. That zeroed exactly the shards taking writes, + /// so a budget built on it read lowest under the heaviest load. fn unflushed_memtable_bytes(&self) -> usize { - // Total unflushed bytes = active memtable + all frozen memtables - self.state - .try_read() - .ok() - .map(|s| { - let active = s.memtable.estimated_size(); - active + s.frozen_memtable_bytes - }) - .unwrap_or(0) + self.memory.unflushed() } fn oldest_memtable_watcher(&self) -> Option { @@ -1304,7 +1586,7 @@ enum WriterMode { MemTable { state: Arc>, writer_state: Arc, - backpressure: BackpressureController, + backpressure: Arc, }, /// WAL-only mode: drainable pending-batch queue + WAL pipeline. No /// MemTable, no indexes, no Lance file flushing. @@ -1312,7 +1594,7 @@ enum WriterMode { state: Arc, wal_flush_tx: mpsc::UnboundedSender, trigger: StdRwLock, - backpressure: BackpressureController, + backpressure: Arc, }, } @@ -1532,10 +1814,13 @@ impl ShardWriter { // means "no entry covered yet." let initial_covered_wal_entry_position = next_wal_position.saturating_sub(1); + let memory = Arc::new(MemoryCounters::new(config.pod_memory_bytes.clone())); + // Replay above may already have filled the memtable. + memory.set_active(memtable.memory_size()); + let state = Arc::new(RwLock::new(WriterState { memtable, last_flushed_wal_entry_position: initial_covered_wal_entry_position, - frozen_memtable_bytes: 0, frozen_flush_watchers: VecDeque::new(), frozen_memtables: VecDeque::new(), flush_requested: false, @@ -1551,8 +1836,6 @@ impl ShardWriter { .with_storage_context(config.store_params.clone(), config.session.clone()), ); - let backpressure = BackpressureController::new(config.clone()); - // Background WAL flush handler — parallel WAL I/O + index updates. let wal_handler = WalFlushHandler::new(wal_flusher.clone(), Some(state.clone()), stats.clone()); @@ -1566,6 +1849,7 @@ impl ShardWriter { // It rebuilds the same secondary indexes on each flushed generation. let memtable_handler = MemTableFlushHandler::new( state.clone(), + memory.clone(), flusher, epoch, index_configs.to_vec(), @@ -1581,6 +1865,7 @@ impl ShardWriter { // Shared state used by `put()` to dispatch trigger checks. let writer_state = Arc::new(SharedWriterState::new( state.clone(), + memory, wal_flusher, wal_flush_tx, memtable_flush_tx, @@ -1593,6 +1878,10 @@ impl ShardWriter { index_configs.to_vec(), )); + // After `writer_state`: the default valve reads its byte counters. + let backpressure = + resolve_backpressure(config, || LocalSource::Memtable(writer_state.clone())); + Ok(WriterMode::MemTable { state, writer_state, @@ -1617,15 +1906,16 @@ impl ShardWriter { wal_flush_rx, )?; - // Reuse `BackpressureController` (which is keyed off - // `max_unflushed_memtable_bytes`) as the WAL-only backpressure - // budget. WAL-only callers feed it `WalOnlyState::estimated_size()`. - // Keeps the config knob meaningful in WAL-only mode and prevents - // the pending queue from growing unbounded under non-durable writes. - let backpressure = BackpressureController::new(config.clone()); + let state = Arc::new(WalOnlyState::default()); + + // Reuse the memtable valve (keyed off `max_unflushed_memtable_bytes`) + // as the WAL-only budget, fed `WalOnlyState::estimated_size()`. Keeps + // the config knob meaningful in WAL-only mode and prevents the pending + // queue from growing unbounded under non-durable writes. + let backpressure = resolve_backpressure(config, || LocalSource::WalOnly(state.clone())); Ok(WriterMode::WalOnly { - state: Arc::new(WalOnlyState::default()), + state, wal_flush_tx, trigger: StdRwLock::new(WalOnlyTriggerState::default()), backpressure, @@ -1820,7 +2110,7 @@ impl ShardWriter { batches: Vec, state_lock: &Arc>, writer_state: &Arc, - backpressure: &BackpressureController, + backpressure: &Arc, ) -> Result { let (result, watcher) = self .put_memtable_no_wait(batches, state_lock, writer_state, backpressure) @@ -1840,7 +2130,7 @@ impl ShardWriter { batches: Vec, state_lock: &Arc>, writer_state: &Arc, - backpressure: &BackpressureController, + backpressure: &Arc, ) -> Result<(WriteResult, Option)> { // Reject writes on a fenced writer before mutating the memtable, so a // poisoned writer can't drift further from the durable WAL. @@ -1848,12 +2138,10 @@ impl ShardWriter { // Apply backpressure if needed (before acquiring main lock) backpressure - .maybe_apply_backpressure(|| { - ( - writer_state.unflushed_memtable_bytes(), - writer_state.oldest_memtable_watcher(), - ) - }) + .maybe_apply_backpressure( + batches_memory_size(&batches), + ShardMemory(ShardMemorySource::Memtable(writer_state.memory.clone())), + ) .await?; let start = std::time::Instant::now(); @@ -1870,13 +2158,17 @@ impl ShardWriter { let end_pos = results.last().map(|(pos, _, _)| pos + 1).unwrap_or(0); let batch_positions = start_pos..end_pos; - // 2. Track last batch for WAL durability + // 2. Publish the memtable's new resident size before any freeze + // below reclassifies it, so `seal` reads the same value it moves. + writer_state.memory.set_active(state.memtable.memory_size()); + + // 3. Track last batch for WAL durability let durable_watcher = writer_state.track_batch_for_wal(end_pos.saturating_sub(1)); - // 3. Check if WAL flush should be triggered + // 4. Check if WAL flush should be triggered writer_state.maybe_trigger_wal_flush(&mut state); - // 4. Check if memtable flush is needed + // 5. Check if memtable flush is needed if let Err(e) = writer_state.maybe_trigger_memtable_flush(&mut state) { warn!("Failed to trigger memtable flush: {}", e); } @@ -1915,7 +2207,7 @@ impl ShardWriter { state: &Arc, wal_flush_tx: &mpsc::UnboundedSender, trigger: &StdRwLock, - backpressure: &BackpressureController, + backpressure: &Arc, ) -> Result { // Reject writes on a fenced writer before enqueuing — see // `put_memtable_no_wait`. @@ -1927,7 +2219,10 @@ impl ShardWriter { // shape as MemTable mode. WAL-only mode has no per-frozen-MemTable // watcher, so the backpressure loop falls back to its short sleep. backpressure - .maybe_apply_backpressure(|| (state.estimated_size(), None)) + .maybe_apply_backpressure( + batches_memory_size(&batches), + ShardMemory(ShardMemorySource::WalOnly(state.clone())), + ) .await?; let start = std::time::Instant::now(); @@ -2060,6 +2355,33 @@ impl ShardWriter { self.stats.clone() } + /// Resident bytes of the active memtable — row data plus its in-memory + /// indexes. `0` in WAL-only mode (no memtable). + /// + /// A relaxed load of a counter maintained under the write lock, so it stays + /// accurate for a shard that is actively taking writes. Read this together + /// with [`Self::frozen_bytes`] and **sum** them for the unflushed total; + /// never subtract one from the other, as two non-atomic reads can cross and + /// underflow. + /// + /// Only a flush reclaims these bytes, so this is the pool to bound against + /// OOM. It can far exceed `max_memtable_size`, which gates row data alone. + pub fn active_bytes(&self) -> usize { + match &self.mode { + WriterMode::MemTable { writer_state, .. } => writer_state.memory.active(), + WriterMode::WalOnly { .. } => 0, + } + } + + /// Resident bytes of sealed memtables whose flush has not yet committed. + /// `0` in WAL-only mode. See [`Self::active_bytes`]. + pub fn frozen_bytes(&self) -> usize { + match &self.mode { + WriterMode::MemTable { writer_state, .. } => writer_state.memory.frozen(), + WriterMode::WalOnly { .. } => 0, + } + } + /// Get the current shard manifest. pub async fn manifest(&self) -> Result> { self.manifest_store.read_latest().await @@ -2564,6 +2886,9 @@ impl WalFlushHandler { /// handler flushes in the background. struct MemTableFlushHandler { state: Arc>, + /// Shared with `SharedWriterState`; this handler releases a memtable's + /// frozen bytes once its flush commits. + memory: Arc, flusher: Arc, epoch: u64, /// Secondary index configs to rebuild on each flushed generation. When @@ -2579,8 +2904,10 @@ struct MemTableFlushHandler { } impl MemTableFlushHandler { + #[allow(clippy::too_many_arguments)] fn new( state: Arc>, + memory: Arc, flusher: Arc, epoch: u64, index_configs: Vec, @@ -2589,6 +2916,7 @@ impl MemTableFlushHandler { ) -> Self { Self { state, + memory, flusher, epoch, index_configs, @@ -2660,7 +2988,11 @@ impl MemTableFlushHandler { memtable: Arc, ) -> Result { let start = Instant::now(); - let memtable_size = memtable.estimated_size(); + // Read before the flush below takes the indexes, so this matches the + // `memory_size` that `freeze_memtable` charged to `frozen` for this + // same (now immutable) memtable. Reading it after the flush would + // release less than was charged and leak the counter upward. + let memtable_size = memtable.memory_size(); let flush_result = async { // Step 1: Wait for WAL flush completion (already queued at freeze time). @@ -2728,10 +3060,13 @@ impl MemTableFlushHandler { { let mut state = self.state.write().await; // Backpressure drain: unconditional so `wait_for_flush_drain` - // sees the watcher's error signal, not a dropped channel. - if let Some((_size, _watcher)) = state.frozen_flush_watchers.pop_front() { - state.frozen_memtable_bytes = - state.frozen_memtable_bytes.saturating_sub(memtable_size); + // sees the watcher's error signal, not a dropped channel. The + // popped entry gates the release but does not size it — flushes + // complete out of order, so the queue front may belong to another + // generation. Releasing *this* memtable's charge keeps the total + // right regardless of completion order. + if state.frozen_flush_watchers.pop_front().is_some() { + self.memory.release_frozen(memtable_size); } // Retire the frozen handle on commit success, keyed by generation // (non-FIFO completion is fine). Zero grace evicts here; otherwise @@ -4362,15 +4697,30 @@ mod tests { writer.close().await.unwrap(); } + fn fake_local( + config: &ShardWriterConfig, + read: impl Fn() -> (usize, Option) + Send + Sync + 'static, + ) -> LocalBackpressureController { + LocalBackpressureController::new(LocalSource::Fake(Box::new(read)), config) + } + + fn shard_memory(counters: Arc) -> ShardMemory { + ShardMemory(ShardMemorySource::Memtable(counters)) + } + + fn empty_shard_memory() -> ShardMemory { + shard_memory(Arc::new(MemoryCounters::new(None))) + } + #[tokio::test] async fn test_no_backpressure_when_under_threshold() { let config = ShardWriterConfig::default().with_max_unflushed_memtable_bytes(1024 * 1024); // 1MB - let controller = BackpressureController::new(config); + let controller = fake_local(&config, || (100, None)); // Should return immediately - well under threshold (100 bytes < 1MB) controller - .maybe_apply_backpressure(|| (100, None)) + .maybe_apply_backpressure(0, empty_shard_memory()) .await .unwrap(); @@ -4386,19 +4736,19 @@ mod tests { .with_max_unflushed_memtable_bytes(100) // Very low threshold .with_backpressure_log_interval(Duration::from_millis(50)); - let controller = BackpressureController::new(config); - // Simulate: starts at 1000 bytes, drops by 400 each call (simulating memtable flushes) let call_count = Arc::new(AtomicUsize::new(0)); let call_count_clone = call_count.clone(); + let controller = fake_local(&config, move || { + let count = call_count_clone.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // 1000 -> 600 -> 200 -> under threshold (need 3 iterations) + let unflushed = 1000usize.saturating_sub(count * 400); + (unflushed, None) + }); + controller - .maybe_apply_backpressure(move || { - let count = call_count_clone.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - // 1000 -> 600 -> 200 -> under threshold (need 3 iterations) - let unflushed = 1000usize.saturating_sub(count * 400); - (unflushed, None) - }) + .maybe_apply_backpressure(0, empty_shard_memory()) .await .unwrap(); @@ -4408,6 +4758,152 @@ mod tests { assert_eq!(controller.stats().count(), 1); } + /// Records what it saw and answers with a fixed verdict. + #[derive(Debug)] + struct SpyController { + seen: Arc>>, + reject: bool, + } + + #[async_trait::async_trait] + impl BackpressureController for SpyController { + async fn maybe_apply_backpressure( + &self, + incoming_bytes: usize, + shard: ShardMemory, + ) -> Result<()> { + self.seen + .write() + .unwrap() + .push((incoming_bytes, shard.unflushed())); + if self.reject { + return Err(Error::backpressure("full")); + } + Ok(()) + } + } + + /// An injected controller *replaces* the built-in valve rather than + /// stacking on it, and is handed both the incoming batch size and what the + /// calling shard already holds — everything it needs to own the whole + /// policy, per-shard rules included. + #[tokio::test] + async fn test_injected_controller_replaces_default_and_sees_write_context() { + let seen = Arc::new(StdRwLock::new(Vec::new())); + let spy = Arc::new(SpyController { + seen: seen.clone(), + reject: false, + }); + let config = ShardWriterConfig::default().with_backpressure(spy); + + // A source that would trip the built-in valve immediately, to prove it + // is not the thing being consulted. + let default_polls = Arc::new(AtomicUsize::new(0)); + let default_polls_clone = default_polls.clone(); + let controller = resolve_backpressure(&config, || { + LocalSource::Fake(Box::new(move || { + default_polls_clone.fetch_add(1, Ordering::Relaxed); + (usize::MAX, None) + })) + }); + + let counters = Arc::new(MemoryCounters::new(None)); + counters.set_active(700); + counters.seal(300, 700); + controller + .maybe_apply_backpressure(4096, shard_memory(counters)) + .await + .unwrap(); + + assert_eq!( + default_polls.load(Ordering::Relaxed), + 0, + "the built-in valve must not run once a controller is injected" + ); + assert_eq!(*seen.read().unwrap(), vec![(4096, 1000)]); + } + + /// `ShardMemory` must stay live across polls: a controller that delays is + /// waiting for these bytes to fall, so a captured copy would spin forever. + #[tokio::test] + async fn test_shard_memory_reflects_drain_across_polls() { + #[derive(Debug)] + struct DrainWaiter { + polls: AtomicUsize, + } + + #[async_trait::async_trait] + impl BackpressureController for DrainWaiter { + async fn maybe_apply_backpressure(&self, _: usize, shard: ShardMemory) -> Result<()> { + while shard.unflushed() > 0 { + self.polls.fetch_add(1, Ordering::Relaxed); + tokio::task::yield_now().await; + } + Ok(()) + } + } + + let counters = Arc::new(MemoryCounters::new(None)); + counters.set_active(0); + counters.seal(4_096, 0); + + let controller = Arc::new(DrainWaiter { + polls: AtomicUsize::new(0), + }); + let gate = controller.clone(); + let view = shard_memory(counters.clone()); + let waiting = tokio::spawn(async move { gate.maybe_apply_backpressure(1, view).await }); + + // Let the gate observe the full pool, then drain it as a flush commit + // would. A stale copy would never see this and the task would hang. + tokio::task::yield_now().await; + counters.release_frozen(4_096); + + waiting.await.unwrap().unwrap(); + assert!(controller.polls.load(Ordering::Relaxed) > 0); + assert_eq!(counters.unflushed(), 0); + } + + /// With nothing injected, lance keeps its own per-shard valve. + #[tokio::test] + async fn test_default_backpressure_is_used_when_none_injected() { + let config = ShardWriterConfig::default().with_max_unflushed_memtable_bytes(100); + let polls = Arc::new(AtomicUsize::new(0)); + let polls_clone = polls.clone(); + let controller = resolve_backpressure(&config, || { + LocalSource::Fake(Box::new(move || { + polls_clone.fetch_add(1, Ordering::Relaxed); + (0, None) + })) + }); + + controller + .maybe_apply_backpressure(1, empty_shard_memory()) + .await + .unwrap(); + + assert_eq!(polls.load(Ordering::Relaxed), 1); + } + + /// A rejecting controller surfaces as `Error::Backpressure`, which is + /// distinguishable from a real failure without matching on the message. + #[tokio::test] + async fn test_injected_controller_rejects_with_backpressure_error() { + let spy = Arc::new(SpyController { + seen: Arc::new(StdRwLock::new(Vec::new())), + reject: true, + }); + let config = ShardWriterConfig::default().with_backpressure(spy); + let controller = + resolve_backpressure(&config, || LocalSource::Fake(Box::new(|| (0, None)))); + + let err = controller + .maybe_apply_backpressure(1, empty_shard_memory()) + .await + .unwrap_err(); + assert!(err.is_backpressure(), "expected backpressure, got {err:?}"); + } + #[test] fn test_record_put() { let stats = WriteStats::new();