diff --git a/crates/trie/prefetch/src/prefetch.rs b/crates/trie/prefetch/src/prefetch.rs index ef98301fc..860356679 100644 --- a/crates/trie/prefetch/src/prefetch.rs +++ b/crates/trie/prefetch/src/prefetch.rs @@ -20,8 +20,10 @@ use tracing::{debug, trace}; /// Prefetch trie storage when executing transactions. #[derive(Debug, Clone)] pub struct TriePrefetch { - /// Cached storage slot. - cached: HashMap>, + /// Cached accounts. + cached_accounts: HashMap, + /// Cached storages. + cached_storages: HashMap>, /// State trie metrics. #[cfg(feature = "metrics")] metrics: TrieRootMetrics, @@ -37,7 +39,8 @@ impl TriePrefetch { /// Create new `TriePrefetch` instance. pub fn new() -> Self { Self { - cached: HashMap::new(), + cached_accounts: HashMap::new(), + cached_storages: HashMap::new(), #[cfg(feature = "metrics")] metrics: TrieRootMetrics::default(), } @@ -52,10 +55,13 @@ impl TriePrefetch { ) where DB: Database, { + let mut count= 0u64; loop { tokio::select! { hashed_state = prefetch_rx.recv() => { if let Some(hashed_state) = hashed_state { + count += 1; + let provider_ro = Arc::clone(&provider_ro); let hashed_state = self.deduplicate_and_update_cached(&hashed_state); if let Err(e) = self.prefetch_once::(provider_ro, hashed_state) { @@ -64,7 +70,7 @@ impl TriePrefetch { } } _ = &mut interrupt_rx => { - debug!(target: "trie::trie_prefetch", "Interrupted trie prefetch task"); + debug!(target: "trie::trie_prefetch", "Interrupted trie prefetch task. Processed {:?}, left {:?}", count, prefetch_rx.len()); return } } @@ -75,8 +81,19 @@ impl TriePrefetch { fn deduplicate_and_update_cached(&mut self, hashed_state: &HashedPostState) -> HashedPostState { let mut new_hashed_state = HashedPostState::default(); + // deduplicate accounts if their keys are not present in storages + for (address, account) in &hashed_state.accounts { + if !hashed_state.storages.contains_key(address) && + !self.cached_accounts.contains_key(address) + { + self.cached_accounts.insert(*address, true); + new_hashed_state.accounts.insert(*address, *account); + } + } + + // deduplicate storages for (address, storage) in &hashed_state.storages { - let cached_entry = self.cached.entry(*address).or_default(); + let cached_entry = self.cached_storages.entry(*address).or_default(); // Collect the keys to be added to `new_storage` after filtering let keys_to_add: Vec<_> = storage @@ -193,12 +210,12 @@ impl TriePrefetch { #[cfg(feature = "metrics")] self.metrics.record(stats); - trace!( + debug!( target: "trie::trie_prefetch", duration = ?stats.duration(), branches_added = stats.branches_added(), leaves_added = stats.leaves_added(), - "prefetched trie storages" + "prefetched account trie" ); Ok(())