feat: enable prefetch on the new engine

Author: pythonberg1997

Cargo.lock

@@ -30,7 +30,6 @@ use std::{
     clone::Clone,
     collections::BTreeMap,
     ops::{Deref, DerefMut},
-    sync::Arc,
     time::Instant,
 };
 
@@ -231,14 +230,9 @@ impl AppendableChain {
 
         let initial_execution_outcome = ExecutionOutcome::from((state, block.number));
 
-        // stop the prefetch task.
-        if let Some(interrupt_tx) = interrupt_tx {
-            let _ = interrupt_tx.send(());
-        }
-
         // 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() {
@@ -283,7 +277,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.
@@ -356,18 +357,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(Arc::new(consistent_view), prefetch_rx, interrupt_rx).await;
+                trie_prefetch.run(provider_factory, prefetch_rx, interrupt_rx).await;
             }
         });
 

crates/blockchain-tree/src/chain.rs

@@ -80,6 +80,7 @@ where
         invalid_block_hook: Box<dyn InvalidBlockHook>,
         sync_metrics_tx: MetricEventsSender,
         skip_state_root_validation: bool,
+        enable_prefetch: bool,
     ) -> Self {
         let engine_kind =
             if chain_spec.is_optimism() { EngineApiKind::OpStack } else { EngineApiKind::Ethereum };
@@ -104,6 +105,7 @@ where
             invalid_block_hook,
             engine_kind,
             skip_state_root_validation,
+            enable_prefetch,
         );
 
         let engine_handler = EngineApiRequestHandler::new(to_tree_tx, from_tree);
@@ -217,6 +219,7 @@ mod tests {
             Box::new(NoopInvalidBlockHook::default()),
             sync_metrics_tx,
             false,
+            false,
         );
     }
 }

crates/engine/service/src/service.rs

@@ -33,6 +33,7 @@ reth-stages-api.workspace = true
 reth-tasks.workspace = true
 reth-trie.workspace = true
 reth-trie-parallel.workspace = true
+reth-trie-prefetch.workspace = true
 
 # alloy
 alloy-primitives.workspace = true

crates/engine/tree/Cargo.toml

@@ -33,7 +33,8 @@ use reth_payload_builder::PayloadBuilderHandle;
 use reth_payload_primitives::{PayloadAttributes, PayloadBuilder, PayloadBuilderAttributes};
 use reth_payload_validator::ExecutionPayloadValidator;
 use reth_primitives::{
-    Block, GotExpected, Header, SealedBlock, SealedBlockWithSenders, SealedHeader,
+    revm_primitives::EvmState, Block,  GotExpected, Header, SealedBlock,
+    SealedBlockWithSenders, SealedHeader,
 };
 use reth_provider::{
     providers::ConsistentDbView, BlockReader, DatabaseProviderFactory, ExecutionOutcome,
@@ -44,6 +45,7 @@ use reth_revm::database::StateProviderDatabase;
 use reth_stages_api::ControlFlow;
 use reth_trie::{updates::TrieUpdates, HashedPostState, TrieInput};
 use reth_trie_parallel::parallel_root::{ParallelStateRoot, ParallelStateRootError};
+use reth_trie_prefetch::TriePrefetch;
 use std::{
     cmp::Ordering,
     collections::{btree_map, hash_map, BTreeMap, VecDeque},
@@ -505,6 +507,8 @@ pub struct EngineApiTreeHandler<P, E, T: EngineTypes, Spec> {
     engine_kind: EngineApiKind,
     /// Flag indicating whether the state root validation should be skipped.
     skip_state_root_validation: bool,
+    /// Flag indicating whether to enable prefetch.
+    enable_prefetch: bool,
 }
 
 impl<P: Debug, E: Debug, T: EngineTypes + Debug, Spec: Debug> std::fmt::Debug
@@ -527,6 +531,8 @@ impl<P: Debug, E: Debug, T: EngineTypes + Debug, Spec: Debug> std::fmt::Debug
             .field("metrics", &self.metrics)
             .field("invalid_block_hook", &format!("{:p}", self.invalid_block_hook))
             .field("engine_kind", &self.engine_kind)
+            .field("skip_state_root_validation", &self.skip_state_root_validation)
+            .field("enable_prefetch", &self.enable_prefetch)
             .finish()
     }
 }
@@ -555,6 +561,7 @@ where
         config: TreeConfig,
         engine_kind: EngineApiKind,
         skip_state_root_validation: bool,
