Root hash thread pool, new trie benches.

crates/eth-sparse-mpt/Cargo.toml

@@ -7,6 +7,7 @@ license.workspace = true
 homepage.workspace = true
 repository.workspace = true
 
+
 [dependencies]
 thiserror = "1.0.61"
 serde = { version = "1.0.203", features = ["derive"] }
@@ -33,15 +34,21 @@ alloy-primitives.workspace = true
 alloy-rlp.workspace = true
 alloy-trie.workspace = true
 
-# test only dependencies but included here to be accessible from benches/
-hash-db = "0.15.2"
-triehash = "0.8.4"
+
+# test deps
+hash-db = { version = "0.15.2", optional = true }
+triehash = { version = "0.8.4", optional = true }
+flate2 = { version = "1.0.35", optional = true }
+eyre = { workspace = true, optional = true}
+
+[features]
+benchmark-utils = ["dep:hash-db", "dep:triehash", "dep:flate2", "dep:eyre"]
 
 [dev-dependencies]
 criterion = { version = "0.4", features = ["html_reports"] }
-eyre = "0.6.12"
 rand = { version = "0.8.5", features = ["small_rng"] }
 proptest = "1.5.0"
+eth-sparse-mpt = { path = ".", features = ["benchmark-utils"] }
 
 [[bench]]
 name = "trie_insert_bench"
@@ -51,3 +58,7 @@ harness = false
 name = "trie_nodes_benches"
 harness = false
 
+[[bench]]
+name = "trie_example_bench"
+harness = false
+

crates/eth-sparse-mpt/benches/trie_example_bench.rs

@@ -0,0 +1,85 @@
+use std::path::PathBuf;
+
+use eth_sparse_mpt::{
+    reth_sparse_trie::{change_set::ETHTrieChangeSet, hash::EthSparseTries},
+    test_utils::deserialize_from_json_gzip,
+    RootHashThreadPool,
+};
+use std::time::Instant;
+
+fn main() {
+    let mut examples = Vec::new();
+
+    for i in 0..7 {
+        let dir: PathBuf = format!("./test_data/prepared_tries/example{}/", i)
+            .parse()
+            .unwrap();
+
+        let change_set: ETHTrieChangeSet = {
+            let mut p = dir.clone();
+            p.push("change_set.json.gz");
+            deserialize_from_json_gzip(p).expect("change set")
+        };
+
+        let tries: EthSparseTries = {
+            let mut p = dir.clone();
+            p.push("tries.json.gz");
+            deserialize_from_json_gzip(p).expect("sparse trie")
+        };
+
+        examples.push((change_set, tries));
+    }
+
+    const WARMUP_RUNS: usize = 100;
+    const REAL_RUNS: usize = 1000;
+
+    const PAR_ACCOUNT_TRIE: bool = true;
+    const PAR_STORAGE_TRIES: bool = true;
+
+    let threadpool = RootHashThreadPool::try_new(4).unwrap();
+
+    println!("example,min,max,p50,p99,MAX/MIN,p99/p50");
+    for (i, (change_set, tries)) in examples.into_iter().enumerate() {
+        let mut measures = Vec::new();
+
+        for _ in 0..WARMUP_RUNS {
+            let (change_set, mut tries) = (change_set.clone(), tries.clone());
+
+            threadpool.rayon_pool.install(|| {
+                tries
+                    .calculate_root_hash(change_set, PAR_STORAGE_TRIES, PAR_ACCOUNT_TRIE)
+                    .unwrap();
+            });
+        }
+
+        for _ in 0..REAL_RUNS {
+            let (change_set, mut tries) = (change_set.clone(), tries.clone());
+
+            let start = Instant::now();
+
+            threadpool.rayon_pool.install(|| {
+                tries
+                    .calculate_root_hash(change_set, PAR_STORAGE_TRIES, PAR_ACCOUNT_TRIE)
+                    .unwrap();
+            });
+            measures.push(start.elapsed().as_micros());
+        }
+
+        measures.sort();
+        let min = *measures.first().unwrap();
+        let max = *measures.last().unwrap();
+        let p50 = measures[measures.len() / 2];
+        let p99 = measures[measures.len() * 99 / 100];
+
+        println!(
+            "{},{},{},{},{},{:.2},{:.2}",
+            i,
+            min,
+            max,
+            p50,
+            p99,
+            max as f64 / min as f64,
+            p99 as f64 / p50 as f64
+        );
+    }
+}

crates/eth-sparse-mpt/benches/trie_insert_bench.rs

