Repository: cberner/redb
Pull request #295 (2022-07-21 04:57:51)

Fix bug that caused stale reads
Re-opening a table during a write transaction caused the new Table
instance to return stale read results


Patch:
diff --git a/src/transactions.rs b/src/transactions.rs
index bd4d93e..4448455 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -104,7 +104,6 @@ pub struct WriteTransaction<'db> {
     freed_tree: BtreeMut<'db, FreedTableKey, [u8]>,
     freed_pages: Rc<RefCell<Vec<PageNumber>>>,
     open_tables: RefCell<HashMap<String, &'static panic::Location<'static>>>,
-    pending_table_updates: RefCell<HashMap<String, Option<PageNumber>>>,
     completed: AtomicBool,
     durability: Durability,
 }
@@ -128,7 +127,6 @@ impl<'db> WriteTransaction<'db> {
             freed_tree: BtreeMut::new(freed_root, db.get_memory(), freed_pages.clone()),
             freed_pages,
             open_tables: RefCell::new(Default::default()),
-            pending_table_updates: RefCell::new(Default::default()),
             completed: Default::default(),
             durability: Durability::Immediate,
         })
@@ -212,9 +210,9 @@ impl<'db> WriteTransaction<'db> {
         table: &mut BtreeMut<K, V>,
     ) {
         self.open_tables.borrow_mut().remove(name).unwrap();
-        self.pending_table_updates
+        self.table_tree
             .borrow_mut()
-            .insert(name.to_string(), table.get_root());
+            .stage_update_table_root(name, table.get_root());
     }
 
     /// Delete the given table
