From 3c132958d1932fd3038ee27fd66cfe39d75a078d Mon Sep 17 00:00:00 2001 From: Federico Gimenez Date: Mon, 9 Dec 2024 20:52:58 +0100 Subject: [PATCH] perf(engine): add StateRootTask bench (#13212) --- crates/engine/tree/Cargo.toml | 4 + crates/engine/tree/benches/state_root_task.rs | 166 ++++++++++++++++++ crates/engine/tree/src/tree/mod.rs | 2 +- crates/engine/tree/src/tree/root.rs | 16 +- 4 files changed, 179 insertions(+), 9 deletions(-) create mode 100644 crates/engine/tree/benches/state_root_task.rs diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index 680b6933e..67cb72850 100644 --- a/crates/engine/tree/Cargo.toml +++ b/crates/engine/tree/Cargo.toml @@ -93,6 +93,10 @@ rand.workspace = true name = "channel_perf" harness = false +[[bench]] +name = "state_root_task" +harness = false + [features] test-utils = [ "reth-blockchain-tree/test-utils", diff --git a/crates/engine/tree/benches/state_root_task.rs b/crates/engine/tree/benches/state_root_task.rs new file mode 100644 index 000000000..391fd333d --- /dev/null +++ b/crates/engine/tree/benches/state_root_task.rs @@ -0,0 +1,166 @@ +//! Benchmark for `StateRootTask` complete workflow, including sending state +//! updates using the incoming messages sender and waiting for the final result. + +#![allow(missing_docs)] + +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; +use reth_engine_tree::tree::root::{StateRootConfig, StateRootTask}; +use reth_evm::system_calls::OnStateHook; +use reth_primitives::{Account as RethAccount, StorageEntry}; +use reth_provider::{ + providers::ConsistentDbView, + test_utils::{create_test_provider_factory, MockNodeTypesWithDB}, + HashingWriter, ProviderFactory, +}; +use reth_testing_utils::generators::{self, Rng}; +use reth_trie::TrieInput; +use revm_primitives::{ + Account as RevmAccount, AccountInfo, AccountStatus, Address, EvmState, EvmStorageSlot, HashMap, + B256, KECCAK_EMPTY, U256, +}; +use std::sync::Arc; + +#[derive(Debug, Clone)] +struct BenchParams { + num_accounts: usize, + updates_per_account: usize, + storage_slots_per_account: usize, +} + +fn create_bench_state_updates(params: &BenchParams) -> Vec { + let mut rng = generators::rng(); + let all_addresses: Vec
= (0..params.num_accounts).map(|_| rng.gen()).collect(); + let mut updates = Vec::new(); + + for _ in 0..params.updates_per_account { + let num_accounts_in_update = rng.gen_range(1..=params.num_accounts); + let mut state_update = EvmState::default(); + + let selected_addresses = &all_addresses[0..num_accounts_in_update]; + + for &address in selected_addresses { + let mut storage = HashMap::default(); + for _ in 0..params.storage_slots_per_account { + let slot = U256::from(rng.gen::()); + storage.insert( + slot, + EvmStorageSlot::new_changed(U256::ZERO, U256::from(rng.gen::())), + ); + } + + let account = RevmAccount { + info: AccountInfo { + balance: U256::from(rng.gen::()), + nonce: rng.gen::(), + code_hash: KECCAK_EMPTY, + code: Some(Default::default()), + }, + storage, + status: AccountStatus::Touched, + }; + + state_update.insert(address, account); + } + + updates.push(state_update); + } + + updates +} + +fn convert_revm_to_reth_account(revm_account: &RevmAccount) -> RethAccount { + RethAccount { + balance: revm_account.info.balance, + nonce: revm_account.info.nonce, + bytecode_hash: if revm_account.info.code_hash == KECCAK_EMPTY { + None + } else { + Some(revm_account.info.code_hash) + }, + } +} + +fn setup_provider( + factory: &ProviderFactory, + state_updates: &[EvmState], +) -> Result<(), Box> { + let provider_rw = factory.provider_rw()?; + + for update in state_updates { + let account_updates = update + .iter() + .map(|(address, account)| (*address, Some(convert_revm_to_reth_account(account)))); + provider_rw.insert_account_for_hashing(account_updates)?; + + let storage_updates = update.iter().map(|(address, account)| { + let storage_entries = account.storage.iter().map(|(slot, value)| StorageEntry { + key: B256::from(*slot), + value: value.present_value, + }); + (*address, storage_entries) + }); + provider_rw.insert_storage_for_hashing(storage_updates)?; + } + + provider_rw.commit()?; + Ok(()) +} + +fn bench_state_root(c: &mut Criterion) { + let mut group = c.benchmark_group("state_root"); + + let scenarios = vec![ + BenchParams { num_accounts: 100, updates_per_account: 5, storage_slots_per_account: 10 }, + BenchParams { num_accounts: 1000, updates_per_account: 10, storage_slots_per_account: 20 }, + ]; + + for params in scenarios { + group.bench_with_input( + BenchmarkId::new( + "state_root_task", + format!( + "accounts_{}_updates_{}_slots_{}", + params.num_accounts, + params.updates_per_account, + params.storage_slots_per_account + ), + ), + ¶ms, + |b, params| { + b.iter_with_setup( + || { + let factory = create_test_provider_factory(); + let state_updates = create_bench_state_updates(params); + setup_provider(&factory, &state_updates).expect("failed to setup provider"); + + let trie_input = Arc::new(TrieInput::from_state(Default::default())); + + let config = StateRootConfig { + consistent_view: ConsistentDbView::new(factory, None), + input: trie_input, + }; + + (config, state_updates) + }, + |(config, state_updates)| { + let task = StateRootTask::new(config); + let mut hook = task.state_hook(); + let handle = task.spawn(); + + for update in state_updates { + hook.on_state(&update) + } + drop(hook); + + black_box(handle.wait_for_result().expect("task failed")); + }, + ) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_state_root); +criterion_main!(benches); diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 763d5d990..5fc07abf7 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -76,7 +76,7 @@ pub use invalid_block_hook::{InvalidBlockHooks, NoopInvalidBlockHook}; pub use persistence_state::PersistenceState; pub use reth_engine_primitives::InvalidBlockHook; -mod root; +pub mod root; /// Keeps track of the state of the tree. /// diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index ae22b036b..53a881387 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/crates/engine/tree/src/tree/root.rs @@ -37,7 +37,7 @@ pub(crate) type StateRootResult = Result<(B256, TrieUpdates), ParallelStateRootE /// Handle to a spawned state root task. #[derive(Debug)] #[allow(dead_code)] -pub(crate) struct StateRootHandle { +pub struct StateRootHandle { /// Channel for receiving the final result. rx: mpsc::Receiver, } @@ -50,14 +50,14 @@ impl StateRootHandle { } /// Waits for the state root calculation to complete. - pub(crate) fn wait_for_result(self) -> StateRootResult { + pub fn wait_for_result(self) -> StateRootResult { self.rx.recv().expect("state root task was dropped without sending result") } } /// Common configuration for state root tasks #[derive(Debug)] -pub(crate) struct StateRootConfig { +pub struct StateRootConfig { /// View over the state in the database. pub consistent_view: ConsistentDbView, /// Latest trie input. @@ -67,7 +67,7 @@ pub(crate) struct StateRootConfig { /// Messages used internally by the state root task #[derive(Debug)] #[allow(dead_code)] -pub(crate) enum StateRootMessage { +pub enum StateRootMessage { /// New state update from transaction execution StateUpdate(EvmState), /// Proof calculation completed for a specific state update @@ -223,7 +223,7 @@ fn evm_state_to_hashed_post_state(update: EvmState) -> HashedPostState { /// to the tree. /// Then it updates relevant leaves according to the result of the transaction. #[derive(Debug)] -pub(crate) struct StateRootTask { +pub struct StateRootTask { /// Task configuration. config: StateRootConfig, /// Receiver for state root related messages. @@ -250,7 +250,7 @@ where + 'static, { /// Creates a new state root task with the unified message channel - pub(crate) fn new(config: StateRootConfig) -> Self { + pub fn new(config: StateRootConfig) -> Self { let (tx, rx) = channel(); Self { @@ -264,7 +264,7 @@ where } /// Spawns the state root task and returns a handle to await its result. - pub(crate) fn spawn(self) -> StateRootHandle { + pub fn spawn(self) -> StateRootHandle { let (tx, rx) = mpsc::sync_channel(1); std::thread::Builder::new() .name("State Root Task".to_string()) @@ -279,7 +279,7 @@ where } /// Returns a state hook to be used to send state updates to this task. - pub(crate) fn state_hook(&self) -> impl OnStateHook { + pub fn state_hook(&self) -> impl OnStateHook { let state_hook = StateHookSender::new(self.tx.clone()); move |state: &EvmState| {