@@ -6,7 +6,7 @@ use eth_sparse_mpt::{
         shared_cache::RethSparseTrieShareCacheInternal, SparseTrieSharedCache,
     },
     sparse_mpt::{DiffTrie, FixedTrie},
-    utils::{get_test_change_set, get_test_multiproofs},
+    test_utils::{get_test_change_set, get_test_multiproofs},
 };
 
 fn get_storage_tries(changes: &ETHTrieChangeSet, tries: &EthSparseTries) -> Vec<DiffTrie> {

crates/eth-sparse-mpt/benches/trie_nodes_benches.rs

@@ -1,10 +1,7 @@
 use alloy_primitives::{keccak256, Bytes, B256, U256};
 use alloy_rlp::Encodable;
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use eth_sparse_mpt::{
-    sparse_mpt::DiffTrie,
-    utils::{HashMap, KeccakHasher},
-};
+use eth_sparse_mpt::{sparse_mpt::DiffTrie, test_utils::KeccakHasher, utils::HashMap};
 
 // hashing this trie it roughly equivalent to updating the trie for the block
 const TRIE_SIZE: usize = 3000;

crates/eth-sparse-mpt/src/lib.rs

@@ -10,9 +10,11 @@
 
 pub mod reth_sparse_trie;
 pub mod sparse_mpt;
+#[cfg(any(test, feature = "benchmark-utils"))]
+pub mod test_utils;
 pub mod utils;
 
 pub use reth_sparse_trie::{
     calculate_root_hash_with_sparse_trie, prefetch_tries_for_accounts, ChangedAccountData,
-    SparseTrieSharedCache,
+    RootHashThreadPool, SparseTrieSharedCache,
 };

crates/eth-sparse-mpt/src/reth_sparse_trie/hash.rs

@@ -7,10 +7,14 @@ use alloy_primitives::{Bytes, B256};
 use alloy_rlp::Encodable;
 use rayon::prelude::*;
 use reth_trie::TrieAccount;
+use serde::{Deserialize, Serialize};
+use serde_with::{serde_as, Seq};
 
-#[derive(Default, Clone)]
+#[serde_as]
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct EthSparseTries {
     pub account_trie: DiffTrie,
+    #[serde_as(as = "Seq<(_, _)>")]
     pub storage_tries: HashMap<Bytes, DiffTrie>,
 }
 

crates/eth-sparse-mpt/src/reth_sparse_trie/mod.rs

@@ -1,10 +1,12 @@
 use alloy_primitives::{Address, B256};
 use change_set::{prepare_change_set, prepare_change_set_for_prefetch};
 use hash::RootHashError;
+use rayon::{ThreadPoolBuildError, ThreadPoolBuilder};
 use reth_provider::{
     providers::ConsistentDbView, BlockReader, DatabaseProviderFactory, ExecutionOutcome,
     StateCommitmentProvider,
 };
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 pub mod change_set;
@@ -90,13 +92,56 @@ where
     Err(SparseTrieError::FailedToFetchData)
 }
 
+#[derive(Clone, Debug)]
+pub struct RootHashThreadPool {
+    pub rayon_pool: Arc<rayon::ThreadPool>,
+}
+
+impl RootHashThreadPool {
+    pub fn try_new(threads: usize) -> Result<RootHashThreadPool, ThreadPoolBuildError> {
+        let rayon_pool = ThreadPoolBuilder::new()
+            .num_threads(threads)
+            .thread_name(|idx| format!("sparse_mpt:{}", idx))
+            .build()?;
+        Ok(RootHashThreadPool {
+            rayon_pool: Arc::new(rayon_pool),
+        })
+    }
+}
+
+impl Default for RootHashThreadPool {
+    fn default() -> Self {
+        let cpus = rayon::current_num_threads();
+        Self::try_new(cpus).expect("failed to create default root hash threadpool")
+    }
+}
+
 /// Calculate root hash for the given outcome on top of the block defined by consistent_db_view.
 /// * shared_cache should be created once for each parent block and it stores fetched parts of the trie
 /// * It uses rayon for parallelism and the thread pool should be configured from outside.
 pub fn calculate_root_hash_with_sparse_trie<Provider>(
     consistent_db_view: ConsistentDbView<Provider>,
     outcome: &ExecutionOutcome,
     shared_cache: SparseTrieSharedCache,
+    thread_pool: &Option<RootHashThreadPool>,
+) -> (Result<B256, SparseTrieError>, SparseTrieMetrics)
+where
+    Provider: DatabaseProviderFactory<Provider: BlockReader> + Send + Sync,
+    Provider: StateCommitmentProvider,
+{
+    if let Some(thread_pool) = thread_pool {
+        thread_pool.rayon_pool.install(|| {
+            calculate_root_hash_with_sparse_trie_internal(consistent_db_view, outcome, shared_cache)
+        })
+    } else {
+        calculate_root_hash_with_sparse_trie_internal(consistent_db_view, outcome, shared_cache)
+    }
+}
+
+fn calculate_root_hash_with_sparse_trie_internal<Provider>(
+    consistent_db_view: ConsistentDbView<Provider>,
+    outcome: &ExecutionOutcome,
+    shared_cache: SparseTrieSharedCache,
 ) -> (Result<B256, SparseTrieError>, SparseTrieMetrics)
 where
     Provider: DatabaseProviderFactory<Provider: BlockReader> + Send + Sync,

