storage: avoid B-tree cache rescans

cmd/unbounded-storage/src/storage/btree/cow.rs

@@ -59,10 +59,11 @@ const MAX_TRAVERSAL_NODES: u64 = 1 << 24;
 
 /// Immutable decoded cache of internal B+tree nodes reachable from
 /// a published root. Leaf pages are deliberately excluded so lookup
-/// memory scales with the upper levels only.
+/// memory scales with the upper levels only. `bytes` estimates decoded
+/// node payload and excludes map buckets and `Arc` bookkeeping.
 #[derive(Default)]
 pub struct InternalNodeCache {
-    nodes: HashMap<Lba, CachedInternalNode>,
+    nodes: HashMap<Lba, Arc<CachedInternalNode>>,
     bytes: usize,
 }
 
@@ -77,10 +78,16 @@ impl InternalNodeCache {
     }
 
     fn get(&self, lba: Lba) -> Option<&CachedInternalNode> {
-        self.nodes.get(&lba)
+        self.nodes.get(&lba).map(Arc::as_ref)
     }
 }
 
+fn internal_node_bytes(node: &CachedInternalNode) -> usize {
+    size_of::<CachedInternalNode>()
+        + node.keys.len() * size_of::<PageKey>()
+        + node.children.len() * size_of::<Lba>()
+}
+
 /// Build a cache for every internal node reachable from `root_lba`.
 /// The traversal first discovers the tree height by following the
 /// leftmost spine, then reads only levels above the leaves. That keeps
@@ -121,16 +128,12 @@ pub async fn build_internal_cache<B: BlockDevice>(
                         stack.push((child, remaining_depth - 1));
                     }
                 }
