feat: optimize prefetch

Author: pythonberg1997

crates/blockchain-tree/src/chain.rs

@@ -3,35 +3,37 @@
 //! A [`Chain`] contains the state of accounts for the chain after execution of its constituent
 //! blocks, as well as a list of the blocks the chain is composed of.
 
-use super::externals::TreeExternals;
-use crate::BundleStateDataRef;
+use std::{
+    collections::{BTreeMap, HashMap},
+    ops::{Deref, DerefMut},
+    time::Instant,
+};
+
 use reth_blockchain_tree_api::{
-    error::{BlockchainTreeError, InsertBlockErrorKind},
-    BlockAttachment, BlockValidationKind,
+    BlockAttachment,
+    BlockValidationKind, error::{BlockchainTreeError, InsertBlockErrorKind},
 };
 use reth_consensus::{Consensus, ConsensusError, PostExecutionInput};
 use reth_db_api::database::Database;
 use reth_evm::execute::{BlockExecutorProvider, Executor};
 use reth_execution_errors::BlockExecutionError;
 use reth_execution_types::{Chain, ExecutionOutcome};
 use reth_primitives::{
-    BlockHash, BlockNumber, ForkBlock, GotExpected, Header, SealedBlockWithSenders, SealedHeader,
-    B256, U256,
+    B256, BlockHash, BlockNumber, ForkBlock, GotExpected, Header, SealedBlockWithSenders,
+    SealedHeader, U256,
 };
 use reth_provider::{
-    providers::{BundleStateProvider, ConsistentDbView},
-    FullExecutionDataProvider, ProviderError, StateRootProvider,
+    FullExecutionDataProvider,
+    ProviderError, providers::{BundleStateProvider, ConsistentDbView}, StateRootProvider,
 };
 use reth_revm::{database::StateProviderDatabase, primitives::EvmState};
-use reth_trie::{updates::TrieUpdates, HashedPostState};
+use reth_trie::{HashedPostState, updates::TrieUpdates};
 use reth_trie_parallel::parallel_root::ParallelStateRoot;
 use reth_trie_prefetch::TriePrefetch;
-use std::{
-    collections::{BTreeMap, HashMap},
-    ops::{Deref, DerefMut},
-    sync::Arc,
-    time::Instant,
-};
+
+use crate::BundleStateDataRef;
+
+use super::externals::TreeExternals;
 
 /// A chain in the blockchain tree that has functionality to execute blocks and append them to
 /// itself.