+        enable_prefetch: bool,
     ) -> Self {
         let (incoming_tx, incoming) = std::sync::mpsc::channel();
 
@@ -577,6 +584,7 @@ where
             invalid_block_hook: Box::new(NoopInvalidBlockHook),
             engine_kind,
             skip_state_root_validation,
+            enable_prefetch,
         }
     }
 
@@ -603,6 +611,7 @@ where
         invalid_block_hook: Box<dyn InvalidBlockHook>,
         kind: EngineApiKind,
         skip_state_root_validation: bool,
+        enable_prefetch: bool,
     ) -> (Sender<FromEngine<EngineApiRequest<T>>>, UnboundedReceiver<EngineApiEvent>) {
         let best_block_number = provider.best_block_number().unwrap_or(0);
         let header = provider.sealed_header(best_block_number).ok().flatten().unwrap_or_default();
@@ -634,10 +643,20 @@ where
             config,
             kind,
             skip_state_root_validation,
+            enable_prefetch,
         );
         task.set_invalid_block_hook(invalid_block_hook);
         let incoming = task.incoming_tx.clone();
-        std::thread::Builder::new().name("Tree Task".to_string()).spawn(|| task.run()).unwrap();
+        std::thread::Builder::new()
+            .name("Tree Task".to_string())
+            .spawn(move || {
+                let runtime =
+                    tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap();
+                runtime.block_on(async {
+                    task.run();
+                });
+            })
+            .unwrap();
         (incoming, outgoing)
     }
 
@@ -2175,8 +2194,16 @@ where
         }
 
         trace!(target: "engine::tree", block=?block.num_hash(), "Executing block");
-        let executor =
-            self.executor_provider.executor(StateProviderDatabase::new(&state_provider), None);
+        let (prefetch_tx, interrupt_tx) =
+            if self.enable_prefetch && !self.skip_state_root_validation {
+                self.setup_prefetch()
+            } else {
+                (None, None)
+            };
+
+        let executor = self
+            .executor_provider
+            .executor(StateProviderDatabase::new(&state_provider), prefetch_tx);
 
         let block_number = block.number;
         let block_hash = block.hash();
@@ -2247,6 +2274,11 @@ where
                 state_provider.state_root_with_updates(hashed_state.clone())?
             };
 
+            // stop the prefetch task.
+            if let Some(interrupt_tx) = interrupt_tx {
+                let _ = interrupt_tx.send(());
+            };
+
             if state_root != block.state_root {
                 // call post-block hook
                 self.invalid_block_hook.on_invalid_block(
@@ -2575,6 +2607,22 @@ where
         );
         Ok(())
     }
+
+    fn setup_prefetch(&self) -> (Option<UnboundedSender<EvmState>>, Option<oneshot::Sender<()>>) {
+        let (prefetch_tx, prefetch_rx) = tokio::sync::mpsc::unbounded_channel();
+        let (interrupt_tx, interrupt_rx) = oneshot::channel();
+
+        let mut trie_prefetch = TriePrefetch::new();
+        let provider_factory = self.provider.clone();
+
+        tokio::spawn({
+            async move {
+                trie_prefetch.run(provider_factory, prefetch_rx, interrupt_rx).await;
+            }
+        });
+
+        (Some(prefetch_tx), Some(interrupt_tx))
+    }
 }
 
 /// This is an error that can come from advancing persistence. Either this can be a
@@ -2728,6 +2776,7 @@ mod tests {
                 TreeConfig::default(),
                 EngineApiKind::Ethereum,
                 false,
+                false,
             );
 
             let block_builder = TestBlockBuilder::default().with_chain_spec((*chain_spec).clone());

crates/engine/tree/src/tree/mod.rs

@@ -237,6 +237,7 @@ where
                     ctx.invalid_block_hook()?,
                     ctx.sync_metrics_tx(),
                     ctx.node_config().skip_state_root_validation,
+                    ctx.node_config().enable_prefetch,
                 );
                 eth_service
             }
@@ -271,6 +272,7 @@ where
                     ctx.invalid_block_hook()?,
                     ctx.sync_metrics_tx(),
                     ctx.node_config().skip_state_root_validation,
+                    ctx.node_config().enable_prefetch,
                 );
                 eth_service
             }

crates/node/builder/src/launch/engine.rs

@@ -148,14 +148,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(rx) => rx.recv().map_err(|_| {
                             ParallelStateRootError::StorageRoot(StorageRootError::Database(
@@ -199,15 +202,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(),

crates/trie/parallel/src/parallel_root.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;