perf: reuse accounts trie in payload processing (#16181)

Co-authored-by: Alexey Shekhirin <[email protected]>
Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com>

Author: 3 people

crates/engine/tree/benches/state_root_task.rs

@@ -227,7 +227,7 @@ fn bench_state_root(c: &mut Criterion) {
 
                         (genesis_hash, payload_processor, provider, state_updates)
                     },
-                    |(genesis_hash, payload_processor, provider, state_updates)| {
+                    |(genesis_hash, mut payload_processor, provider, state_updates)| {
                         black_box({
                             let mut handle = payload_processor.spawn(
                                 Default::default(),

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

@@ -2283,7 +2283,7 @@ where
             // background task or try to compute it in parallel
             if use_state_root_task {
                 match handle.state_root() {
-                    Ok(StateRootComputeOutcome { state_root, trie_updates }) => {
+                    Ok(StateRootComputeOutcome { state_root, trie_updates, trie }) => {
                         let elapsed = execution_finish.elapsed();
                         info!(target: "engine::tree", ?state_root, ?elapsed, "State root task finished");
                         // we double check the state root here for good measure
@@ -2297,6 +2297,9 @@ where
                                 "State root task returned incorrect state root"
                             );
                         }
+
+                        // hold on to the sparse trie for the next payload
+                        self.payload_processor.set_sparse_trie(trie);
                     }
                     Err(error) => {
                         debug!(target: "engine::tree", %error, "Background parallel state root computation failed");

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

@@ -28,6 +28,7 @@ use reth_trie_parallel::{
     proof_task::{ProofTaskCtx, ProofTaskManager},
     root::ParallelStateRootError,
 };
+use reth_trie_sparse::SparseTrieState;
 use std::{
     collections::VecDeque,
     sync::{
@@ -67,6 +68,9 @@ where
     precompile_cache_disabled: bool,
     /// Precompile cache map.
     precompile_cache_map: PrecompileCacheMap<SpecFor<Evm>>,
+    /// A sparse trie, kept around to be used for the state root computation so that allocations
+    /// can be minimized.
+    sparse_trie: Option<SparseTrieState>,
     _marker: std::marker::PhantomData<N>,
 }
 
@@ -91,6 +95,7 @@ where
             evm_config,
             precompile_cache_disabled: config.precompile_cache_disabled(),
             precompile_cache_map,
+            sparse_trie: None,
             _marker: Default::default(),
         }
     }
@@ -134,7 +139,7 @@ where
     /// This returns a handle to await the final state root and to interact with the tasks (e.g.
     /// canceling)
     pub fn spawn<P>(
-        &self,
+        &mut self,
         header: SealedHeaderFor<N>,
         transactions: VecDeque<Recovered<N::SignedTx>>,
         provider_builder: StateProviderBuilder<N, P>,
@@ -191,11 +196,15 @@ where
             multi_proof_task.run();
         });
 
-        let mut sparse_trie_task = SparseTrieTask::new(
+        // take the sparse trie if it was set
+        let sparse_trie = self.sparse_trie.take();
+
+        let mut sparse_trie_task = SparseTrieTask::new_with_stored_trie(
             self.executor.clone(),
             sparse_trie_rx,
             proof_task.handle(),
             self.trie_metrics.clone(),
+            sparse_trie,
         );
 
         // wire the sparse trie to the state root response receiver
@@ -241,6 +250,11 @@ where
         PayloadHandle { to_multi_proof: None, prewarm_handle, state_root: None }
     }
 
+    /// Sets the sparse trie to be kept around for the state root computation.
+    pub(super) fn set_sparse_trie(&mut self, sparse_trie: SparseTrieState) {
+        self.sparse_trie = Some(sparse_trie);
+    }
+
     /// Spawn prewarming optionally wired to the multiproof task for target updates.
     fn spawn_caching_with<P>(
         &self,
@@ -566,7 +580,7 @@ mod tests {
             }
         }
 
-        let payload_processor = PayloadProcessor::<EthPrimitives, _>::new(
+        let mut payload_processor = PayloadProcessor::<EthPrimitives, _>::new(
             WorkloadExecutor::default(),
             EthEvmConfig::new(factory.chain_spec()),
             &TreeConfig::default(),

crates/engine/tree/src/tree/payload_processor/sparse_trie.rs

@@ -11,7 +11,7 @@ use reth_trie_parallel::root::ParallelStateRootError;
 use reth_trie_sparse::{
     blinded::{BlindedProvider, BlindedProviderFactory},
     errors::{SparseStateTrieResult, SparseTrieErrorKind},
-    SparseStateTrie,
+    SparseStateTrie, SparseTrieState,
 };
 use std::{
     sync::mpsc,
@@ -63,6 +63,43 @@ where
         }
     }
 
+    /// Creates a new sparse trie, populating the accounts trie with the given cleared
+    /// `SparseTrieState` if it exists.
+    pub(super) fn new_with_stored_trie(
+        executor: WorkloadExecutor,
+        updates: mpsc::Receiver<SparseTrieUpdate>,
+        blinded_provider_factory: BPF,
+        trie_metrics: MultiProofTaskMetrics,
+        sparse_trie_state: Option<SparseTrieState>,
+    ) -> Self {
+        if let Some(sparse_trie_state) = sparse_trie_state {
+            Self::with_accounts_trie(
+                executor,
+                updates,
+                blinded_provider_factory,
+                trie_metrics,
+                sparse_trie_state,
+            )
+        } else {
+            Self::new(executor, updates, blinded_provider_factory, trie_metrics)
+        }
+    }
+
+    /// Creates a new sparse trie task, using the given cleared `SparseTrieState` for the accounts
+    /// trie.
+    pub(super) fn with_accounts_trie(
+        executor: WorkloadExecutor,
+        updates: mpsc::Receiver<SparseTrieUpdate>,
+        blinded_provider_factory: BPF,
+        metrics: MultiProofTaskMetrics,
+        sparse_trie_state: SparseTrieState,
+    ) -> Self {
+        let mut trie = SparseStateTrie::new(blinded_provider_factory).with_updates(true);
+        trie.populate_from(sparse_trie_state);
+
+        Self { executor, updates, metrics, trie }
+    }
+
     /// Runs the sparse trie task to completion.
     ///
     /// This waits for new incoming [`SparseTrieUpdate`].
@@ -109,7 +146,10 @@ where
         self.metrics.sparse_trie_final_update_duration_histogram.record(start.elapsed());
         self.metrics.sparse_trie_total_duration_histogram.record(now.elapsed());
 
-        Ok(StateRootComputeOutcome { state_root, trie_updates })
+        // take the account trie
+        let trie = self.trie.take_cleared_account_trie_state();
+
+        Ok(StateRootComputeOutcome { state_root, trie_updates, trie })
     }
 }
 
@@ -121,6 +161,8 @@ pub struct StateRootComputeOutcome {
     pub state_root: B256,
     /// The trie updates.
     pub trie_updates: TrieUpdates,
+    /// The account state trie.
+    pub trie: SparseTrieState,
 }
 
 /// Updates the sparse trie with the given proofs and state, and returns the elapsed time.

crates/trie/sparse/src/state.rs

@@ -1,6 +1,6 @@
 use crate::{
     blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
-    LeafLookup, RevealedSparseTrie, SparseTrie, TrieMasks,
+    LeafLookup, RevealedSparseTrie, SparseTrie, SparseTrieState, TrieMasks,
 };
 use alloc::{collections::VecDeque, vec::Vec};
 use alloy_primitives::{
@@ -107,6 +107,19 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         self.revealed_account_paths.contains(&Nibbles::unpack(account))
     }
 
+    /// Uses the input `SparseTrieState` to populate the backing data structures in the `state`
+    /// trie.
+    pub fn populate_from(&mut self, trie: SparseTrieState) {
+        if let Some(new_trie) = self.state.as_revealed_mut() {
+            new_trie.use_allocated_state(trie);
+        } else {
+            self.state = SparseTrie::revealed_with_provider(
+                self.provider_factory.account_node_provider(),
+                trie,
+            )
+        }
+    }
+
     /// Was the account witness for `address` complete?
     pub fn check_valid_account_witness(&self, address: B256) -> bool {
         let path = Nibbles::unpack(address);
@@ -343,7 +356,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
     ) -> SparseStateTrieResult<()> {
         let FilteredProofNodes {
             nodes,
-            new_nodes,
+            new_nodes: _,
             total_nodes: _total_nodes,
             skipped_nodes: _skipped_nodes,
         } = filter_revealed_nodes(account_subtree, &self.revealed_account_paths)?;
@@ -366,9 +379,6 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
                 self.retain_updates,
             )?;
 
-            // Reserve the capacity for new nodes ahead of time.
-            trie.reserve_nodes(new_nodes);
-
             // Reveal the remaining proof nodes.
             for (path, node) in account_nodes {
                 let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
@@ -650,7 +660,7 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         &mut self,
     ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
         match self.state {
-            SparseTrie::Blind => {
+            SparseTrie::Blind | SparseTrie::AllocatedEmpty { .. } => {
                 let (root_node, hash_mask, tree_mask) = self
                     .provider_factory
                     .account_node_provider()
@@ -868,6 +878,12 @@ impl<F: BlindedProviderFactory> SparseStateTrie<F> {
         storage_trie.remove_leaf(slot)?;
         Ok(())
     }
+
+    /// Clears and takes the account trie.
+    pub fn take_cleared_account_trie_state(&mut self) -> SparseTrieState {
+        let trie = core::mem::take(&mut self.state);
+        trie.cleared()
+    }
 }
 
 /// Result of [`filter_revealed_nodes`].