diff --git a/doc/snapshot-recovery.md b/doc/snapshot-recovery.md index 904d1c7..ef43a0f 100644 --- a/doc/snapshot-recovery.md +++ b/doc/snapshot-recovery.md @@ -44,6 +44,33 @@ This is done by looking up the inner mapping we just loaded; it tells us where t If the child ptr points to an inner node, we need to correct the child node's child pointers recursively; if it points to a leaf node, i.e., it is a page ID rather than a virtual memory address, we don't need to correct it, as page ID translations are handled by the page table (described below). +### Allocation and snapshot sweep + +A newly allocated page must already be write-locked when its mapping becomes +visible to the snapshot sweep. Publishing an unlocked mapping and then calling +`try_write().unwrap()` allows the sweep to acquire a reader lock in between, +causing allocation to panic. Merely changing that call to a blocking write lock +still allows readers to observe a page before its initialization is complete. + +Both base-page allocation and mini-page insertion acquire the new entry's write +lock inside `MappingTable::insert_with`, while the insertion mutex still prevents +iterators from observing the new ID. The returned guard protects initialization +until the caller releases it. The callback must not reenter the mapping table; +it only locks the fresh, unpublished entry. + +The Shuttle regressions in `storage::tests::allocation_publication` explore up to +10,000 schedules per allocation path. A snapshot-style iterator probes for a reader +lock on the newly published entry and, if successful, holds it while waiting for +allocation to finish. The nonblocking probe avoids starving the writer in an +unfair DFS schedule. The allocator must complete without a panic or deadlock, +and the reader must observe completed initialization. On the unfixed code, both tests reproduce the +allocation `try_write().unwrap()` panic without relying on spurious weak-CAS +failures. + +```text +cargo test --release --lib --features shuttle storage::tests::allocation_publication -- --test-threads=1 +``` + ### Leaf mapping Leaf mapping is a page table that maps page ID to disk offset. diff --git a/src/storage.rs b/src/storage.rs index 187af44..251afbc 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -145,9 +145,8 @@ impl PageTable { counter!(AllocDiskID); let loc = PageLocation::Base(self.vfs.alloc_offset(self.config.leaf_page_size)); // Allocate space in disk for a full leaf page let entry = RwLock::new(loc); - let (id, value) = self.table.insert(entry); + let (id, lock_guard) = self.table.insert_with(entry, |value| value.write()); let pid = PageID::from_id(id); - let lock_guard = value.try_write().unwrap(); let base_ptr = LeafNode::make_base_page(self.config.leaf_page_size, snapshot_version); let x_locked = LeafEntryXLocked::with_buffer( lock_guard, @@ -178,9 +177,8 @@ impl PageTable { } } let entry = RwLock::new(mini_loc); - let (id, value) = self.table.insert(entry); + let (id, lock_guard) = self.table.insert_with(entry, |value| value.write()); let pid = PageID::from_id(id); - let lock_guard = value.try_write().unwrap(); let x_locked = LeafEntryXLocked::new( lock_guard, @@ -475,4 +473,104 @@ mod tests { } } } + + #[cfg(feature = "shuttle")] + mod allocation_publication { + use crate::{ + nodes::{leaf_node::MiniPageNextLevel, LeafNode, PageID}, + snapshot::INVALID_SNAPSHOT_VERSION, + storage::PageLocation, + sync::{ + atomic::{AtomicBool, Ordering}, + thread, Arc, + }, + BfTree, Config, + }; + + fn check_allocation(mini_page: bool) { + shuttle::check_dfs( + move || { + let mut config = Config::new(":memory:", 4096 * 4); + config.cache_only(mini_page); + let tree = Arc::new(BfTree::with_config(config, None).unwrap()); + let next_id = tree.storage.page_table.table.peek_next_id(); + let initialized = Arc::new(AtomicBool::new(false)); + + let writer_tree = tree.clone(); + let writer_initialized = initialized.clone(); + let mut writer = Some(thread::spawn(move || { + let storage = &writer_tree.storage; + let mini_guard = mini_page.then(|| { + storage + .alloc_mini_page(storage.config.leaf_page_size) + .unwrap() + }); + let (pid, leaf) = if let Some(ref mini_guard) = mini_guard { + LeafNode::initialize_mini_page( + mini_guard, + storage.config.leaf_page_size, + MiniPageNextLevel::new_null(), + true, + INVALID_SNAPSHOT_VERSION, + ); + storage + .page_table + .insert_mini_page_mapping(PageLocation::Mini( + mini_guard.as_ptr() as *mut LeafNode + )) + } else { + storage + .page_table + .alloc_base_page_mapping(INVALID_SNAPSHOT_VERSION) + }; + assert_eq!(pid.as_id(), next_id); + writer_initialized.store(true, Ordering::Release); + thread::yield_now(); + drop(leaf); + })); + + // Snapshot sweep discovers pages through this iterator, not through the tree. + for (lock, pid) in tree.storage.page_table.iter() { + if pid.as_id() != next_id { + continue; + } + // Avoid the blocking lock's spin loop under Shuttle's unfair DFS schedules. + if let Ok(guard) = lock.try_read() { + // The allocator must be able to finish while a published page is read. + writer.take().unwrap().join().unwrap(); + assert!( + initialized.load(Ordering::Acquire), + "a published page must be write-locked until initialization finishes" + ); + assert!(matches!( + &*guard, + PageLocation::Base(_) | PageLocation::Mini(_) + )); + } + } + if let Some(writer) = writer { + writer.join().unwrap(); + } + assert!(initialized.load(Ordering::Acquire)); + assert_eq!(tree.storage.page_table.table.peek_next_id(), next_id + 1); + + tree.storage + .page_table + .get_mut(&PageID::from_id(next_id)) + .dealloc_self(&tree.storage, mini_page); + }, + Some(10_000), + ); + } + + #[test] + fn base_allocation_is_locked_before_publication() { + check_allocation(false); + } + + #[test] + fn mini_allocation_is_locked_before_publication() { + check_allocation(true); + } + } } diff --git a/src/utils/mapping_table.rs b/src/utils/mapping_table.rs index 05b0092..e041d17 100644 --- a/src/utils/mapping_table.rs +++ b/src/utils/mapping_table.rs @@ -189,6 +189,16 @@ impl MappingTable { /// /// The id can be used to retrieve the record later using `get`. pub fn insert(&self, val: T) -> (u64, &T) { + self.insert_with(val, |record| record) + } + + /// Initialize a record before its ID becomes visible to iterators. + /// The callback runs under the insertion mutex and must not reenter the table. + pub(crate) fn insert_with<'a, R>( + &'a self, + val: T, + initialize: impl FnOnce(&'a T) -> R, + ) -> (u64, R) { let mut states = self.states.lock().unwrap(); let page_id = states.next_id; @@ -196,6 +206,6 @@ impl MappingTable { self.set(page_id, val, &mut states); - (page_id, self.get(page_id)) + (page_id, initialize(self.get(page_id))) } }