@@ -227,22 +229,25 @@ impl AppendableChain {
         let block_hash = block.hash();
         let block = block.unseal();
 
+        let execute_start = Instant::now();
         let state = executor.execute((&block, U256::MAX, ancestor_blocks).into())?;
         externals.consensus.validate_block_post_execution(
             &block,
             PostExecutionInput::new(&state.receipts, &state.requests),
         )?;
 
-        let initial_execution_outcome = ExecutionOutcome::from((state, block.number));
+        tracing::debug!(
+            target: "blockchain_tree::chain",
+            number = block.number,
+            duration = ?execute_start.elapsed(),
+            "executed and validated block"
+        );
 
-        // stop the prefetch task.
-        if let Some(interrupt_tx) = interrupt_tx {
-            let _ = interrupt_tx.send(());
-        }
+        let initial_execution_outcome = ExecutionOutcome::from((state, block.number));
 
         // check state root if the block extends the canonical chain __and__ if state root
         // validation was requested.
-        if block_validation_kind.is_exhaustive() {
+        let result = if block_validation_kind.is_exhaustive() {
             // calculate and check state root
             let start = Instant::now();
             let (state_root, trie_updates) = if block_attachment.is_canonical() {
@@ -285,7 +290,14 @@ impl AppendableChain {
             Ok((initial_execution_outcome, trie_updates))
         } else {
             Ok((initial_execution_outcome, None))
-        }
+        };
+
+        // stop the prefetch task.
+        if let Some(interrupt_tx) = interrupt_tx {
+            let _ = interrupt_tx.send(());
+        };
+
+        result
     }
 
     /// Validate and execute the given block, and append it to this chain.
@@ -358,18 +370,11 @@ impl AppendableChain {
         let (interrupt_tx, interrupt_rx) = tokio::sync::oneshot::channel();
 
         let mut trie_prefetch = TriePrefetch::new();
-        let consistent_view = if let Ok(view) =
-            ConsistentDbView::new_with_latest_tip(externals.provider_factory.clone())
-        {
-            view
-        } else {
-            tracing::debug!("Failed to create consistent view for trie prefetch");
-            return (None, None)
-        };
+        let provider_factory = externals.provider_factory.clone();
 
         tokio::spawn({
             async move {
-                trie_prefetch.run::<DB>(Arc::new(consistent_view), prefetch_rx, interrupt_rx).await;
+                trie_prefetch.run::<DB>(provider_factory, prefetch_rx, interrupt_rx).await;
             }
         });
 

crates/trie/parallel/src/parallel_root.rs

@@ -1,24 +1,27 @@
-#[cfg(feature = "metrics")]
-use crate::metrics::ParallelStateRootMetrics;
-use crate::{stats::ParallelTrieTracker, storage_root_targets::StorageRootTargets};
+use std::collections::HashMap;
+
 use alloy_rlp::{BufMut, Encodable};
 use rayon::prelude::*;
+use thiserror::Error;
+use tracing::*;
+
 use reth_db_api::database::Database;
 use reth_execution_errors::StorageRootError;
 use reth_primitives::B256;
-use reth_provider::{providers::ConsistentDbView, DatabaseProviderFactory, ProviderError};
+use reth_provider::{DatabaseProviderFactory, ProviderError, providers::ConsistentDbView};
 use reth_trie::{
+    HashBuilder,
     hashed_cursor::{HashedCursorFactory, HashedPostStateCursorFactory},
+    HashedPostState,
+    Nibbles,
     node_iter::{TrieElement, TrieNodeIter},
-    trie_cursor::TrieCursorFactory,
-    updates::TrieUpdates,
-    walker::TrieWalker,
-    HashBuilder, HashedPostState, Nibbles, StorageRoot, TrieAccount,
+    StorageRoot, trie_cursor::TrieCursorFactory, TrieAccount, updates::TrieUpdates, walker::TrieWalker,
 };
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
-use std::collections::HashMap;
-use thiserror::Error;
-use tracing::*;
+
+use crate::{stats::ParallelTrieTracker, storage_root_targets::StorageRootTargets};
+#[cfg(feature = "metrics")]
+use crate::metrics::ParallelStateRootMetrics;
 
 /// Parallel incremental state root calculator.
 ///
@@ -129,14 +132,17 @@ where
             hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,
         );
 
+        let account_tree_start = std::time::Instant::now();
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
         let mut account_rlp = Vec::with_capacity(128);
         while let Some(node) = account_node_iter.try_next().map_err(ProviderError::Database)? {
             match node {
                 TrieElement::Branch(node) => {
+                    tracker.inc_branch();
                     hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                 }
                 TrieElement::Leaf(hashed_address, account) => {
+                    tracker.inc_leaf();
                     let (storage_root, _, updates) = match storage_roots.remove(&hashed_address) {
                         Some(result) => result,
                         // Since we do not store all intermediate nodes in the database, there might
@@ -174,15 +180,18 @@ where
             prefix_sets.destroyed_accounts,
         );
 
+        let account_tree_duration = account_tree_start.elapsed();
         let stats = tracker.finish();
 
         #[cfg(feature = "metrics")]
         self.metrics.record_state_trie(stats);
 
-        trace!(
+        debug!(
             target: "trie::parallel_state_root",
             %root,
             duration = ?stats.duration(),
+            account_tree_duration = ?account_tree_duration,
+            storage_trees_duration = ?(stats.duration() - account_tree_duration),
             branches_added = stats.branches_added(),
             leaves_added = stats.leaves_added(),
             missed_leaves = stats.missed_leaves(),
@@ -218,11 +227,13 @@ impl From<ParallelStateRootError> for ProviderError {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use rand::Rng;
-    use reth_primitives::{keccak256, Account, Address, StorageEntry, U256};
-    use reth_provider::{test_utils::create_test_provider_factory, HashingWriter};
-    use reth_trie::{test_utils, HashedStorage};
+
+    use reth_primitives::{Account, Address, keccak256, StorageEntry, U256};
+    use reth_provider::{HashingWriter, test_utils::create_test_provider_factory};
+    use reth_trie::{HashedStorage, test_utils};
+
+    use super::*;
 
     #[tokio::test]
     async fn random_parallel_root() {

crates/trie/prefetch/src/lib.rs

@@ -7,8 +7,11 @@
 )]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
+pub use prefetch::TriePrefetch;
 pub use reth_trie_parallel::StorageRootTargets;
 
+/// Trie prefetch stats.
+pub mod stats;
+
 /// Implementation of trie prefetch.
-mod prefetch;
-pub use prefetch::TriePrefetch;
+pub mod prefetch;