@@ -268,11 +266,7 @@ impl<'db> WriteTransaction<'db> {
     /// All writes performed in this transaction will be visible to future transactions, and are
     /// durable as consistent with the [`Durability`] level set by [`Self::set_durability`]
     pub fn commit(mut self) -> Result {
-        for (name, root) in self.pending_table_updates.borrow_mut().drain() {
-            self.table_tree
-                .borrow_mut()
-                .update_table_root(&name, root)?;
-        }
+        self.table_tree.borrow_mut().flush_table_root_updates()?;
         match self.commit_inner() {
             Ok(_) => {
                 self.db.deallocate_write_transaction(self.transaction_id);
@@ -315,6 +309,7 @@ impl<'db> WriteTransaction<'db> {
     pub fn abort(self) -> Result {
         #[cfg(feature = "logging")]
         info!("Aborting transaction id={}", self.transaction_id);
+        self.table_tree.borrow_mut().clear_table_root_updates();
         self.mem.rollback_uncommitted_writes()?;
         self.db.deallocate_write_transaction(self.transaction_id);
         self.completed.store(true, Ordering::Release);
@@ -329,6 +324,8 @@ impl<'db> WriteTransaction<'db> {
             .oldest_live_read_transaction()
             .unwrap_or(self.transaction_id);
 
+        let root = self.table_tree.borrow_mut().flush_table_root_updates()?;
+
         self.process_freed_pages(oldest_live_read)?;
         if oldest_live_read < self.transaction_id {
             self.store_freed_pages()?;
@@ -342,7 +339,6 @@ impl<'db> WriteTransaction<'db> {
             }
         }
 
-        let root = self.table_tree.borrow().get_root();
         let freed_root = self.freed_tree.get_root();
 
         self.mem
@@ -352,11 +348,12 @@ impl<'db> WriteTransaction<'db> {
 
     // Commit without a durability guarantee
     pub(crate) fn non_durable_commit(&mut self) -> Result {
+        let root = self.table_tree.borrow_mut().flush_table_root_updates()?;
+
         // Store all freed pages for a future commit(), since we can't free pages during a
         // non-durable commit (it's non-durable, so could be rolled back anytime in the future)
         self.store_freed_pages()?;
 
-        let root = self.table_tree.borrow().get_root();
         let freed_root = self.freed_tree.get_root();
 
         self.mem
@@ -462,7 +459,13 @@ impl<'db> WriteTransaction<'db> {
 
     #[allow(dead_code)]
     pub(crate) fn print_debug(&self) {
-        if let Some(page) = self.table_tree.borrow().get_root() {
+        // Flush any pending updates to make sure we get the latest root
+        if let Some(page) = self
+            .table_tree
+            .borrow_mut()
+            .flush_table_root_updates()
+            .unwrap()
+        {
             eprintln!("Master tree:");
             let master_tree: Btree<str, InternalTableDefinition> = Btree::new(Some(page), self.mem);
             master_tree.print_debug(true);
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index 52e7da5..6fd10d0 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -861,6 +861,7 @@ impl TransactionalMemory {
     }
 
     // Safety: the caller must ensure that no references to the memory in `page` exist
+    // TODO: add debug_assertion to check for double-free
     pub(crate) unsafe fn free(&self, page: PageNumber) -> Result {
         // Zero fill the page to ensure that deleted data is not stored in the file
         let mut mut_page = self.get_page_mut(page);
@@ -889,6 +890,7 @@ impl TransactionalMemory {
 
     // Frees the page if it was allocated since the last commit. Returns true, if the page was freed
     // Safety: the caller must ensure that no references to the memory in `page` exist
+    // TODO: add debug_assertion to check for double-free
     pub(crate) unsafe fn free_if_uncommitted(&self, page: PageNumber) -> Result<bool> {
         if self.allocated_since_commit.lock().unwrap().remove(&page) {
             // Zero fill the page to ensure that deleted data is not stored in the file
diff --git a/src/tree_store/table_tree.rs b/src/tree_store/table_tree.rs
index 78014c1..2f4c628 100644
--- a/src/tree_store/table_tree.rs
+++ b/src/tree_store/table_tree.rs
@@ -7,6 +7,7 @@ use crate::types::{
 use crate::{DatabaseStats, Error, Result};
 use std::cell::RefCell;
 use std::cmp::max;
+use std::collections::HashMap;
 use std::mem::size_of;
 use std::ops::RangeFull;
 use std::rc::Rc;
@@ -247,6 +248,8 @@ impl<'a> Iterator for TableNameIter<'a> {
 pub(crate) struct TableTree<'txn> {
     tree: BtreeMut<'txn, str, InternalTableDefinition>,
     mem: &'txn TransactionalMemory,
+    // Cached updates from tables that have been closed. These must be flushed to the btree
+    pending_table_updates: HashMap<String, Option<PageNumber>>,
     freed_pages: Rc<RefCell<Vec<PageNumber>>>,
 }
 
@@ -259,32 +262,37 @@ impl<'txn> TableTree<'txn> {
         Self {
             tree: BtreeMut::new(master_root, mem, freed_pages.clone()),
             mem,
+            pending_table_updates: Default::default(),
             freed_pages,
         }
     }
 
-    pub(crate) fn get_root(&self) -> Option<PageNumber> {
-        self.tree.get_root()
+    // Queues an update to the table root
+    pub(crate) fn stage_update_table_root(&mut self, name: &str, table_root: Option<PageNumber>) {
+        self.pending_table_updates
+            .insert(name.to_string(), table_root);
     }
 
-    pub(crate) fn update_table_root(
-        &mut self,
-        name: &str,
-        table_root: Option<PageNumber>,
-    ) -> Result {
-        // Bypass .get_table() since the table types are dynamic
-        // TODO: optimize away this get()
-        let mut definition = self.tree.get(name).unwrap().unwrap();
-        // No-op if the root has not changed
-        if definition.table_root == table_root {
-            return Ok(());
-        }
-        definition.table_root = table_root;
-        // Safety: References into the master table are never returned to the user
-        unsafe {
-            self.tree.insert(name, &definition)?;
+    pub(crate) fn clear_table_root_updates(&mut self) {
+        self.pending_table_updates.clear();
+    }
+
+    pub(crate) fn flush_table_root_updates(&mut self) -> Result<Option<PageNumber>> {
+        for (name, table_root) in self.pending_table_updates.drain() {
+            // Bypass .get_table() since the table types are dynamic
+            // TODO: optimize away this get()
+            let mut definition = self.tree.get(&name).unwrap().unwrap();
+            // No-op if the root has not changed
+            if definition.table_root == table_root {
+                continue;
+            }
+            definition.table_root = table_root;
+            // Safety: References into the master table are never returned to the user
+            unsafe {
+                self.tree.insert(&name, &definition)?;
+            }
         }
-        Ok(())
+        Ok(self.tree.get_root())
     }
 
     // root_page: the root of the master table
@@ -303,7 +311,7 @@ impl<'txn> TableTree<'txn> {
         name: &str,
         table_type: TableType,
     ) -> Result<Option<InternalTableDefinition>> {
-        if let Some(definition) = self.tree.get(name)? {
+        if let Some(mut definition) = self.tree.get(name)? {
             if definition.get_type() != table_type {
                 return Err(Error::TableTypeMismatch(format!(
                     "{:?} is not of type {:?}",
@@ -323,6 +331,10 @@ impl<'txn> TableTree<'txn> {
                 )));
             }
 
+            if let Some(updated_root) = self.pending_table_updates.get(name) {
+                definition.table_root = *updated_root;
+            }
+
             Ok(Some(definition))
         } else {
             Ok(None)
@@ -349,6 +361,8 @@ impl<'txn> TableTree<'txn> {
                 }
             }
 
+            self.pending_table_updates.remove(name);
+
             // Safety: References into the master table are never returned to the user
             let found = unsafe { self.tree.remove(name)?.is_some() };
             return Ok(found);
@@ -395,7 +409,11 @@ impl<'txn> TableTree<'txn> {
 
         let mut iter = self.tree.range::<RangeFull, &str>(..)?;
         while let Some(entry) = iter.next() {
-            let definition = InternalTableDefinition::from_bytes(entry.value());
+            let mut definition = InternalTableDefinition::from_bytes(entry.value());
+            if let Some(updated_root) = self.pending_table_updates.get(str::from_bytes(entry.key()))
+            {
+                definition.table_root = *updated_root;
+            }
             let subtree_stats = btree_stats(
                 definition.table_root,
                 self.mem,


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 8763056..8c8f795 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -713,6 +713,29 @@ fn regression11() {
     tx.commit().unwrap();
 }
 
+#[test]
+// Test that for stale read bug when re-opening a table during a write
+fn regression12() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 1024 * 1024;
+    let db = unsafe { Database::create(tmpfile.path(), db_size).unwrap() };
+
+    let table_def: TableDefinition<u64, u64> = TableDefinition::new("x");
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        t.insert(&0, &0).unwrap();
+        assert_eq!(t.get(&0).unwrap().unwrap(), 0);
+        drop(t);
+
+        let t2 = tx.open_table(table_def).unwrap();
+        assert_eq!(t2.get(&0).unwrap().unwrap(), 0);
+    }
+    tx.commit().unwrap();
+}
+
 #[test]
 fn non_durable_read_isolation() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();