fix(rekey): preserve retained historical roots

Author: presempathy-awb

src/btree/tree/maintenance.rs

@@ -1,5 +1,7 @@
 //! Maintenance walks: rekey under a new epoch and reachable-page collection.
 
+use std::collections::BTreeMap;
+
 use crate::Result;
 use crate::errors::PagedbError;
 use crate::pager::format::page_kind::PageKind;
@@ -21,20 +23,46 @@ impl<V: Vfs> BTree<V> {
     ///
     /// Returns the count of pages touched.
     pub async fn rekey_walk(&self) -> Result<u64> {
+        self.rekey_walk_unique(&mut BTreeMap::new()).await
+    }
+
+    /// 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.
+    pub(crate) async fn rekey_walk_unique(
+        &self,
+        visited: &mut BTreeMap<u64, PageKind>,
+    ) -> Result<u64> {
         if self.root_page_id == 0 {
             return Ok(0);
         }
-        let mut stack: Vec<u64> = vec![self.root_page_id];
+        let mut stack: Vec<(u64, u64)> = vec![(0, self.root_page_id)];
         let mut count: u64 = 0;
-        while let Some(page_id) = stack.pop() {
+        while let Some((parent_page_id, page_id)) = stack.pop() {
+            if is_reserved(page_id) {
+                return Err(PagedbError::reserved_page_referenced(
+                    parent_page_id,
+                    page_id,
+                ));
+            }
+            if let Some(kind) = visited.get(&page_id) {
+                match kind {
+                    PageKind::BTreeLeaf | PageKind::BTreeInternal => continue,
+                    _ => return Err(PagedbError::IllegalPageKind),
+                }
+            }
             // Determine kind by reading the node under its own header kind byte.
-            let (is_leaf, body_bytes) = {
-                let (g, _page_kind) = self.pager.read_main_node(page_id, self.realm_id).await?;
+            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, body.to_vec())
+                (is_leaf, page_kind, body.to_vec())
             };
+            visited.insert(page_id, page_kind);
 
             if is_leaf {
                 // Collect overflow chains referenced by this leaf.
@@ -45,37 +73,9 @@ impl<V: Vfs> BTree<V> {
                         ..
                     } = v
                     {
-                        // Rewrite the root page (`PageKind::OverflowRoot`).
-                        let root_info =
-                            overflow::read_root_page(&self.pager, self.realm_id, *ov_root).await?;
-                        self.pager
-                            .rewrite_page_under_current_epoch(
-                                *ov_root,
-                                self.realm_id,
-                                PageKind::OverflowRoot,
-                            )
+                        count += self
+                            .rekey_overflow_unique(page_id, *ov_root, visited)
                             .await?;
-                        count += 1;
-                        // Walk and rewrite chain pages (always PageKind::Overflow).
-                        let mut next = root_info.next;
-                        while next != 0 {
-                            let ov_guard = self
-                                .pager
-                                .read_main_page(next, self.realm_id, PageKind::Overflow)
-                                .await?;
-                            let ov_body = ov_guard.body();
-                            let (ov_next, _) = overflow::decode_overflow(&ov_body)?;
-                            drop(ov_guard);
-                            self.pager
-                                .rewrite_page_under_current_epoch(
-                                    next,
-                                    self.realm_id,
-                                    PageKind::Overflow,
-                                )
-                                .await?;
-                            count += 1;
-                            next = ov_next;
-                        }
                     }
                 }
                 // Rewrite the leaf page.
@@ -86,9 +86,13 @@ impl<V: Vfs> BTree<V> {
             } else {
                 // Internal node: push children onto stack.
                 let internal = internal::Internal::decode(&body_bytes)?;
-                stack.push(internal.leftmost_child);
+                if internal.leftmost_child != 0 {
+                    stack.push((page_id, internal.leftmost_child));
+                }
                 for entry in &internal.entries {
-                    stack.push(entry.right_child);
+                    if entry.right_child != 0 {
+                        stack.push((page_id, entry.right_child));
+                    }
                 }
                 // Rewrite the internal page.
                 self.pager
@@ -104,6 +108,55 @@ impl<V: Vfs> BTree<V> {
         Ok(count)
     }
 
+    async fn rekey_overflow_unique(
+        &self,
+        leaf_page_id: u64,
+        root: u64,
+        visited: &mut BTreeMap<u64, PageKind>,
+    ) -> Result<u64> {
+        if is_reserved(root) {
+            return Err(PagedbError::reserved_page_referenced(leaf_page_id, root));
+        }
+        if let Some(kind) = visited.get(&root) {
+            return match kind {
+                PageKind::OverflowRoot => Ok(0),
+                _ => Err(PagedbError::IllegalPageKind),
+            };
+        }
+        visited.insert(root, PageKind::OverflowRoot);
+
+        let root_info = overflow::read_root_page(&self.pager, self.realm_id, root).await?;
+        self.pager
+            .rewrite_page_under_current_epoch(root, self.realm_id, PageKind::OverflowRoot)
+            .await?;
+        let mut count = 1;
+        let mut next = root_info.next;
+        while next != 0 {
+            if is_reserved(next) {
+                return Err(PagedbError::reserved_page_referenced(root, next));
+            }
+            if let Some(kind) = visited.get(&next) {
+                return match kind {
+                    PageKind::Overflow => Err(PagedbError::overflow_chain_cycle(root, next)),
+                    _ => Err(PagedbError::IllegalPageKind),
+                };
+            }
+            visited.insert(next, PageKind::Overflow);
+            let guard = self
+                .pager
+                .read_main_page(next, self.realm_id, PageKind::Overflow)
+                .await?;
+            let (following, _) = overflow::decode_overflow(guard.body_ref())?;
+            drop(guard);
+            self.pager
+                .rewrite_page_under_current_epoch(next, self.realm_id, PageKind::Overflow)
+                .await?;
+            count += 1;
+            next = following;
+        }
+        Ok(count)
+    }
+
     /// Collect all page IDs reachable from this tree's root (internal nodes,
     /// leaves, overflow chains) into `out`. Used by the deep-walk integrity
     /// checker to identify orphan pages.

src/txn/db/rekey/main.rs

@@ -1,5 +1,6 @@
 //! Main-database rekey transition and durable intent publication.
 
+use std::collections::BTreeMap;
 use std::sync::atomic::Ordering;
 
 use subtle::ConstantTimeEq;
@@ -16,7 +17,9 @@ use crate::vfs::Vfs;
 
 #[cfg(test)]
 use super::super::core::RekeyTestFault;
-use super::super::core::{Db, WriterState, encode_free_list_root, encode_root_ref};
+use super::super::core::{
+    Db, WriterState, decode_commit_meta, encode_free_list_root, encode_root_ref,
+};
 use super::intent::{intent_proof, migrate_legacy, validate_intent_for_current_cipher};
 
 impl<V: Vfs + Clone> Db<V> {
@@ -192,23 +195,54 @@ impl<V: Vfs + Clone> Db<V> {
         // target header was published.
         self.pager
             .set_active_mk_epoch(target_master_key.clone(), intent.target_mk_epoch);
+        let mut rewritten = BTreeMap::new();
         let main_tree = BTree::open(
             self.pager.clone(),
             self.realm_id,
             state.root_page_id,
             state.next_page_id,
             self.page_size,
         );
-        main_tree.rekey_walk().await?;
-        if state.commit_history_root_page_id != 0 {
+        main_tree.rekey_walk_unique(&mut rewritten).await?;
+        let retained_history = if state.commit_history_root_page_id != 0 {
             let history_tree = BTree::open(
                 self.pager.clone(),
                 self.realm_id,
                 state.commit_history_root_page_id,
                 state.next_page_id,
                 self.page_size,
             );
-            history_tree.rekey_walk().await?;
+            let entries = history_tree.collect_all().await?;
+            history_tree.rekey_walk_unique(&mut rewritten).await?;
+            entries
+                .into_iter()
+                .map(|(_, value)| decode_commit_meta(&value))
+                .collect::<Result<Vec<_>>>()?
+        } else {
+            Vec::new()
+        };
+        for historical in retained_history {
+            for root_page_id in [
+                historical.active_root_page_id,
+                historical.catalog_root_page_id,
+            ] {
+                if root_page_id == 0 {
+                    continue;
+                }
+                BTree::open(
+                    self.pager.clone(),
+                    self.realm_id,
+                    root_page_id,
+                    historical.next_page_id,
+                    self.page_size,
+                )
+                .rekey_walk_unique(&mut rewritten)
+                .await?;
+            }
+            // Historical readers traverse only retained data and catalog
+            // roots. The recorded free-list root is writer-only metadata whose
+            // superseded chain pages may already have been recycled. Only the
+            // current header's live chain below is safe to follow.
         }
         if state.free_list_root_page_id != 0 {
             let (_, chain_pages) = crate::pager::freelist::read_chain(

tests/rekey_basic.rs

@@ -1,7 +1,10 @@
 //! Rekey integration tests: main-db and immutable-segment transitions.
 
+use pagedb::options::RetainPolicy;
 use pagedb::vfs::memory::MemVfs;
-use pagedb::{Db, Evictable, PagedbError, RealmId, SegmentKind, SegmentPageKind};
+use pagedb::{
+    Db, Evictable, OpenOptions, PagedbError, RealmId, RealmQuotas, SegmentKind, SegmentPageKind,
+};
 
 const PAGE: usize = 4096;
 const KEK0: [u8; 32] = [0xAA; 32];
@@ -57,6 +60,100 @@ async fn rekey_main_db_only() {
     drop(rx);
 }
 
+/// Rekey must rewrite roots named only by retained commit-history metadata
+/// before the source epoch is retired.
+#[tokio::test(flavor = "current_thread")]
+async fn rekey_preserves_retained_historical_snapshots_after_reopen() {
+    let vfs = MemVfs::new();
+    let options = OpenOptions::default().with_commit_history_retain(RetainPolicy::Unbounded);
+    let db = Db::open_internal_with_options(vfs.clone(), KEK0, PAGE, REALM, options)
+        .await
+        .unwrap();
+    db.set_realm_quotas(
+        REALM,
+        RealmQuotas {
+            max_pages: Some(10),
+            ..RealmQuotas::default()
+        },
+    )
+    .await
+    .unwrap();
+    let historical_large = vec![0xA5; PAGE * 3];
+    let latest_large = vec![0x5A; PAGE * 2];
+
+    let first = {
+        let mut tx = db.begin_write().await.unwrap();
+        for index in 0u16..256 {
+            tx.put(
+                format!("history-key-{index:04}").as_bytes(),
+                &index.to_le_bytes(),
+            )
+            .await
+            .unwrap();
+        }
+        tx.put(b"versioned", b"before-rekey").await.unwrap();
+        tx.put(b"historical-overflow", &historical_large)
+            .await
+            .unwrap();
+        tx.commit().await.unwrap()
+    };
+    db.set_realm_quotas(
+        REALM,
+        RealmQuotas {
+            max_pages: Some(20),
+            ..RealmQuotas::default()
+        },
+    )
+    .await
+    .unwrap();
+    {
+        let mut tx = db.begin_write().await.unwrap();
+        tx.put(b"versioned", b"latest").await.unwrap();
+        tx.put(b"historical-overflow", &latest_large).await.unwrap();
+        tx.commit().await.unwrap();
+    }
+
+    db.rekey_db(KEK0, 1).await.unwrap();
+    drop(db);
+
+    let reopened = Db::open_existing(vfs, KEK0, PAGE, REALM).await.unwrap();
+    let historical = reopened.begin_read_at(first).await.unwrap();
+    assert_eq!(
+        historical.get(b"versioned").await.unwrap().as_deref(),
+        Some(b"before-rekey".as_slice())
+    );
+    assert_eq!(
+        historical
+            .get(b"historical-overflow")
+            .await
+            .unwrap()
+            .as_deref(),
+        Some(historical_large.as_slice())
+    );
+    let first_index = 0u16.to_le_bytes();
+    assert_eq!(
+        historical
+            .get(b"history-key-0000")
+            .await
+            .unwrap()
+            .as_deref(),
+        Some(first_index.as_slice())
+    );
+    assert!(matches!(
+        historical.open_segment("absent").await,
+        Err(PagedbError::NotFound)
+    ));
+    let latest = reopened.begin_read().await.unwrap();
+    assert_eq!(
+        latest.get(b"versioned").await.unwrap().as_deref(),
+        Some(b"latest".as_slice())
+    );
+    assert_eq!(
+        latest.get(b"historical-overflow").await.unwrap().as_deref(),
+        Some(latest_large.as_slice())
+    );
+}
+
 // ── Test 2 ─────────────────────────────────────────────────────────────────
 
 /// Linked immutable segments are replaced under the target epoch while