feat(DbTx): add get_by_encoded_key (#13171)

crates/storage/db-api/src/mock.rs

@@ -7,7 +7,7 @@ use crate::{
         ReverseWalker, Walker,
     },
     database::Database,
-    table::{DupSort, Table, TableImporter},
+    table::{DupSort, Encode, Table, TableImporter},
     transaction::{DbTx, DbTxMut},
     DatabaseError,
 };
@@ -49,6 +49,13 @@ impl DbTx for TxMock {
         Ok(None)
     }
 
+    fn get_by_encoded_key<T: Table>(
+        &self,
+        _key: &<T::Key as Encode>::Encoded,
+    ) -> Result<Option<T::Value>, DatabaseError> {
+        Ok(None)
+    }
+
     fn commit(self) -> Result<bool, DatabaseError> {
         Ok(true)
     }

crates/storage/db-api/src/transaction.rs

@@ -1,6 +1,6 @@
 use crate::{
     cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW},
-    table::{DupSort, Table},
+    table::{DupSort, Encode, Table},
     DatabaseError,
 };
 
@@ -11,8 +11,15 @@ pub trait DbTx: Send + Sync {
     /// `DupCursor` type for this read-only transaction
     type DupCursor<T: DupSort>: DbDupCursorRO<T> + DbCursorRO<T> + Send + Sync;
 
-    /// Get value
+    /// Get value by an owned key
     fn get<T: Table>(&self, key: T::Key) -> Result<Option<T::Value>, DatabaseError>;
+    /// Get value by a reference to the encoded key, especially useful for "raw" keys
+    /// that encode to themselves like Address and B256. Doesn't need to clone a
+    /// reference key like `get`.
+    fn get_by_encoded_key<T: Table>(
+        &self,
+        key: &<T::Key as Encode>::Encoded,
+    ) -> Result<Option<T::Value>, DatabaseError>;
     /// Commit for read only transaction will consume and free transaction and allows
     /// freeing of memory pages
     fn commit(self) -> Result<bool, DatabaseError>;

crates/storage/db/Cargo.toml

@@ -129,3 +129,8 @@ harness = false
 name = "iai"
 required-features = ["test-utils"]
 harness = false
+
+[[bench]]
+name = "get"
+required-features = ["test-utils"]
+harness = false

crates/storage/db/benches/get.rs

@@ -0,0 +1,52 @@
+#![allow(missing_docs)]
+
+use alloy_primitives::TxHash;
+use criterion::{criterion_group, criterion_main, Criterion};
+use pprof::criterion::{Output, PProfProfiler};
+use reth_db::{test_utils::create_test_rw_db_with_path, Database, TransactionHashNumbers};
+use reth_db_api::transaction::DbTx;
+use std::{fs, sync::Arc};
+
+mod utils;
+use utils::BENCH_DB_PATH;
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().with_profiler(PProfProfiler::new(1, Output::Flamegraph(None)));
+    targets = get
+}
+criterion_main!(benches);
+
+// Small benchmark showing that [get_by_encoded_key] is slightly faster than [get]
+// for a reference key, as [get] requires copying or cloning the key first.
+fn get(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Get");
+
+    // Random keys to get
+    let mut keys = Vec::new();
+    for _ in 0..10_000_000 {
+        let key = TxHash::random();
+        keys.push(key);
+    }
+
+    // We don't bother mock the DB to reduce noise from DB I/O, value decoding, etc.
+    let _ = fs::remove_dir_all(BENCH_DB_PATH);
+    let db = Arc::try_unwrap(create_test_rw_db_with_path(BENCH_DB_PATH)).unwrap();
+    let tx = db.tx().expect("tx");
+
+    group.bench_function("get", |b| {
+        b.iter(|| {
+            for key in &keys {
+                tx.get::<TransactionHashNumbers>(*key).unwrap();
+            }
+        })
+    });
+
+    group.bench_function("get_by_encoded_key", |b| {
+        b.iter(|| {
+            for key in &keys {
+                tx.get_by_encoded_key::<TransactionHashNumbers>(key).unwrap();
+            }
+        })
+    });
+}

crates/storage/db/src/implementation/mdbx/tx.rs

@@ -283,8 +283,15 @@ impl<K: TransactionKind> DbTx for Tx<K> {
     type DupCursor<T: DupSort> = Cursor<K, T>;
 
     fn get<T: Table>(&self, key: T::Key) -> Result<Option<<T as Table>::Value>, DatabaseError> {
+        self.get_by_encoded_key::<T>(&key.encode())
+    }
+
+    fn get_by_encoded_key<T: Table>(
+        &self,
+        key: &<T::Key as Encode>::Encoded,
+    ) -> Result<Option<T::Value>, DatabaseError> {
         self.execute_with_operation_metric::<T, _>(Operation::Get, None, |tx| {
-            tx.get(self.get_dbi::<T>()?, key.encode().as_ref())
+            tx.get(self.get_dbi::<T>()?, key.as_ref())
                 .map_err(|e| DatabaseError::Read(e.into()))?
                 .map(decode_one::<T>)
                 .transpose()

testing/ef-tests/src/models.rs

@@ -219,9 +219,12 @@ impl Account {
     ///
     /// In case of a mismatch, `Err(Error::Assertion)` is returned.
     pub fn assert_db(&self, address: Address, tx: &impl DbTx) -> Result<(), Error> {
-        let account = tx.get::<tables::PlainAccountState>(address)?.ok_or_else(|| {
-            Error::Assertion(format!("Expected account ({address}) is missing from DB: {self:?}"))
-        })?;
+        let account =
+            tx.get_by_encoded_key::<tables::PlainAccountState>(&address)?.ok_or_else(|| {
+                Error::Assertion(format!(
+                    "Expected account ({address}) is missing from DB: {self:?}"
+                ))
+            })?;
 
         assert_equal(self.balance, account.balance, "Balance does not match")?;
         assert_equal(self.nonce.to(), account.nonce, "Nonce does not match")?;