refactor: Hide Db/Tx lifetimes via Deref/Generics

Author: BrianBland

src/database.rs

@@ -11,6 +11,7 @@ use parking_lot::Mutex;
 use std::{
     fs::File,
     io,
+    ops::Deref,
     path::{Path, PathBuf},
 };
 
@@ -194,21 +195,6 @@ impl Database {
         Ok(())
     }
 
-    pub fn begin_rw(&self) -> Result<Transaction<'_, RW>, TransactionError> {
-        let context = self.storage_engine.write_context();
-        let min_snapshot_id = self.transaction_manager.lock().begin_rw(context.snapshot_id)?;
-        if min_snapshot_id > 0 {
-            self.storage_engine.unlock(min_snapshot_id - 1);
-        }
-        Ok(Transaction::new(context, self))
-    }
-
-    pub fn begin_ro(&self) -> Result<Transaction<'_, RO>, TransactionError> {
-        let context = self.storage_engine.read_context();
-        self.transaction_manager.lock().begin_ro(context.snapshot_id);
-        Ok(Transaction::new(context, self))
-    }
-
     pub fn state_root(&self) -> B256 {
         self.storage_engine.read_context().root_node_hash
     }
@@ -244,13 +230,32 @@ impl Database {
     }
 }
 
+pub fn begin_ro<DB: Deref<Target = Database>>(
+    db: DB,
+) -> Result<Transaction<DB, RO>, TransactionError> {
+    let context = db.storage_engine.read_context();
+    db.transaction_manager.lock().begin_ro(context.snapshot_id);
+    Ok(Transaction::new(context, db))
+}
+
+pub fn begin_rw<DB: Deref<Target = Database>>(
+    db: DB,
+) -> Result<Transaction<DB, RW>, TransactionError> {
+    let context = db.storage_engine.write_context();
+    let min_snapshot_id = db.transaction_manager.lock().begin_rw(context.snapshot_id)?;
+    if min_snapshot_id > 0 {
+        db.storage_engine.unlock(min_snapshot_id - 1);
+    }
+    Ok(Transaction::new(context, db))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::{account::Account, path::AddressPath};
     use alloy_primitives::{address, Address, U256};
     use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY};
-    use std::fs;
+    use std::{fs, sync::Arc};
     use tempdir::TempDir;
 
     #[test]
@@ -315,16 +320,16 @@ mod tests {
         let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
 
         let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY);
-        let mut tx = db.begin_rw().unwrap();
+        let mut tx = begin_rw(&db).unwrap();
         tx.set_account(AddressPath::for_address(address), Some(account1.clone())).unwrap();
 
         tx.commit().unwrap();
 
         let account2 = Account::new(456, U256::from(123), EMPTY_ROOT_HASH, KECCAK_EMPTY);
-        let mut tx = db.begin_rw().unwrap();
+        let mut tx = begin_rw(&db).unwrap();
         tx.set_account(AddressPath::for_address(address), Some(account2.clone())).unwrap();
 
-        let mut ro_tx = db.begin_ro().unwrap();
+        let mut ro_tx = begin_ro(&db).unwrap();
         tx.commit().unwrap();
 
         // The read transaction was created before the write was committed, so it should not see the
@@ -334,7 +339,7 @@ mod tests {
         assert_eq!(account1, read_account.unwrap());
 
         // The writer transaction is committed, so the read transaction should see the changes.
-        let mut ro_tx = db.begin_ro().unwrap();
+        let mut ro_tx = begin_ro(&db).unwrap();
 
         let read_account = ro_tx.get_account(AddressPath::for_address(address)).unwrap();
 
@@ -372,29 +377,29 @@ mod tests {
         let address1 = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
         let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY);
 
-        let mut tx = db.begin_rw().unwrap();
+        let mut tx = begin_rw(&db).unwrap();
         tx.set_account(AddressPath::for_address(address1), Some(account1.clone())).unwrap();
 
         tx.commit().unwrap();
         db.close().unwrap();
 
         let db = Database::open(&file_path).unwrap();
-        let mut tx = db.begin_ro().unwrap();
+        let mut tx = begin_ro(&db).unwrap();
         let account = tx.get_account(AddressPath::for_address(address1)).unwrap().unwrap();
         assert_eq!(account, account1);
 
         tx.commit().unwrap();
 
         let address2 = address!("0x1234567890abcdef1234567890abcdef12345678");
         let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY);
-        let mut tx = db.begin_rw().unwrap();
+        let mut tx = begin_rw(&db).unwrap();
         tx.set_account(AddressPath::for_address(address2), Some(account2.clone())).unwrap();
 
         tx.commit().unwrap();
         db.close().unwrap();
 
         let db = Database::open(&file_path).unwrap();
-        let mut tx = db.begin_ro().unwrap();
+        let mut tx = begin_ro(&db).unwrap();
 
         let account = tx.get_account(AddressPath::for_address(address1)).unwrap().unwrap();
         assert_eq!(account, account1);
