From fdff4f18f251ce6ce56f03b31e4b31ddadf0b75a Mon Sep 17 00:00:00 2001 From: Hai | RISE <150876604+hai-rise@users.noreply.github.com> Date: Fri, 6 Dec 2024 20:58:17 +0700 Subject: [PATCH] feat(`DbTx`): add `get_by_encoded_key` (#13171) --- crates/storage/db-api/src/mock.rs | 9 +++- crates/storage/db-api/src/transaction.rs | 11 +++- crates/storage/db/Cargo.toml | 5 ++ crates/storage/db/benches/get.rs | 52 +++++++++++++++++++ .../storage/db/src/implementation/mdbx/tx.rs | 9 +++- testing/ef-tests/src/models.rs | 9 ++-- 6 files changed, 88 insertions(+), 7 deletions(-) create mode 100644 crates/storage/db/benches/get.rs diff --git a/crates/storage/db-api/src/mock.rs b/crates/storage/db-api/src/mock.rs index e972821d8..5580727fd 100644 --- a/crates/storage/db-api/src/mock.rs +++ b/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( + &self, + _key: &::Encoded, + ) -> Result, DatabaseError> { + Ok(None) + } + fn commit(self) -> Result { Ok(true) } diff --git a/crates/storage/db-api/src/transaction.rs b/crates/storage/db-api/src/transaction.rs index f39cf92fb..6dc79670a 100644 --- a/crates/storage/db-api/src/transaction.rs +++ b/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: DbDupCursorRO + DbCursorRO + Send + Sync; - /// Get value + /// Get value by an owned key fn get(&self, key: T::Key) -> Result, 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( + &self, + key: &::Encoded, + ) -> Result, DatabaseError>; /// Commit for read only transaction will consume and free transaction and allows /// freeing of memory pages fn commit(self) -> Result; diff --git a/crates/storage/db/Cargo.toml b/crates/storage/db/Cargo.toml index 4a4eff471..fd313a40a 100644 --- a/crates/storage/db/Cargo.toml +++ b/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 diff --git a/crates/storage/db/benches/get.rs b/crates/storage/db/benches/get.rs new file mode 100644 index 000000000..04eda02e0 --- /dev/null +++ b/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::(*key).unwrap(); + } + }) + }); + + group.bench_function("get_by_encoded_key", |b| { + b.iter(|| { + for key in &keys { + tx.get_by_encoded_key::(key).unwrap(); + } + }) + }); +} diff --git a/crates/storage/db/src/implementation/mdbx/tx.rs b/crates/storage/db/src/implementation/mdbx/tx.rs index 2ff2789ea..09be53a5f 100644 --- a/crates/storage/db/src/implementation/mdbx/tx.rs +++ b/crates/storage/db/src/implementation/mdbx/tx.rs @@ -283,8 +283,15 @@ impl DbTx for Tx { type DupCursor = Cursor; fn get(&self, key: T::Key) -> Result::Value>, DatabaseError> { + self.get_by_encoded_key::(&key.encode()) + } + + fn get_by_encoded_key( + &self, + key: &::Encoded, + ) -> Result, DatabaseError> { self.execute_with_operation_metric::(Operation::Get, None, |tx| { - tx.get(self.get_dbi::()?, key.encode().as_ref()) + tx.get(self.get_dbi::()?, key.as_ref()) .map_err(|e| DatabaseError::Read(e.into()))? .map(decode_one::) .transpose() diff --git a/testing/ef-tests/src/models.rs b/testing/ef-tests/src/models.rs index 160b0ec1d..7f6c0cdae 100644 --- a/testing/ef-tests/src/models.rs +++ b/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::(address)?.ok_or_else(|| { - Error::Assertion(format!("Expected account ({address}) is missing from DB: {self:?}")) - })?; + let account = + tx.get_by_encoded_key::(&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")?;