-                bytes += size_of::<CachedInternalNode>()
-                    + keys.len() * size_of::<PageKey>()
-                    + children.len() * size_of::<Lba>();
-                nodes.insert(
-                    node_lba,
-                    CachedInternalNode {
-                        keys: keys.into_boxed_slice(),
-                        children: children.into_boxed_slice(),
-                    },
-                );
+                let node = Arc::new(CachedInternalNode {
+                    keys: keys.into_boxed_slice(),
+                    children: children.into_boxed_slice(),
+                });
+                bytes += internal_node_bytes(&node);
+                nodes.insert(node_lba, node);
             }
             Decoded::Internal { .. } | Decoded::Empty | Decoded::Meta { .. } if strict => {
                 return Err(Error::Corrupt);
@@ -413,6 +416,33 @@ pub struct PathCopyResult {
     pub new_root: Lba,
     pub new_pages: Vec<Lba>,
     pub retired_pages: Vec<Lba>,
+    new_internal_nodes: HashMap<Lba, Arc<CachedInternalNode>>,
+}
+
+/// Derive the next immutable internal-node cache from the parent cache and
+/// the pages produced by one path-copy commit. Untouched nodes retain their
+/// shared decoded representation; retired nodes are removed and freshly
+/// encoded internal nodes are inserted without reading them back from disk.
+pub fn update_internal_cache(
+    parent: &Arc<InternalNodeCache>,
+    result: &PathCopyResult,
+) -> Arc<InternalNodeCache> {
+    let mut nodes = parent.nodes.clone();
+    let mut bytes = parent.bytes;
+
+    for lba in &result.retired_pages {
+        if let Some(node) = nodes.remove(lba) {
+            bytes -= internal_node_bytes(&node);
+        }
+    }
+    for (&lba, node) in &result.new_internal_nodes {
+        if let Some(replaced) = nodes.insert(lba, node.clone()) {
+            bytes -= internal_node_bytes(&replaced);
+        }
+        bytes += internal_node_bytes(node);
+    }
+
+    Arc::new(InternalNodeCache { nodes, bytes })
 }
 
 /// Apply `sorted_ops` to the tree rooted at `parent_root` using
@@ -441,6 +471,7 @@ pub async fn apply_path_copy<B: BlockDevice>(
         txn_id,
         new_pages: RefCell::new(Vec::new()),
         retired_pages: RefCell::new(Vec::new()),
+        new_internal_nodes: RefCell::new(HashMap::new()),
         leaf_cap: max_leaf_entries(ps),
         internal_cap: max_internal_keys(ps),
         page_size: ps,
@@ -453,6 +484,7 @@ pub async fn apply_path_copy<B: BlockDevice>(
             new_root,
             new_pages: ctx.new_pages.into_inner(),
             retired_pages: ctx.retired_pages.into_inner(),
+            new_internal_nodes: ctx.new_internal_nodes.into_inner(),
         }),
         Err(e) => {
             free_all(ctx.allocator, &ctx.new_pages.borrow());
@@ -517,6 +549,7 @@ struct PathCopyCtx<'a, B: BlockDevice> {
     txn_id: u64,
     new_pages: RefCell<Vec<Lba>>,
     retired_pages: RefCell<Vec<Lba>>,
+    new_internal_nodes: RefCell<HashMap<Lba, Arc<CachedInternalNode>>>,
     leaf_cap: usize,
     internal_cap: usize,
     page_size: usize,
@@ -696,6 +729,13 @@ impl<'a, B: BlockDevice> PathCopyCtx<'a, B> {
             let lba = self.allocator.alloc()?;
             self.new_pages.borrow_mut().push(lba);
             let page = page::encode_internal(self.page_size, self.txn_id, &keys, &kids)?;
+            self.new_internal_nodes.borrow_mut().insert(
+                lba,
+                Arc::new(CachedInternalNode {
+                    keys: keys.into_boxed_slice(),
+                    children: kids.into_boxed_slice(),
+                }),
+            );
             out.push((chunk[0].0, lba));
             pages.push((lba, page));
         }

cmd/unbounded-storage/src/storage/btree/mod.rs

@@ -446,9 +446,10 @@ impl<B: BlockDevice> BTreeIndex<B> {
         .await
     }
 
-    /// Benchmark-only lookup path backed by the committed in-memory
-    /// mirror. This skips the terminal leaf read, so it must not be used
-    /// for production recovery/corruption semantics.
+    /// Look up an entry in the committed in-memory mirror. Mutation
+    /// bookkeeping may use this because the single mutator owns the mirror.
+    /// Client reads must normally hit the on-disk leaf to preserve corruption
+    /// detection; the engine only bypasses that check in benchmark mode.
     pub fn lookup_committed_mirror(&self, key: &PageKey) -> Option<LeafEntry> {
         self.mutator.borrow().entries.get(key).copied()
     }
@@ -487,7 +488,8 @@ impl<B: BlockDevice> BTreeIndex<B> {
         // Path-copy consumes the operations, so retain only the bounded
         // batch needed to update the mirror after the durable commit.
         let mirror_ops = sorted_ops.clone();
-        let parent_root = self.root.load().root_lba;
+        let parent = self.root.load();
+        let parent_root = parent.root_lba;
         let result = cow::apply_path_copy(
             &*self.device,
             &self.scratch,
@@ -497,17 +499,7 @@ impl<B: BlockDevice> BTreeIndex<B> {
             sorted_ops,
         )
         .await?;
-
-        let internal_cache =
-            match cow::build_internal_cache(&*self.device, &self.scratch, result.new_root, true)
-                .await
-            {
-                Ok(cache) => cache,
-                Err(e) => {
-                    cow::free_all(&self.allocator, &result.new_pages);
-                    return Err(e);
-                }
-            };
+        let internal_cache = cow::update_internal_cache(&parent.internal_cache, &result);
 
         let active = self.active_meta.get();
         // apply_path_copy above has already marked the new pages
@@ -597,8 +589,8 @@ impl<B: BlockDevice> BTreeIndex<B> {
         self.mutator.borrow().entries.len()
     }
 
-    /// Heap footprint of the immutable lookup cache owned by the
-    /// currently published snapshot.
+    /// Estimated decoded-node payload of the immutable lookup cache owned
+    /// by the current snapshot. Map and shared-reference overhead is excluded.
     pub fn lookup_cache_bytes(&self) -> usize {
         self.root.load().lookup_cache_bytes()
     }

cmd/unbounded-storage/src/storage/btree/tests.rs

@@ -145,6 +145,88 @@ fn large_batch_spans_multiple_leaves() {
     assert!(block_on(idx.lookup(&key(300))).unwrap().is_none());
 }
 
+#[test]
+fn commit_updates_internal_cache_without_rescanning_tree() {
+    let cfg = MockDeviceConfig {
+        page_size: 512,
+        capacity_pages: 512,
+        ..Default::default()
+    };
+    let dev = Arc::new(MockDevice::new(cfg));
+    let alloc = Arc::new(Allocator::new(512));
+    let scratch = ScratchPool::new(&*dev, 512, 8).expect("scratch pool");
+    let idx = block_on(BTreeIndex::open(
+        dev.clone(),
+        alloc,
+        scratch,
+        512,
+        false,
+        false,
+    ))
+    .unwrap();
+    let reads_before = dev.reads();
+
+    block_on(
+        idx.apply_batch(
+            (0..100u32)
+                .map(|i| Mutation::Insert {
+                    key: key(i),
+                    value: entry(1000 + i as u64),
+                })
+                .collect(),
+        ),
+    )
+    .unwrap();
+
+    assert_eq!(
+        dev.reads() - reads_before,
+        1,
+        "path copy should read only the previous single-leaf root",
+    );
+    assert!(idx.lookup_cache_bytes() > 0);
+
+    let reads_before_overwrite = dev.reads();
+    block_on(idx.apply_batch(vec![Mutation::Insert {
+        key: key(5),
+        value: entry(2005),
+    }]))
+    .unwrap();
+    assert_eq!(
+        dev.reads() - reads_before_overwrite,
+        3,
+        "second path copy should read one root, one internal node, and one leaf",
+    );
+
+    let reads_before_lookup = dev.reads();
+    assert_eq!(block_on(idx.lookup(&key(5))).unwrap(), Some(entry(2005)));
+    assert_eq!(block_on(idx.lookup(&key(95))).unwrap(), Some(entry(1095)));
+    assert_eq!(
+        dev.reads() - reads_before_lookup,
+        2,
+        "lookups should read terminal leaves but no replaced or shared internal nodes",
+    );
+
+    let pinned = idx.root.load_full();
+    let pinned_cache_bytes = pinned.lookup_cache_bytes();
+    block_on(
+        idx.apply_batch(
+            (0..100u32)
+                .map(|i| Mutation::Delete { key: key(i) })
+                .collect(),
+        ),
+    )
+    .unwrap();
+
+    assert_eq!(idx.live_entries(), 0);
+    assert_eq!(idx.lookup_cache_bytes(), 0);
+    assert!(block_on(idx.lookup(&key(5))).unwrap().is_none());
+    assert_eq!(
+        pinned.lookup_cache_bytes(),
+        pinned_cache_bytes,
+        "the old snapshot must retain its shared internal-node cache",
+    );
+}
+
 #[test]
 fn restart_from_meta_restores_entries() {
     let (dev, alloc) = fresh(128);

cmd/unbounded-storage/src/storage/engine.rs

@@ -792,7 +792,7 @@ impl<B: BlockDevice> StorageEngine<B> {
                     let prior = if let Some(prior) = states.get(key) {
                         *prior
                     } else {
-                        let prior = self.btree.lookup(key).await.ok().flatten();
+                        let prior = self.btree.lookup_committed_mirror(key);
                         states.insert(*key, prior);
                         prior
                     };
@@ -811,7 +811,7 @@ impl<B: BlockDevice> StorageEngine<B> {
                         let current = if let Some(current) = states.get(&victim.key) {
                             *current
                         } else {
-                            let current = self.btree.lookup(&victim.key).await.ok().flatten();
+                            let current = self.btree.lookup_committed_mirror(&victim.key);
                             states.insert(victim.key, current);
                             current
                         };
@@ -1040,11 +1040,13 @@ mod tests {
             done: overwrite,
         }]));
         let eviction = MutatorReply::new();
+        let reads_before = eng.device.reads();
         block_on(eng.process_batch(vec![MutatorReq::Delete {
             victims: vec![resident(key, 200)],
             done: eviction.clone(),
         }]));
 
+        assert_eq!(eng.device.reads(), reads_before);
         assert_eq!(eng.btree.lookup_committed_mirror(&key), Some(entry(201)));
         assert!(matches!(
             block_on(eviction.wait()),

cmd/unbounded-storage/tests/storage/tests.rs

@@ -156,22 +156,14 @@ proptest! {
     ///      are gone from the LRU but still live in the btree, so
     ///      `btree_entries` exceeds `resident_pages` by up to
     ///      `EVICT_SWEEP_TARGET` per failure.
-    ///   2. A corrupted btree-internal read of the prior-LBA
-    ///      probe in the mutator's `process_batch`. A flipped
-    ///      byte that makes `btree::lookup` return `Ok(None)`
-    ///      when a prior entry existed causes the engine to skip
-    ///      `retire_range(old)`; the new LBA is admitted to both
-    ///      sides but the old LBA stays in the LRU and `reverse`
-    ///      map even though its btree key was overwritten by the
-    ///      new insert. That orphans one LRU entry per
-    ///      corruption, so `resident_pages` can exceed
-    ///      `btree_entries` by up to `device_corruptions_injected`.
-    ///      A corruption that hits the path-copy descent inside
-    ///      `apply_batch` itself aborts the commit (the engine's
-    ///      `apply_node` surfaces `Decoded::Empty` as
-    ///      `Error::Corrupt` so a subtree is never silently
-    ///      dropped from the new tree), so it doesn't contribute
-    ///      to the gap.
+    ///   2. Corruption during the path-copy descent aborts the
+    ///      whole mutator batch (`apply_node` surfaces
+    ///      `Decoded::Empty` as `Error::Corrupt`) rather than
+    ///      publishing a partial tree. Corruption injection is
+    ///      retained as defensive slack for the failed batch's
+    ///      surrounding eviction bookkeeping. Prior-LBA probes do
+    ///      not contribute: the single mutator now reads those
+    ///      from the exact committed mirror rather than disk.
     ///   The data-write failure paths in `write_page_from` either
     ///   rewind both sides or touch neither, so they don't
     ///   contribute. `pending_free_len` is added as a small
@@ -221,8 +213,7 @@ proptest! {
             diff <= bound,
             "|resident_pages ({}) - btree_entries ({})| = {} exceeds bound {} \
              (device_io_errors={}, device_corruptions_injected={}, pending_free_len={}); \
-             LRU and index diverged beyond what failed evictions and corrupted btree \
-             reads can explain",
+             LRU and index diverged beyond what failed mutator batches can explain",
             resident, btree, diff, bound,
             report.device_io_errors, report.device_corruptions_injected,
             report.pending_free_len,