@@ -434,7 +439,7 @@ mod tests {
         assert_eq!(db.storage_engine.page_manager.size(), 0);
 
         // Add 1000 accounts
-        let mut tx = db.begin_rw().expect("rw transaction creation failed");
+        let mut tx = begin_rw(&db).expect("rw transaction creation failed");
         let initial_accounts = random_accounts(1000).collect::<Vec<_>>();
         for (address, account) in &initial_accounts {
             tx.set_account(address.clone(), account.clone()).expect("adding account failed");
@@ -456,7 +461,7 @@ mod tests {
         }
 
         // Add 1000 more accounts
-        let mut tx = db.begin_rw().expect("rw transaction creation failed");
+        let mut tx = begin_rw(&db).expect("rw transaction creation failed");
         let more_accounts = random_accounts(1000).collect::<Vec<_>>();
         for (address, account) in &more_accounts {
             tx.set_account(address.clone(), account.clone()).expect("adding account failed");
@@ -487,7 +492,7 @@ mod tests {
         }
 
         // Obtain a read transaction, and verify that it can access all the initial accounts
-        let mut read_tx = db.begin_ro().expect("ro transaction creation failed");
+        let mut read_tx = begin_ro(&db).expect("ro transaction creation failed");
         for (address, account) in &initial_accounts {
             assert_eq!(
                 read_tx.get_account(address.clone()).expect("error while reading account"),
@@ -496,7 +501,7 @@ mod tests {
         }
 
         // Delete the initial accounts
-        let mut tx = db.begin_rw().expect("rw transaction creation failed");
+        let mut tx = begin_rw(&db).expect("rw transaction creation failed");
         for (address, _) in &initial_accounts {
             tx.set_account(address.clone(), None).expect("deleting account failed");
         }
@@ -512,4 +517,26 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_db_arc_tx() {
+        let tmp_dir = TempDir::new("test_db").unwrap();
+        let file_path = tmp_dir.path().join("test.db");
+        let db = Database::create_new(&file_path).unwrap();
+
+        let db_arc = Arc::new(db);
+
+        let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045");
+        let mut tx = begin_rw(db_arc.clone()).unwrap();
+        tx.set_account(
+            AddressPath::for_address(address),
+            Some(Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY)),
+        )
+        .unwrap();
+        tx.commit().unwrap();
+
+        let mut tx = begin_ro(db_arc).unwrap();
+        let account = tx.get_account(AddressPath::for_address(address)).unwrap().unwrap();
+        assert_eq!(account, Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY));
+    }
 }

src/transaction.rs

@@ -14,7 +14,7 @@ use alloy_trie::Nibbles;
 pub use error::TransactionError;
 pub use manager::TransactionManager;
 use sealed::sealed;
-use std::{collections::HashMap, fmt::Debug};
+use std::{collections::HashMap, fmt::Debug, ops::Deref, sync::Arc};
 
 #[sealed]
 pub trait TransactionKind: Debug {}
@@ -34,21 +34,23 @@ impl TransactionKind for RO {}
 // Compile-time assertion to ensure that `Transaction` is `Send`
 const _: fn() = || {
     fn consumer<T: Send>() {}
-    consumer::<Transaction<'_, RO>>();
-    consumer::<Transaction<'_, RW>>();
+    consumer::<Transaction<&Database, RO>>();
+    consumer::<Transaction<&Database, RW>>();
+    consumer::<Transaction<Arc<Database>, RO>>();
+    consumer::<Transaction<Arc<Database>, RW>>();
 };
 
 #[derive(Debug)]
-pub struct Transaction<'tx, K: TransactionKind> {
+pub struct Transaction<DB, K: TransactionKind> {
     committed: bool,
     context: TransactionContext,
-    database: &'tx Database,
+    database: DB,
     pending_changes: HashMap<Nibbles, Option<TrieValue>>,
     _marker: std::marker::PhantomData<K>,
 }
 
-impl<'tx, K: TransactionKind> Transaction<'tx, K> {
-    pub(crate) fn new(context: TransactionContext, database: &'tx Database) -> Self {
+impl<DB: Deref<Target = Database>, K: TransactionKind> Transaction<DB, K> {
+    pub(crate) fn new(context: TransactionContext, database: DB) -> Self {
         Self {
             committed: false,
             context,
@@ -137,7 +139,7 @@ impl<'tx, K: TransactionKind> Transaction<'tx, K> {
     }
 }
 
-impl Transaction<'_, RW> {
+impl<DB: Deref<Target = Database>> Transaction<DB, RW> {
     pub fn set_account(
         &mut self,
         address_path: AddressPath,
@@ -186,7 +188,7 @@ impl Transaction<'_, RW> {
     }
 }
 
-impl Transaction<'_, RO> {
+impl<DB: Deref<Target = Database>> Transaction<DB, RO> {
     pub fn commit(mut self) -> Result<(), TransactionError> {
         let mut transaction_manager = self.database.transaction_manager.lock();
         transaction_manager.remove_tx(self.context.snapshot_id, false);
@@ -196,7 +198,7 @@ impl Transaction<'_, RO> {
     }
 }
 
-impl<K: TransactionKind> Drop for Transaction<'_, K> {
+impl<DB, K: TransactionKind> Drop for Transaction<DB, K> {
     fn drop(&mut self) {
         // TODO: panic if the transaction is not committed
     }