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

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<String>) -> 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)

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)

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<Vec<u32>>: strong + weak refcounts, the Vec header, then the ids.
+        let published = 2 * std::mem::size_of::<usize>()
+            + std::mem::size_of::<Vec<u32>>()
+            + max_neighbors * std::mem::size_of::<u32>();
+        published + max_neighbors * std::mem::size_of::<ScoredPoint>()
+    }
+
     fn publish_from_ranked(&self, ranked: &[ScoredPoint]) {
         self.published.store(Arc::new(
             ranked.iter().map(|point| point.id).collect::<Vec<_>>(),
@@ -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::<LevelLinks>()
+            + (0..=target_level)
+                .map(|level| LevelLinks::allocated_bytes(max_neighbors(m, level)))
+                .sum::<usize>()
+    }
+
     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<VisitedList>,
     packed_level0: ArcSwap<PackedLevel>,
+    /// 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(&params, &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::<VisitedList>()
+                + capacity.div_ceil(WORD_BITS) * std::mem::size_of::<usize>());
 
         Ok(Self {
             params,
+            base_bytes: capacity * std::mem::size_of::<Node>() + 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::<usize>()
+            + neighbors.capacity() * std::mem::size_of::<u32>();
+
         // 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(())
     }
 }

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::<StoredArrowBatch>()
+            + self.capacity * (std::mem::size_of::<RowLookup>() + std::mem::size_of::<u64>())
+    }
+
     /// Number of committed vectors.
     pub fn committed_len(&self) -> usize {
         self.committed_len.load(Ordering::Acquire)

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

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<K> {
     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<K> SkipListCore<K> {
             arena: UnsafeCell::new(Arena::new()),
             height: AtomicUsize::new(1),
             len: AtomicUsize::new(0),
+            arena_bytes: AtomicUsize::new(0),
         }
     }
 
@@ -303,7 +310,8 @@ impl<K: Ord> SkipListWriter<K> {
         // only mutator, so no link changes between read and publish.
         let layout = node_layout::<K>(height);
         // SAFETY: single-writer exclusive access to the arena.
-        let node = unsafe { (*self.core.arena.get()).alloc(layout) } as *mut Node<K>;
+        let node = unsafe { (*self.core.arena.get()).alloc(layout, &self.core.arena_bytes) }
+            as *mut Node<K>;
         // 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<K> {
 }
 
 impl<K: Ord> SkipListReader<K> {
+    /// 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.

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<Vec<RowPosition>>) -> usize {
+            nulls
+                .lock()
+                .map(|n| n.capacity() * std::mem::size_of::<RowPosition>())
+                .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<DataType> {
         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