fix(rekey): stream retained-history roots in bounded batches

Rekeying the commit-history index used to collect every retained row
into memory before rewriting the roots it names, making resident cost
proportional to retention depth. Stream the rows in fixed-size batches
via a new bounded scan primitive instead, and detect a page reached
under two incompatible kinds during the shared traversal as a reported
corruption rather than a generic illegal-kind error.

Author: farhan-syah

src/btree/tree/maintenance.rs

@@ -28,10 +28,22 @@ impl<V: Vfs> BTree<V> {
 
     /// Rekey this tree while sharing traversal state with other live roots.
     ///
-    /// Retained snapshots commonly share most of their pages. A caller walking
-    /// several roots supplies one page-kind map so each physical page is
-    /// authenticated and rewritten exactly once without hiding a cross-kind
-    /// alias.
+    /// Copy-on-write snapshots share most of their physical pages, so a caller
+    /// rewriting several roots — the live tree plus every root a retained
+    /// commit-history row still names — passes one map and each page is
+    /// authenticated and re-sealed exactly once. Without it the work is
+    /// `O(roots × pages)` of AEAD rather than `O(unique pages)`.
+    ///
+    /// The map holds each page's *kind*, not just its id, because the two
+    /// failure modes are not the same. A page already walked as a node and
+    /// referenced again as a node is an ordinary snapshot share and is skipped;
+    /// the same id presented under a different role means one of the two
+    /// references survived the page being freed and reused, and is reported as
+    /// [`CorruptionDetail::PageKindAliased`](crate::errors::CorruptionDetail::PageKindAliased).
+    /// A plain id set could not tell them apart.
+    ///
+    /// Returns the number of pages this call rewrote — pages an earlier walk
+    /// already covered are not counted twice.
     pub(crate) async fn rekey_walk_unique(
         &self,
         visited: &mut BTreeMap<u64, PageKind>,
@@ -48,79 +60,108 @@ impl<V: Vfs> BTree<V> {
                     page_id,
                 ));
             }
