diff --git a/src/btree/tree/maintenance.rs b/src/btree/tree/maintenance.rs index 582da82..4b28a6a 100644 --- a/src/btree/tree/maintenance.rs +++ b/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,85 +23,181 @@ impl BTree { /// /// Returns the count of pages touched. pub async fn rekey_walk(&self) -> Result { + self.rekey_walk_unique(&mut BTreeMap::new()).await + } + + /// Rekey this tree while sharing traversal state with other live roots. + /// + /// 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, + ) -> Result { if self.root_page_id == 0 { return Ok(0); } - let mut stack: Vec = 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() { - // 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 body = g.body(); - let header = read_header(&body)?; - let is_leaf = header.kind == NodeKind::Leaf; - (is_leaf, body.to_vec()) + 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(&walked_as) = visited.get(&page_id) { + match walked_as { + PageKind::BTreeLeaf | PageKind::BTreeInternal => continue, + 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", + )); + } + } + } + // `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 - { - // 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, - ) + match node_kind { + NodeKind::Leaf => { + let leaf = Leaf::decode(guard.body_ref())?; + let overflow_roots: Vec = 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, overflow_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. - 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)?; - stack.push(internal.leftmost_child); - for entry in &internal.entries { - stack.push(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, + ) -> Result { + // 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(&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), + other => Err(PagedbError::page_kind_aliased( + root, + other.name(), + PageKind::OverflowRoot.name(), + )), + }; + } + 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(&walked_as) = visited.get(&next) { + return match walked_as { + PageKind::Overflow => Err(PagedbError::overflow_chain_cycle(root, next)), + other => Err(PagedbError::page_kind_aliased( + next, + other.name(), + PageKind::Overflow.name(), + )), + }; + } + 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) } @@ -366,3 +464,189 @@ impl BTree { 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 { + 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) -> BTree { + 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:?}" + ); + } +} diff --git a/src/btree/tree/scan.rs b/src/btree/tree/scan.rs index 77a9a26..e37fcf9 100644 --- a/src/btree/tree/scan.rs +++ b/src/btree/tree/scan.rs @@ -68,6 +68,52 @@ impl BTree { } } + /// Collect at most `limit` records at or after `start`, in ascending key + /// order. + /// + /// The bounded counterpart to [`Self::collect_all`], for callers that must + /// traverse a whole tree without holding it in memory at once. Like + /// `collect_all` it has no upper key bound, for the same reason: no + /// concrete maximum key is outside the valid domain, so a bounded "scan to + /// the end" would silently drop records at the top of the keyspace. + /// + /// Resume by passing the last returned key with a `0x00` byte appended. + /// That is the immediate successor in the key ordering — no key can sort + /// strictly between `k` and `k ‖ 0x00` — so paging this way never skips a + /// record and never returns one twice. A short batch means the tree ended. + pub async fn collect_batch_from( + &self, + start: &[u8], + limit: usize, + ) -> Result, Vec)>> { + if self.root_page_id == 0 || limit == 0 { + return Ok(Vec::new()); + } + let mut path = self.path_to_leaf_for_key(start).await?; + let mut out: Vec<(Vec, Vec)> = Vec::with_capacity(limit); + loop { + let leaf_id = *path.last().expect("non-empty path"); + let leaf = self.read_leaf(leaf_id).await?; + for (k, v) in &leaf.records { + if k.as_slice() < start { + continue; + } + if out.len() == limit { + return Ok(out); + } + let val = self.resolve_leaf_value(v).await?; + out.push((k.clone(), val)); + } + if out.len() == limit { + return Ok(out); + } + match self.next_leaf_after(&path).await? { + Some(next_path) => path = next_path, + None => return Ok(out), + } + } + } + /// Return the smallest key in the tree, or `None` if the tree is empty. /// Descends the leftmost spine only — O(tree height), not O(tree size). pub async fn first_key(&self) -> Result>> { diff --git a/src/errors.rs b/src/errors.rs index 23234ec..08a5b33 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -282,6 +282,19 @@ pub enum CorruptionDetail { /// has no terminator. Distinct from a truncated chain: the links /// authenticate, they just form a loop. OverflowChainCycle { root_page_id: u64, page_id: u64 }, + /// One physical page was reached twice in a single traversal under two + /// incompatible page kinds. + /// + /// Distinct from [`Self::NodeKindMismatch`], which is a disagreement inside + /// one page. Here every read authenticates and every body decodes: two live + /// references simply claim the same page for different roles, so at least + /// one of them points at a page that was freed and handed to another + /// object while still linked. + PageKindAliased { + page_id: u64, + walked_as: &'static str, + referenced_as: &'static str, + }, } /// Quota failure reason, distinguishing which resource was exhausted. @@ -398,6 +411,21 @@ impl PagedbError { }) } + /// Canonical constructor for one page claimed by two incompatible + /// references in a single traversal. + #[must_use] + pub const fn page_kind_aliased( + page_id: u64, + walked_as: &'static str, + referenced_as: &'static str, + ) -> Self { + Self::Corruption(CorruptionDetail::PageKindAliased { + page_id, + walked_as, + referenced_as, + }) + } + /// Canonical constructor for an incremental snapshot that cannot be /// applied to this handle's current identity or reader-visible state. #[must_use] diff --git a/src/pager/format/page_kind.rs b/src/pager/format/page_kind.rs index 676d47f..3bde8a6 100644 --- a/src/pager/format/page_kind.rs +++ b/src/pager/format/page_kind.rs @@ -56,6 +56,29 @@ impl PageKind { self as u8 } + /// Stable lowercase name, for naming a page's role in an error that a + /// human reads — both sides of a + /// [`CorruptionDetail::PageKindAliased`](crate::errors::CorruptionDetail::PageKindAliased), + /// for instance. Part of the diagnostic surface, not the format: the + /// discriminant byte is what persists. + #[must_use] + pub const fn name(self) -> &'static str { + match self { + Self::BTreeInternal => "btree_internal", + Self::BTreeLeaf => "btree_leaf", + Self::Free => "free", + Self::Spill => "spill", + Self::Overflow => "overflow", + Self::Counter => "counter", + Self::Catalog => "catalog", + Self::ApplyJournal => "apply_journal", + Self::OverflowRoot => "overflow_root", + Self::SegmentData => "segment_data", + Self::SegmentIndex => "segment_index", + Self::SegmentOverflow => "segment_overflow", + } + } + /// True iff this kind is legal in a `main.db` Format A page. #[must_use] pub fn is_main_db(self) -> bool { @@ -117,6 +140,28 @@ mod tests { } } + /// Two kinds sharing a name would make an alias error name the same role on + /// both sides and read as a contradiction of itself. + #[test] + fn names_are_distinct_across_every_kind() { + let kinds = [ + PageKind::BTreeInternal, + PageKind::BTreeLeaf, + PageKind::Free, + PageKind::Spill, + PageKind::Overflow, + PageKind::Counter, + PageKind::Catalog, + PageKind::ApplyJournal, + PageKind::OverflowRoot, + PageKind::SegmentData, + PageKind::SegmentIndex, + PageKind::SegmentOverflow, + ]; + let names: std::collections::BTreeSet<&str> = kinds.iter().map(|k| k.name()).collect(); + assert_eq!(names.len(), kinds.len(), "page kind names must be unique"); + } + #[test] fn domain_split() { assert!(PageKind::BTreeLeaf.is_main_db() && !PageKind::BTreeLeaf.is_segment()); diff --git a/src/txn/db/rekey/main.rs b/src/txn/db/rekey/main.rs index b872db1..53dfc24 100644 --- a/src/txn/db/rekey/main.rs +++ b/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,9 +17,16 @@ 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}; +/// Commit-history rows read per batch while rekeying the roots they name. +/// Rows are 40 bytes, so this is a few KiB resident regardless of how deep +/// retention runs. +const HISTORY_ROOT_BATCH: usize = 512; + impl Db { /// Rekey the reachable main database and every catalog-linked immutable /// segment under `new_mk_epoch`. @@ -191,6 +199,14 @@ impl Db { // target header was published. self.pager .set_active_mk_epoch(target_master_key.clone(), intent.target_mk_epoch); + // One traversal set spans every root below. Retained snapshots share + // most of their pages by construction, so this is what keeps the walk + // proportional to unique reachable pages instead of to the number of + // retained commits. It is bounded by the page count of the live + // database, strictly below the dirty-page set the same walk is already + // accumulating in the buffer pool for the single `flush_main` at the + // end — a page id and kind against a whole decrypted page. + let mut rewritten = BTreeMap::new(); let main_tree = BTree::open( self.pager.clone(), self.realm_id, @@ -198,16 +214,10 @@ impl Db { state.next_page_id, self.page_size, ); - main_tree.rekey_walk().await?; + main_tree.rekey_walk_unique(&mut rewritten).await?; 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?; + self.rewrite_retained_history_roots(state, &mut rewritten) + .await?; } if state.free_list_root_page_id != 0 { let (_, chain_pages) = crate::pager::freelist::read_chain( @@ -238,6 +248,84 @@ impl Db { Ok(()) } + /// Rewrite the commit-history index and every reader-visible root its + /// retained rows still name. + /// + /// `begin_read_at` resolves a commit through this index and then opens the + /// data and catalog roots recorded in the row. Copy-on-write means those + /// roots are usually unreachable from the current header, so a rekey that + /// walked only the current header would return `Ok(())`, retire the source + /// epoch, and leave every retained snapshot naming pages that no live key + /// can open. Rewriting them is part of the success contract, not extra + /// safety: failing here is correct, because completing after a partial walk + /// recreates exactly that defect. + /// + /// The row's own `next_page_id` bounds each historical tree. Using the + /// current allocator bound instead would silently widen an old tree's + /// addressable page space beyond what its metadata claims. + async fn rewrite_retained_history_roots( + &self, + state: &WriterState, + rewritten: &mut BTreeMap, + ) -> Result<()> { + 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_unique(rewritten).await?; + + // Retention can be unbounded, so the rows are streamed in fixed-size + // batches rather than collected: the resident cost is one batch, not + // one entry per retained commit. Rewriting the index first means each + // batch is read back through the pages this walk has already re-sealed. + let mut cursor: Vec = Vec::new(); + loop { + let batch = history_tree + .collect_batch_from(&cursor, HISTORY_ROOT_BATCH) + .await?; + let Some((last_key, _)) = batch.last() else { + return Ok(()); + }; + cursor.clear(); + cursor.extend_from_slice(last_key); + cursor.push(0); + let exhausted = batch.len() < HISTORY_ROOT_BATCH; + + for (_, value) in &batch { + let historical = decode_commit_meta(value)?; + // Only the data and catalog roots are reader-visible. The row's + // free-list root is writer-only metadata whose superseded chain + // pages a later commit may already have recycled; following it + // would reinterpret a live page as `PageKind::Free`, so only the + // current header's live chain is safe to walk. + 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(rewritten) + .await?; + } + } + + if exhausted { + return Ok(()); + } + } + } + async fn publish_rekey_target_header( &self, state: &mut WriterState, diff --git a/tests/rekey_basic.rs b/tests/rekey_basic.rs index 1af29ce..2e95327 100644 --- a/tests/rekey_basic.rs +++ b/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