perf(engine): add StateRootTask bench

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",

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<EvmState> {
+    let mut rng = generators::rng();
+    let all_addresses: Vec<Address> = (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::<u64>());
+                storage.insert(
+                    slot,
+                    EvmStorageSlot::new_changed(U256::ZERO, U256::from(rng.gen::<u64>())),
+                );
+            }
+
+            let account = RevmAccount {
+                info: AccountInfo {
+                    balance: U256::from(rng.gen::<u64>()),
+                    nonce: rng.gen::<u64>(),
+                    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<MockNodeTypesWithDB>,
+    state_updates: &[EvmState],
+) -> Result<(), Box<dyn std::error::Error>> {
+    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
+                ),
+            ),
+            &params,
+            |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);

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.
 ///

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<StateRootResult>,
 }
@@ -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<Factory> {
+pub struct StateRootConfig<Factory> {
     /// View over the state in the database.
     pub consistent_view: ConsistentDbView<Factory>,
     /// Latest trie input.
@@ -67,7 +67,7 @@ pub(crate) struct StateRootConfig<Factory> {
 /// 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<Factory> {
+pub struct StateRootTask<Factory> {
     /// Task configuration.
     config: StateRootConfig<Factory>,
     /// 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<Factory>) -> Self {
+    pub fn new(config: StateRootConfig<Factory>) -> 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| {