-            if let Some(kind) = visited.get(&page_id) {
-                match kind {
+            if let Some(&walked_as) = visited.get(&page_id) {
+                match walked_as {
                     PageKind::BTreeLeaf | PageKind::BTreeInternal => continue,
-                    _ => return Err(PagedbError::IllegalPageKind),
+                    other => {
+                        // Named for the role, not a kind: which of the two node
+                        // kinds this page is has not been read yet, and cannot
+                        // be — the alias is decided before any read.
+                        return Err(PagedbError::page_kind_aliased(
+                            page_id,
+                            other.name(),
+                            "btree_node",
+                        ));
+                    }
                 }
             }
-            // Determine kind by reading the node under its own header kind byte.
-            let (is_leaf, page_kind, body_bytes) = {
-                let (g, page_kind) = self.pager.read_main_node(page_id, self.realm_id).await?;
-                let body = g.body();
-                let header = read_header(&body)?;
-                let is_leaf = header.kind == NodeKind::Leaf;
-                (is_leaf, page_kind, body.to_vec())
+            // `read_node_guard` is the only accessor that proves the
+            // authenticated envelope kind and the encrypted body header agree.
+            // This walk both records a page's kind and re-seals the page under
+            // it, so taking those two from different sources would let a
+            // mis-routed page be laundered into a freshly authenticated one and
+            // would leave `visited` describing a kind that is no longer on disk.
+            let (guard, node_kind) = self.read_node_guard(page_id).await?;
+            let page_kind = match node_kind {
+                NodeKind::Leaf => PageKind::BTreeLeaf,
+                NodeKind::Internal => PageKind::BTreeInternal,
             };
             visited.insert(page_id, page_kind);
 
-            if is_leaf {
-                // Collect overflow chains referenced by this leaf.
-                let leaf = Leaf::decode(&body_bytes)?;
-                for (_k, v) in &leaf.records {
-                    if let LeafValue::Overflow {
-                        root_page_id: ov_root,
-                        ..
-                    } = v
-                    {
+            match node_kind {
+                NodeKind::Leaf => {
+                    let leaf = Leaf::decode(guard.body_ref())?;
+                    let overflow_roots: Vec<u64> = leaf
+                        .records
+                        .iter()
+                        .filter_map(|(_, value)| match value {
+                            LeafValue::Overflow { root_page_id, .. } => Some(*root_page_id),
+                            LeafValue::Inline(_) => None,
+                        })
+                        .collect();
+                    drop(guard);
+
+                    for overflow_root in overflow_roots {
                         count += self
-                            .rekey_overflow_unique(page_id, *ov_root, visited)
+                            .rekey_overflow_unique(page_id, overflow_root, visited)
                             .await?;
                     }
                 }
-                // Rewrite the leaf page.
-                self.pager
-                    .rewrite_page_under_current_epoch(page_id, self.realm_id, PageKind::BTreeLeaf)
-                    .await?;
-                count += 1;
-            } else {
-                // Internal node: push children onto stack.
-                let internal = internal::Internal::decode(&body_bytes)?;
-                if internal.leftmost_child != 0 {
-                    stack.push((page_id, internal.leftmost_child));
-                }
-                for entry in &internal.entries {
-                    if entry.right_child != 0 {
-                        stack.push((page_id, entry.right_child));
+                NodeKind::Internal => {
+                    let node = internal::Internal::decode(guard.body_ref())?;
+                    drop(guard);
+
+                    // A zero child id is an absent slot, not a pointer.
+                    if node.leftmost_child != 0 {
+                        stack.push((page_id, node.leftmost_child));
+                    }
+                    for entry in &node.entries {
+                        if entry.right_child != 0 {
+                            stack.push((page_id, entry.right_child));
+                        }
                     }
                 }
-                // Rewrite the internal page.
-                self.pager
-                    .rewrite_page_under_current_epoch(
-                        page_id,
-                        self.realm_id,
-                        PageKind::BTreeInternal,
-                    )
-                    .await?;
-                count += 1;
             }
+
+            self.pager
+                .rewrite_page_under_current_epoch(page_id, self.realm_id, page_kind)
+                .await?;
+            count += 1;
         }
         Ok(count)
     }
 
+    /// Rewrite the overflow chain rooted at `root` — referenced by leaf
+    /// `leaf_page_id` — returning the number of pages this call rewrote.
+    ///
+    /// Overflow roots are reference-counted, so a chain reached a second time
+    /// through a different leaf (in this tree or in another snapshot's tree)
+    /// arrives at the *same* root and is skipped whole. A chain *page* reached
+    /// twice therefore cannot be a legitimate share: either the chain loops, or
+    /// two distinct roots claim one page. Both mean the chain has no honest
+    /// terminator, which is why neither is treated as a stopping condition.
     async fn rekey_overflow_unique(
         &self,
         leaf_page_id: u64,
         root: u64,
         visited: &mut BTreeMap<u64, PageKind>,
     ) -> Result<u64> {
+        // Unlike an internal child slot or a chain terminator, zero is not a
+        // valid overflow root: an `Overflow` leaf value always owns at least
+        // its root page.
         if is_reserved(root) {
             return Err(PagedbError::reserved_page_referenced(leaf_page_id, root));
         }
-        if let Some(kind) = visited.get(&root) {
-            return match kind {
+        if let Some(&walked_as) = visited.get(&root) {
+            return match walked_as {
+                // Already rewritten with its whole chain, by a leaf holding the
+                // other reference to this refcounted value.
                 PageKind::OverflowRoot => Ok(0),
-                _ => Err(PagedbError::IllegalPageKind),
+                other => Err(PagedbError::page_kind_aliased(
+                    root,
+                    other.name(),
+                    PageKind::OverflowRoot.name(),
+                )),
             };
         }
         visited.insert(root, PageKind::OverflowRoot);
@@ -135,10 +176,14 @@ impl<V: Vfs> BTree<V> {
             if is_reserved(next) {
                 return Err(PagedbError::reserved_page_referenced(root, next));
             }
-            if let Some(kind) = visited.get(&next) {
-                return match kind {
+            if let Some(&walked_as) = visited.get(&next) {
+                return match walked_as {
                     PageKind::Overflow => Err(PagedbError::overflow_chain_cycle(root, next)),
-                    _ => Err(PagedbError::IllegalPageKind),
+                    other => Err(PagedbError::page_kind_aliased(
+                        next,
+                        other.name(),
+                        PageKind::Overflow.name(),
+                    )),
                 };
             }
             visited.insert(next, PageKind::Overflow);
@@ -419,3 +464,189 @@ impl<V: Vfs> BTree<V> {
         None
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use crate::errors::CorruptionDetail;
+    use crate::pager::format::page_kind::PageKind;
+    use crate::vfs::memory::MemVfs;
+    use crate::{Db, PagedbError, RealmId};
+
+    use super::super::core::BTree;
+
+    const PAGE: usize = 4096;
+    const KEK: [u8; 32] = [0x7B; 32];
+    const REALM: RealmId = RealmId::new([0x7C; 16]);
+
+    /// A tree deep enough to hold internal nodes, plus one overflow value so
+    /// the walk covers both the node and the chain path.
+    async fn populated_db() -> Db<MemVfs> {
+        let db = Db::open_internal(MemVfs::new(), KEK, PAGE, REALM)
+            .await
+            .unwrap();
+        let mut txn = db.begin_write().await.unwrap();
+        for index in 0u16..256 {
+            txn.put(format!("key-{index:04}").as_bytes(), &index.to_le_bytes())
+                .await
+                .unwrap();
+        }
+        txn.put(b"overflowing", &vec![0x33; PAGE * 3])
+            .await
+            .unwrap();
+        txn.commit().await.unwrap();
+        db
+    }
+
+    async fn data_tree(db: &Db<MemVfs>) -> BTree<MemVfs> {
+        let state = db.writer.lock().await;
+        BTree::open(
+            db.pager.clone(),
+            REALM,
+            state.root_page_id,
+            state.next_page_id,
+            db.page_size,
+        )
+    }
+
+    /// The whole point of the shared map: snapshots overlap almost entirely, so
+    /// a second root covering already-walked pages must cost nothing. Without
+    /// it, rekey re-authenticates and re-seals every shared page once per
+    /// retained root.
+    #[tokio::test(flavor = "current_thread")]
+    async fn a_second_walk_over_shared_pages_rewrites_nothing() {
+        let db = populated_db().await;
+        let tree = data_tree(&db).await;
+
+        let mut shared = BTreeMap::new();
+        let first = tree.rekey_walk_unique(&mut shared).await.unwrap();
+        assert!(
+            first > 1,
+            "fixture must span more than a single page, walked {first}"
+        );
+        let recorded = shared.len();
+
+        let second = tree.rekey_walk_unique(&mut shared).await.unwrap();
+        assert_eq!(
+            second, 0,
+            "every page was already covered by the first walk"
+        );
+        assert_eq!(
+            shared.len(),
+            recorded,
+            "a repeat walk must not discover new pages"
+        );
+    }
+
+    /// Each walk on its own map does the full work — the dedup above is a
+    /// property of the shared map, not of the tree being walked twice.
+    #[tokio::test(flavor = "current_thread")]
+    async fn an_unshared_walk_covers_every_page_again() {
+        let db = populated_db().await;
+        let tree = data_tree(&db).await;
+
+        let first = tree.rekey_walk_unique(&mut BTreeMap::new()).await.unwrap();
+        let second = tree.rekey_walk().await.unwrap();
+        assert_eq!(first, second);
+    }
+
+    /// The map records kinds, not just ids, so that a page reached under two
+    /// different roles is a reported alias rather than a silent skip. Here a
+    /// real overflow root — already walked as `OverflowRoot` — is presented as
+    /// a B+ tree root, which is what a freed-and-reused page looks like from
+    /// the second reference.
+    #[tokio::test(flavor = "current_thread")]
+    async fn a_page_reached_under_two_kinds_is_reported_not_skipped() {
+        let db = populated_db().await;
+        let tree = data_tree(&db).await;
+
+        let mut shared = BTreeMap::new();
+        tree.rekey_walk_unique(&mut shared).await.unwrap();
+        let (&overflow_root, _) = shared
+            .iter()
+            .find(|(_, kind)| **kind == PageKind::OverflowRoot)
+            .expect("the fixture stores one overflow value");
+
+        let aliased = BTree::open(
+            db.pager.clone(),
+            REALM,
+            overflow_root,
+            overflow_root + 1,
+            db.page_size,
+        );
+        let error = aliased.rekey_walk_unique(&mut shared).await.unwrap_err();
+        assert!(
+            matches!(
+                error,
+                PagedbError::Corruption(CorruptionDetail::PageKindAliased {
+                    page_id,
+                    walked_as: "overflow_root",
+                    referenced_as: "btree_node",
+                }) if page_id == overflow_root
+            ),
+            "expected an alias naming the page and both roles, got {error:?}"
+        );
+    }
+
+    /// Rekey re-seals each page under the kind it walked it as. Taking that
+    /// kind from the body header while the envelope says otherwise would launder
+    /// a mis-routed page into a freshly authenticated one, so the walk must go
+    /// through the accessor that proves the two agree.
+    #[tokio::test(flavor = "current_thread")]
+    async fn a_page_whose_envelope_contradicts_its_body_is_never_re_sealed() {
+        let db = populated_db().await;
+        let (leaf_page_id, forged) = {
+            let state = db.writer.lock().await;
+            let tree = BTree::open(
+                db.pager.clone(),
+                REALM,
+                state.root_page_id,
+                state.next_page_id,
+                db.page_size,
+            );
+            let mut reachable = std::collections::BTreeSet::new();
+            tree.collect_all_page_ids(&mut reachable).await.unwrap();
+            let mut leaf = None;
+            for &page_id in &reachable {
+                if let Ok((_, PageKind::BTreeLeaf)) = db.pager.read_main_node(page_id, REALM).await
+                {
+                    leaf = Some(page_id);
+                    break;
+                }
+            }
+            (leaf.expect("the fixture has leaves"), state.next_page_id)
+        };
+
+        // Copy a live leaf's bytes under the internal-node envelope kind: the
+        // body stays structurally valid, so the only defect is the routing.
+        let guard = db
+            .pager
+            .read_main_node(leaf_page_id, REALM)
+            .await
+            .unwrap()
+            .0;
+        let body = guard.body_ref().to_vec();
+        drop(guard);
+        db.pager
+            .write_main_page(forged, REALM, PageKind::BTreeInternal, &body)
+            .await
+            .unwrap();
+        db.pager.flush_main(REALM).await.unwrap();
+        db.pager.reset_main_pages();
+
+        let tree = BTree::open(db.pager.clone(), REALM, forged, forged + 1, db.page_size);
+        let error = tree.rekey_walk().await.unwrap_err();
+        assert!(
+            matches!(
+                error,
+                PagedbError::Corruption(CorruptionDetail::NodeKindMismatch {
+                    page_id: Some(page_id),
+                    expected: "internal",
+                    found: "leaf",
+                }) if page_id == forged
+            ),
+            "expected the kind disagreement to stop the walk, got {error:?}"
+        );
+    }
+}