From 4fd319add2505ff3076f18c5f34b001f82930210 Mon Sep 17 00:00:00 2001 From: Jordan Olshevski Date: Tue, 21 Jul 2026 17:00:24 +0000 Subject: [PATCH] storage: avoid B-tree cache rescans --- .../src/storage/btree/cow.rs | 66 ++++++++++++--- .../src/storage/btree/mod.rs | 26 ++---- .../src/storage/btree/tests.rs | 82 +++++++++++++++++++ cmd/unbounded-storage/src/storage/engine.rs | 6 +- cmd/unbounded-storage/tests/storage/tests.rs | 27 ++---- 5 files changed, 157 insertions(+), 50 deletions(-) diff --git a/cmd/unbounded-storage/src/storage/btree/cow.rs b/cmd/unbounded-storage/src/storage/btree/cow.rs index 119d83554..5b5a9db71 100644 --- a/cmd/unbounded-storage/src/storage/btree/cow.rs +++ b/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, + nodes: HashMap>, 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::() + + node.keys.len() * size_of::() + + node.children.len() * size_of::() +} + /// 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( stack.push((child, remaining_depth - 1)); } } - bytes += size_of::() - + keys.len() * size_of::() - + children.len() * size_of::(); - 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, pub retired_pages: Vec, + new_internal_nodes: HashMap>, +} + +/// 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, + result: &PathCopyResult, +) -> Arc { + 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( 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( 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>, retired_pages: RefCell>, + new_internal_nodes: RefCell>>, 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)); } diff --git a/cmd/unbounded-storage/src/storage/btree/mod.rs b/cmd/unbounded-storage/src/storage/btree/mod.rs index ee0160b2f..e3a965b6f 100644 --- a/cmd/unbounded-storage/src/storage/btree/mod.rs +++ b/cmd/unbounded-storage/src/storage/btree/mod.rs @@ -446,9 +446,10 @@ impl BTreeIndex { .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 { self.mutator.borrow().entries.get(key).copied() } @@ -487,7 +488,8 @@ impl BTreeIndex { // 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 BTreeIndex { 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 BTreeIndex { 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() } diff --git a/cmd/unbounded-storage/src/storage/btree/tests.rs b/cmd/unbounded-storage/src/storage/btree/tests.rs index a0082b53c..0707310ee 100644 --- a/cmd/unbounded-storage/src/storage/btree/tests.rs +++ b/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); diff --git a/cmd/unbounded-storage/src/storage/engine.rs b/cmd/unbounded-storage/src/storage/engine.rs index fda337724..570851a24 100644 --- a/cmd/unbounded-storage/src/storage/engine.rs +++ b/cmd/unbounded-storage/src/storage/engine.rs @@ -792,7 +792,7 @@ impl StorageEngine { 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 StorageEngine { 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()), diff --git a/cmd/unbounded-storage/tests/storage/tests.rs b/cmd/unbounded-storage/tests/storage/tests.rs index 298b13d58..41ff6bf91 100644 --- a/cmd/unbounded-storage/tests/storage/tests.rs +++ b/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,