Repository: cberner/redb
Pull request #403 (2022-10-17 04:37:58)

Move savepoint methods to WriteTransaction
This allows users to construct certain race-free savepoint scenarios, and fixes several races in the fuzzer


Patch:
diff --git a/fuzz/fuzz_targets/fuzz_redb.rs b/fuzz/fuzz_targets/fuzz_redb.rs
index 6db9fb0..b6da97d 100644
--- a/fuzz/fuzz_targets/fuzz_redb.rs
+++ b/fuzz/fuzz_targets/fuzz_redb.rs
@@ -32,10 +32,11 @@ fn exec_table_inner(db: Arc<Database>, transactions: &[FuzzTransaction], referen
     for transaction in transactions.iter() {
         barrier.wait();
 
+        let mut txn = db.begin_write().unwrap();
+        let mut local_reference = reference.lock().unwrap().clone();
+
         if transaction.create_savepoint {
-            savepoints.push(db.savepoint()?);
-            // XXX: potentially there's race here. We should hold the lock before savepoint(),
-            // but that could deadlock, because savepoint() may block
+            savepoints.push(txn.savepoint()?);
             let guard = reference.lock().unwrap();
             reference_savepoints.push(guard.clone());
             drop(guard);
@@ -47,24 +48,15 @@ fn exec_table_inner(db: Arc<Database>, transactions: &[FuzzTransaction], referen
         let restore_to = transaction.restore_savepoint.value;
         if restore_to > 0 && restore_to <= savepoints.len() {
             let index = savepoints.len() - restore_to;
-            let result = db.restore_savepoint(&savepoints[index]);
-            // XXX: potentially there's race here. We should hold the lock before restore_savepoint(),
-            // but that could deadlock, because restore_savepoint() may block
-            let mut guard = reference.lock().unwrap();
+            let result = txn.restore_savepoint(&savepoints[index]);
             if result.is_ok() {
-                *guard = reference_savepoints[index].clone();
+                local_reference = reference_savepoints[index].clone();
             }
-            drop(guard);
             if result.is_err() && !matches!(result, Err(Error::InvalidSavepoint)) {
                 return result;
             }
         }
-        let mut txn = db.begin_write().unwrap();
-        // XXX: potentially there's race here. We should hold the lock before begin_write(),
-        // but that could deadlock, because begin_write() may block
-        let guard = reference.lock().unwrap();
-        let mut local_reference = guard.clone();
-        drop(guard);
+
         // We're not trying to test crash safety, so don't bother with durability
         if !transaction.durable {
             txn.set_durability(Durability::None);
@@ -181,10 +173,11 @@ fn exec_multimap_table_inner(db: Arc<Database>, transactions: &[FuzzTransaction]
     for transaction in transactions.iter() {
         barrier.wait();
 
+        let mut txn = db.begin_write().unwrap();
+        let mut local_reference = reference.lock().unwrap().clone();
+
         if transaction.create_savepoint {
-            savepoints.push(db.savepoint()?);
-            // XXX: potentially there's race here. We should hold the lock before savepoint(),
-            // but that could deadlock, because savepoint() may block
+            savepoints.push(txn.savepoint()?);
             let guard = reference.lock().unwrap();
             reference_savepoints.push(guard.clone());
             drop(guard);
@@ -196,24 +189,15 @@ fn exec_multimap_table_inner(db: Arc<Database>, transactions: &[FuzzTransaction]
         let restore_to = transaction.restore_savepoint.value;
         if restore_to > 0 && restore_to <= savepoints.len() {
             let index = savepoints.len() - restore_to;
-            let result = db.restore_savepoint(&savepoints[index]);
-            // XXX: potentially there's race here. We should hold the lock before restore_savepoint(),
-            // but that could deadlock, because restore_savepoint() may block
-            let mut guard = reference.lock().unwrap();
+            let result = txn.restore_savepoint(&savepoints[index]);
             if result.is_ok() {
-                *guard = reference_savepoints[index].clone();
+                local_reference = reference_savepoints[index].clone();
             }
-            drop(guard);
             if result.is_err() && !matches!(result, Err(Error::InvalidSavepoint)) {
                 return result;
             }
         }
-        let mut txn = db.begin_write().unwrap();
-        // XXX: potentially there's race here. We should hold the lock before begin_write(),
-        // but that could deadlock, because begin_write() may block
-        let guard = reference.lock().unwrap();
-        let mut local_reference = guard.clone();
-        drop(guard);
+
         // We're not trying to test crash safety, so don't bother with durability
         if !transaction.durable {
             txn.set_durability(Durability::None);
diff --git a/src/db.rs b/src/db.rs
index 3d66c18..7085dc3 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -438,6 +438,14 @@ impl Database {
         })
     }
 
+    fn allocate_read_transaction(&self) -> Result<TransactionId> {
+        let mut guard = self.live_read_transactions.lock().unwrap();
+        let id = self.mem.get_last_committed_transaction_id()?;
+        guard.entry(id).and_modify(|x| *x += 1).or_insert(1);
+
+        Ok(id)
+    }
+
     pub(crate) fn deallocate_read_transaction(&self, id: TransactionId) {
         let mut guard = self.live_read_transactions.lock().unwrap();
         let ref_count = guard.get_mut(&id).unwrap();
@@ -447,17 +455,26 @@ impl Database {
         }
     }
 
+    pub(crate) fn allocate_savepoint(&self) -> Result<(SavepointId, TransactionId)> {
+        let id = self.next_savepoint_id.next();
+        self.valid_savepoints.lock().unwrap().insert(id);
+        Ok((id, self.allocate_read_transaction()?))
+    }
+
+    pub(crate) fn is_valid_savepoint(&self, id: SavepointId) -> bool {
+        self.valid_savepoints.lock().unwrap().contains(&id)
+    }
+
+    pub(crate) fn invalidate_savepoints_after(&self, id: SavepointId) {
+        self.valid_savepoints.lock().unwrap().retain(|x| *x <= id);
+    }
+
     pub(crate) fn deallocate_savepoint(&self, savepoint: &Savepoint) {
         self.valid_savepoints
             .lock()
             .unwrap()
             .remove(&savepoint.get_id());
-        let mut guard = self.live_read_transactions.lock().unwrap();
-        let ref_count = guard.get_mut(&savepoint.get_transaction_id()).unwrap();
-        *ref_count -= 1;
-        if *ref_count == 0 {
-            guard.remove(&savepoint.get_transaction_id());
-        }
+        self.deallocate_read_transaction(savepoint.get_transaction_id());
     }
 
     pub(crate) fn oldest_live_read_transaction(&self) -> Option<TransactionId> {
@@ -505,59 +522,6 @@ impl Database {
         Ok(())
     }
 
-    /// Creates a snapshot of the current database state, which can be used to rollback the database
-    pub fn savepoint(&self) -> Result<Savepoint> {
-        let mut valid_savepoints_guard = self.valid_savepoints.lock().unwrap();
-        // Hold write lock to ensure we get a consistent snapshot
-        let guard = self.live_write_transaction.lock().unwrap();
-        assert!(guard.is_none());
-
-        let id = self.next_savepoint_id.next();
-        let mut live_reads = self.live_read_transactions.lock().unwrap();
-
-        let transaction_id = self.mem.get_last_committed_transaction_id()?;
-        live_reads
-            .entry(transaction_id)
-            .and_modify(|x| *x += 1)
-            .or_insert(1);
-
-        #[cfg(feature = "logging")]
-        info!(
-            "Creating savepoint id={:?}, txn_id={:?}",
-            id, transaction_id
-        );
-        let regional_allocators = self.mem.get_raw_allocator_states();
-        let root = self.mem.get_data_root();
-        let freed_root = self.mem.get_freed_root();
-        drop(live_reads);
-        let savepoint = Savepoint::new(
-            self,
-            id,
-            transaction_id,
-            root,
-            freed_root,
-            regional_allocators,
-        );
-        valid_savepoints_guard.insert(savepoint.get_id());
-        drop(guard);
-
-        Ok(savepoint)
-    }
-
-    pub fn restore_savepoint(&self, savepoint: &Savepoint) -> Result {
-        let mut guard = self.valid_savepoints.lock().unwrap();
-        if !guard.contains(&savepoint.get_id()) {
-            return Err(Error::InvalidSavepoint);
-        }
-        WriteTransaction::commit_savepoint(self, savepoint)?;
-
-        // Invalidate all savepoints that are newer than the one being applied to prevent the user
-        // from later trying to restore a savepoint on "another timeline"
-        guard.retain(|x| *x <= savepoint.get_id());
-
-        Ok(())
-    }
-
     /// Begins a write transaction
     ///
     /// Returns a [`WriteTransaction`] which may be used to read/write to the database. Only a single
@@ -575,9 +539,7 @@ impl Database {
     /// Returns a [`ReadTransaction`] which may be used to read from the database. Read transactions
     /// may exist concurrently with writes
     pub fn begin_read(&self) -> Result<ReadTransaction> {
-        let mut guard = self.live_read_transactions.lock().unwrap();
-        let id = self.mem.get_last_committed_transaction_id()?;
-        guard.entry(id).and_modify(|x| *x += 1).or_insert(1);
+        let id = self.allocate_read_transaction()?;
         #[cfg(feature = "logging")]
         info!("Beginning read transaction id={:?}", id);
         Ok(ReadTransaction::new(self, id))
diff --git a/src/transactions.rs b/src/transactions.rs
index 0de6d0b..33e0576 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -17,6 +17,7 @@ use std::mem::size_of;
 use std::ops::RangeFull;
 use std::panic;
 use std::rc::Rc;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::MutexGuard;
 
 /// Informational storage stats about the database
@@ -106,6 +107,7 @@ pub struct WriteTransaction<'db> {
     freed_pages: Rc<RefCell<Vec<PageNumber>>>,
     open_tables: RefCell<HashMap<String, &'static panic::Location<'static>>>,
     completed: bool,
+    dirty: AtomicBool,
     durability: Durability,
     live_write_transaction: MutexGuard<'db, Option<TransactionId>>,
 }
@@ -135,84 +137,118 @@ impl<'db> WriteTransaction<'db> {
             freed_pages,
             open_tables: RefCell::new(Default::default()),
             completed: false,
+            dirty: AtomicBool::new(false),
             durability: Durability::Immediate,
             live_write_transaction,
         })
     }
 
-    pub(crate) fn commit_savepoint(db: &'db Database, savepoint: &'_ Savepoint<'_>) -> Result {
-        let mut live_write_transaction = db.live_write_transaction.lock().unwrap();
-        assert!(live_write_transaction.is_none());
-        let transaction_id = db.increment_transaction_id();
+    /// Creates a snapshot of the current database state, which can be used to rollback the database
+    ///
+    /// Returns `[Error::InvalidSavepoint`], if the transaction is "dirty" (any tables have been openned)
+    pub fn savepoint(&self) -> Result<Savepoint<'db>> {
+        if self.dirty.load(Ordering::Acquire) {
+            return Err(Error::InvalidSavepoint);
+        }
+
+        let (id, transaction_id) = self.db.allocate_savepoint()?;
+        #[cfg(feature = "logging")]
+        info!(
+            "Creating savepoint id={:?}, txn_id={:?}",
+            id, transaction_id
+        );
+
+        let regional_allocators = self.mem.get_raw_allocator_states();
+        let root = self.mem.get_data_root();
+        let freed_root = self.mem.get_freed_root();
+        let savepoint = Savepoint::new(
+            self.db,
+            id,
+            transaction_id,
+            root,
+            freed_root,
+            regional_allocators,
+        );
+
+        Ok(savepoint)
+    }
+
+    /// Restore the state of the database to the given [`Savepoint`]
+    ///
+    /// Calling this method invalidates all [`Savepoint`]s created after savepoint
+    pub fn restore_savepoint(&mut self, savepoint: &Savepoint) -> Result {
+        if !self.db.is_valid_savepoint(savepoint.get_id()) {
+            return Err(Error::InvalidSavepoint);
+        }
         #[cfg(feature = "logging")]
         info!(
             "Beginning savepoint restore (id={:?}) in transaction id={:?}",
             savepoint.get_id(),
-            transaction_id
+            self.transaction_id
         );
-
         // Restoring a savepoint that reverted a file format or checksum type change could corrupt
         // the database
-        assert_eq!(db.get_memory().get_version(), savepoint.get_version());
+        assert_eq!(self.db.get_memory().get_version(), savepoint.get_version());
         assert_eq!(
-            db.get_memory().checksum_type(),
+            self.db.get_memory().checksum_type(),
             savepoint.get_checksum_type()
         );
+        self.dirty.store(true, Ordering::Release);
 
-        let freed_pages = db
-            .get_memory()
+        let allocated_since_savepoint = self
+            .mem
             .pages_allocated_since_raw_state(savepoint.get_regional_allocator_states());
-        let freed_pages = Rc::new(RefCell::new(freed_pages));
-        let root_page = savepoint.get_root();
-        let freed_root = savepoint.get_freed_root();
+        let mut freed_pages = vec![];
+        for page in allocated_since_savepoint {
+            if self.mem.uncommitted(page) {
+                // Safety: The page is uncommitted and we have a &mut on the transaction
+                unsafe {
+                    self.mem.free(page)?;
+                }
+            } else {
+                freed_pages.push(page);
+            }
+        }
+        *self.freed_pages.borrow_mut() = freed_pages;
+        self.table_tree = RefCell::new(TableTree::new(
+            savepoint.get_root(),
+            self.mem,
+            self.freed_pages.clone(),
+        ));
 
         // Remove any freed pages that have already been processed. Otherwise this would result in a double free
         // We assume below that PageNumber is length 8
-        let old_freed_tree: BtreeMut<FreedTableKey, [u8]> = BtreeMut::new(
-            db.get_memory().get_freed_root(),
-            db.get_memory(),
-            freed_pages.clone(),
+        let oldest_unprocessed_transaction = if let Some(entry) = self
+            .freed_tree
+            .range::<RangeFull, FreedTableKey>(..)?
+            .next()
+        {
+            FreedTableKey::from_bytes(entry.key()).transaction_id
+        } else {
+            self.transaction_id.0
+        };
+
+        self.freed_tree = BtreeMut::new(
+            savepoint.get_freed_root(),
+            self.mem,
+            self.freed_pages.clone(),
         );
-        let oldest_unprocessed_transaction =
-            if let Some(entry) = old_freed_tree.range::<RangeFull, FreedTableKey>(..)?.next() {
-                FreedTableKey::from_bytes(entry.key()).transaction_id
-            } else {
-                transaction_id.0
-            };
-        drop(old_freed_tree);
         let lookup_key = FreedTableKey {
             transaction_id: oldest_unprocessed_transaction,
             pagination_id: 0,
         };
-        let mut freed_tree = BtreeMut::new(freed_root, db.get_memory(), freed_pages.clone());
         let mut to_remove = vec![];
-        for entry in freed_tree.range(..lookup_key)? {
+        for entry in self.freed_tree.range(..lookup_key)? {
             to_remove.push(FreedTableKey::from_bytes(entry.key()));
         }
         for key in to_remove {
             // Safety: all references to the freed table above have already been dropped.
-            unsafe { freed_tree.remove(&key)? };
+            unsafe { self.freed_tree.remove(&key)? };
         }
 
-        *live_write_transaction = Some(transaction_id);
-        let transaction = Self {
-            db,
-            mem: db.get_memory(),
-            transaction_id,
-            table_tree: RefCell::new(TableTree::new(
-                root_page,
-                db.get_memory(),
-                freed_pages.clone(),
-            )),
-            freed_tree,
-            freed_pages,
-            open_tables: RefCell::new(Default::default()),
-            completed: false,
-            durability: Durability::Immediate,
-            live_write_transaction,
-        };
-
-        transaction.commit()?;
+        // Invalidate all savepoints that are newer than the one being applied to prevent the user
+        // from later trying to restore a savepoint "on another timeline"
+        self.db.invalidate_savepoints_after(savepoint.get_id());
 
         Ok(())
     }
@@ -238,6 +274,7 @@ impl<'db> WriteTransaction<'db> {
                 location,
             ));
         }
+        self.dirty.store(true, Ordering::Release);
         self.open_tables
             .borrow_mut()
             .insert(definition.name().to_string(), panic::Location::caller());
@@ -271,6 +308,7 @@ impl<'db> WriteTransaction<'db> {
                 location,
             ));
         }
+        self.dirty.store(true, Ordering::Release);
         self.open_tables
             .borrow_mut()
             .insert(definition.name().to_string(), panic::Location::caller());
@@ -309,6 +347,7 @@ impl<'db> WriteTransaction<'db> {
     ) -> Result<bool> {
         #[cfg(feature = "logging")]
         info!("Deleting table: {}", definition);
+        self.dirty.store(true, Ordering::Release);
         self.table_tree
             .borrow_mut()
             .delete_table::<K, V>(definition.name(), TableType::Normal)
@@ -323,6 +362,7 @@ impl<'db> WriteTransaction<'db> {
     ) -> Result<bool> {
         #[cfg(feature = "logging")]
         info!("Deleting multimap table: {}", definition);
+        self.dirty.store(true, Ordering::Release);
         self.table_tree
             .borrow_mut()
             .delete_table::<K, V>(definition.name(), TableType::Multimap)


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 94d0671..144a86b 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -877,9 +877,8 @@ fn regression18() {
 
     let table_def: TableDefinition<u64, [u8]> = TableDefinition::new("x");
 
-    let savepoint0 = db.savepoint().unwrap();
-
     let tx = db.begin_write().unwrap();
+    let savepoint0 = tx.savepoint().unwrap();
     {
         let mut t = tx.open_table(table_def).unwrap();
         let mut value = t.insert_reserve(&118749, 817).unwrap();
@@ -887,9 +886,8 @@ fn regression18() {
     }
     tx.commit().unwrap();
 
-    let savepoint1 = db.savepoint().unwrap();
-
     let tx = db.begin_write().unwrap();
+    let savepoint1 = tx.savepoint().unwrap();
     {
         let mut t = tx.open_table(table_def).unwrap();
         let mut value = t.insert_reserve(&65373, 1807).unwrap();
@@ -897,9 +895,11 @@ fn regression18() {
     }
     tx.commit().unwrap();
 
-    let savepoint2 = db.savepoint().unwrap();
+    let mut tx = db.begin_write().unwrap();
+    let savepoint2 = tx.savepoint().unwrap();
 
-    db.restore_savepoint(&savepoint2).unwrap();
+    tx.restore_savepoint(&savepoint2).unwrap();
+    tx.commit().unwrap();
 
     drop(savepoint0);
 
@@ -911,7 +911,9 @@ fn regression18() {
     }
     tx.commit().unwrap();
 
-    let savepoint4 = db.savepoint().unwrap();
+    let tx = db.begin_write().unwrap();
+    let savepoint4 = tx.savepoint().unwrap();
+    tx.abort().unwrap();
     drop(savepoint1);
 
     let tx = db.begin_write().unwrap();
@@ -925,12 +927,61 @@ fn regression18() {
     }
     tx.commit().unwrap();
 
-    db.restore_savepoint(&savepoint4).unwrap();
+    let mut tx = db.begin_write().unwrap();
+    tx.restore_savepoint(&savepoint4).unwrap();
+    tx.commit().unwrap();
 
     drop(savepoint2);
     drop(savepoint4);
 }
 
+#[test]
+fn regression19() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 1024 * 1024;
+    let db = unsafe {
+        Database::builder()
+            .set_write_strategy(WriteStrategy::Checksum)
+            .create(tmpfile.path(), db_size)
+            .unwrap()
+    };
+
+    let table_def: TableDefinition<u64, [u8]> = TableDefinition::new("x");
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let value = vec![0xFF; 100];
+        t.insert(&1, &value).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    let savepoint0 = tx.savepoint().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let value = vec![0xFF; 101];
+        t.insert(&1, &value).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let value = vec![0xFF; 102];
+        t.insert(&1, &value).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let mut tx = db.begin_write().unwrap();
+    tx.restore_savepoint(&savepoint0).unwrap();
+    tx.commit().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    tx.open_table(table_def).unwrap();
+}
+
 #[test]
 fn change_invalidate_savepoint() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
@@ -942,10 +993,14 @@ fn change_invalidate_savepoint() {
             .create(tmpfile.path(), db_size)
             .unwrap()
     };
-    let savepoint = db.savepoint().unwrap();
+    let tx = db.begin_write().unwrap();
+    let savepoint = tx.savepoint().unwrap();
+    tx.abort().unwrap();
     db.set_write_strategy(WriteStrategy::TwoPhase).unwrap();
+
+    let mut tx = db.begin_write().unwrap();
     assert!(matches!(
-        db.restore_savepoint(&savepoint),
+        tx.restore_savepoint(&savepoint),
         Err(Error::InvalidSavepoint)
     ));
 }
@@ -1314,28 +1369,31 @@ fn savepoint() {
     }
     txn.commit().unwrap();
 
-    let savepoint = db.savepoint().unwrap();
-
     let txn = db.begin_write().unwrap();
+    let savepoint = txn.savepoint().unwrap();
     {
         let mut table = txn.open_table(definition).unwrap();
         table.remove(&0).unwrap();
     }
     txn.commit().unwrap();
 
-    let savepoint2 = db.savepoint().unwrap();
+    let mut txn = db.begin_write().unwrap();
+    let savepoint2 = txn.savepoint().unwrap();
 
-    db.restore_savepoint(&savepoint).unwrap();
+    txn.restore_savepoint(&savepoint).unwrap();
 
     assert!(matches!(
-        db.restore_savepoint(&savepoint2).err().unwrap(),
+        txn.restore_savepoint(&savepoint2).err().unwrap(),
         Error::InvalidSavepoint
     ));
+    txn.commit().unwrap();
 
     let txn = db.begin_read().unwrap();
     let table = txn.open_table(definition).unwrap();
     assert_eq!(table.get(&0).unwrap().unwrap(), "hello");
 
     // Test that savepoints can be used multiple times
-    db.restore_savepoint(&savepoint).unwrap();
+    let mut txn = db.begin_write().unwrap();
+    txn.restore_savepoint(&savepoint).unwrap();
+    txn.commit().unwrap();
 }