crates/eth-sparse-mpt/src/sparse_mpt/diff_trie/tests.rs

@@ -1,7 +1,8 @@
 use super::*;
 use crate::{
     sparse_mpt::*,
-    utils::{reference_trie_hash, HashSet, StoredFailureCase},
+    test_utils::{reference_trie_hash, StoredFailureCase},
+    utils::HashSet,
 };
 use alloy_primitives::{Bytes, B256};
 use eyre::Context;
@@ -493,7 +494,7 @@ fn check_correct_gather_for_orphan_of_and_orphan() {
 
 #[test]
 fn known_failure_case_0() {
-    let input = StoredFailureCase::load("./test_data/failure_case_0.json");
+    let input = StoredFailureCase::load("./test_data/failure_case_0.json.gz");
 
     let mut prev_value = None;
     for i in 0..10 {

crates/eth-sparse-mpt/src/sparse_mpt/fixed_trie.rs

@@ -462,7 +462,7 @@ fn get_child_ptr(child_ptrs: &[(u8, u64)], nibble: u8) -> Option<u64> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::utils::{get_test_change_set, get_test_multiproofs};
+    use crate::test_utils::{get_test_change_set, get_test_multiproofs};
 
     #[test]
     fn test_insert_and_gather_account_trie() {

crates/eth-sparse-mpt/src/test_utils.rs

@@ -0,0 +1,66 @@
+use std::{fs::File, io::Read, path::Path};
+
+use alloy_primitives::{keccak256, Bytes, B256};
+use flate2::read::GzDecoder;
+use rustc_hash::FxHasher;
+use serde::{de::DeserializeOwned, Deserialize, Serialize};
+
+use crate::{
+    reth_sparse_trie::{change_set::ETHTrieChangeSet, trie_fetcher::MultiProof},
+    sparse_mpt::DiffTrie,
+};
+
+pub fn deserialize_from_json_gzip<T: DeserializeOwned>(path: impl AsRef<Path>) -> eyre::Result<T> {
+    let file = File::open(path)?;
+    let mut gz = GzDecoder::new(file);
+    let mut content = String::new();
+    gz.read_to_string(&mut content)?;
+    Ok(serde_json::from_str(&content)?)
+}
+
+#[derive(Debug)]
+pub struct KeccakHasher {}
+
+impl hash_db::Hasher for KeccakHasher {
+    type Out = B256;
+    type StdHasher = FxHasher;
+    const LENGTH: usize = 32;
+
+    fn hash(x: &[u8]) -> Self::Out {
+        keccak256(x)
+    }
+}
+
+pub fn reference_trie_hash(data: &[(Bytes, Bytes)]) -> B256 {
+    triehash::trie_root::<KeccakHasher, _, _, _>(data.to_vec())
+}
+
+pub fn get_test_multiproofs() -> Vec<MultiProof> {
+    let files = [
+        "./test_data/multiproof_0.json.gz",
+        "./test_data/multiproof_1.json.gz",
+    ];
+    let mut result = Vec::new();
+    for file in files {
+        result.push(deserialize_from_json_gzip(file).expect("parsing multiproof"));
+    }
+    result
+}
+
+pub fn get_test_change_set() -> ETHTrieChangeSet {
+    deserialize_from_json_gzip("./test_data/changeset.json.gz").expect("changeset")
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StoredFailureCase {
+    pub trie: DiffTrie,
+    pub updated_keys: Vec<Bytes>,
+    pub updated_values: Vec<Bytes>,
+    pub deleted_keys: Vec<Bytes>,
+}
+
+impl StoredFailureCase {
+    pub fn load(path: &str) -> StoredFailureCase {
+        deserialize_from_json_gzip(path).expect("stored failure case")
+    }
+}