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,19 +195,12 @@ 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<&Self, RO>, TransactionError> {
+        begin_ro(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 begin_rw(&self) -> Result<Transaction<&Self, RW>, TransactionError> {
+        begin_rw(self)
     }
 
     pub fn state_root(&self) -> B256 {
@@ -244,13 +238,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,13 +328,13 @@ 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();
@@ -512,4 +525,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
     }