From 3e0960cb11d969fb41b07f472b71371a490fb6e9 Mon Sep 17 00:00:00 2001 From: Dan Cline <6798349+Rjected@users.noreply.github.com> Date: Mon, 16 Jun 2025 15:17:00 +0200 Subject: [PATCH] perf: reuse accounts trie in payload processing (#16181) Co-authored-by: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com> Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com> --- crates/engine/tree/benches/state_root_task.rs | 2 +- crates/engine/tree/src/tree/mod.rs | 5 +- .../tree/src/tree/payload_processor/mod.rs | 20 ++- .../src/tree/payload_processor/sparse_trie.rs | 46 +++++- crates/trie/sparse/src/state.rs | 28 +++- crates/trie/sparse/src/trie.rs | 156 +++++++++++++++--- 6 files changed, 221 insertions(+), 36 deletions(-) diff --git a/crates/engine/tree/benches/state_root_task.rs b/crates/engine/tree/benches/state_root_task.rs index d705bfecd8..1eeb7a47f5 100644 --- a/crates/engine/tree/benches/state_root_task.rs +++ b/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(), diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 7b8454175e..26cc096535 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/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"); diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs index 5c782fbd4b..118a77521b 100644 --- a/crates/engine/tree/src/tree/payload_processor/mod.rs +++ b/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>, + /// A sparse trie, kept around to be used for the state root computation so that allocations + /// can be minimized. + sparse_trie: Option, _marker: std::marker::PhantomData, } @@ -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

( - &self, + &mut self, header: SealedHeaderFor, transactions: VecDeque>, provider_builder: StateProviderBuilder, @@ -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

( &self, @@ -566,7 +580,7 @@ mod tests { } } - let payload_processor = PayloadProcessor::::new( + let mut payload_processor = PayloadProcessor::::new( WorkloadExecutor::default(), EthEvmConfig::new(factory.chain_spec()), &TreeConfig::default(), diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs index 93f0049109..c8de07c1ec 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/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, + blinded_provider_factory: BPF, + trie_metrics: MultiProofTaskMetrics, + sparse_trie_state: Option, + ) -> 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, + 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. diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs index 39e305f498..dc8ac3506f 100644 --- a/crates/trie/sparse/src/state.rs +++ b/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 SparseStateTrie { 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 SparseStateTrie { ) -> 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 SparseStateTrie { 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 SparseStateTrie { &mut self, ) -> SparseStateTrieResult<&mut RevealedSparseTrie> { 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 SparseStateTrie { 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`]. diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 5759b5d4b8..8636383a05 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -52,6 +52,19 @@ impl TrieMasks { } } +/// A struct for keeping the hashmaps from `RevealedSparseTrie`. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct SparseTrieState { + /// Map from a path (nibbles) to its corresponding sparse trie node. + nodes: HashMap, + /// When a branch is set, the corresponding child subtree is stored in the database. + branch_node_tree_masks: HashMap, + /// When a bit is set, the corresponding child is stored as a hash in the database. + branch_node_hash_masks: HashMap, + /// Map from leaf key paths to their values. + values: HashMap>, +} + /// A sparse trie that is either in a "blind" state (no nodes are revealed, root node hash is /// unknown) or in a "revealed" state (root node has been revealed and the trie can be updated). /// @@ -64,8 +77,15 @@ impl TrieMasks { /// 2. Update tracking - changes to the trie structure can be tracked and selectively persisted /// 3. Incremental operations - nodes can be revealed as needed without loading the entire trie. /// This is what gives rise to the notion of a "sparse" trie. -#[derive(PartialEq, Eq, Default)] +#[derive(PartialEq, Eq, Default, Clone)] pub enum SparseTrie

{ + /// This is a variant that can be used to store a previously allocated trie. In these cases, + /// the trie will still be treated as blind, but the allocated trie will be reused if the trie + /// becomes revealed. + AllocatedEmpty { + /// This is the state of the allocated trie. + allocated: SparseTrieState, + }, /// The trie is blind -- no nodes have been revealed /// /// This is the default state. In this state, @@ -83,6 +103,7 @@ pub enum SparseTrie

{ impl

fmt::Debug for SparseTrie

{ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { + Self::AllocatedEmpty { .. } => write!(f, "AllocatedEmpty"), Self::Blind => write!(f, "Blind"), Self::Revealed(revealed) => write!(f, "Revealed({revealed:?})"), } @@ -184,17 +205,54 @@ impl

SparseTrie

{ masks: TrieMasks, retain_updates: bool, ) -> SparseTrieResult<&mut RevealedSparseTrie

> { + // we take the allocated state here, which will make sure we are either `Blind` or + // `Revealed`, and giving us the allocated state if we were `AllocatedEmpty`. + let allocated = self.take_allocated_state(); + + // if `Blind`, we initialize the revealed trie if self.is_blind() { - *self = Self::Revealed(Box::new(RevealedSparseTrie::from_provider_and_root( - provider, - root, - masks, - retain_updates, - )?)) + let mut revealed = + RevealedSparseTrie::from_provider_and_root(provider, root, masks, retain_updates)?; + + // If we had an allocated state, we use its maps internally. use_allocated_state copies + // over any information we had from revealing. + if let Some(allocated) = allocated { + revealed.use_allocated_state(allocated); + } + + *self = Self::Revealed(Box::new(revealed)); } Ok(self.as_revealed_mut().unwrap()) } + /// Take the allocated state if this is `AllocatedEmpty`, otherwise returns `None`. + /// + /// Converts this `SparseTrie` into `Blind` if this was `AllocatedEmpty`. + pub fn take_allocated_state(&mut self) -> Option { + if let Self::AllocatedEmpty { allocated } = self { + let state = core::mem::take(allocated); + *self = Self::Blind; + Some(state) + } else { + None + } + } + + /// Creates a new trie with the given provider and sparse trie state. + pub fn revealed_with_provider(provider: P, revealed_state: SparseTrieState) -> Self { + let revealed = RevealedSparseTrie { + provider, + nodes: revealed_state.nodes, + branch_node_tree_masks: revealed_state.branch_node_tree_masks, + branch_node_hash_masks: revealed_state.branch_node_hash_masks, + values: revealed_state.values, + prefix_set: PrefixSetMut::default(), + updates: None, + rlp_buf: Vec::new(), + }; + Self::Revealed(Box::new(revealed)) + } + /// Wipes the trie by removing all nodes and values, /// and resetting the trie to only contain an empty root node. /// @@ -205,6 +263,16 @@ impl

SparseTrie

{ Ok(()) } + /// Returns a `SparseTrieState` obtained by clearing the sparse trie state and reusing the + /// allocated state if it was `AllocatedEmpty` or `Revealed`. + pub fn cleared(self) -> SparseTrieState { + match self { + Self::Revealed(revealed) => revealed.cleared_state(), + Self::AllocatedEmpty { allocated } => allocated, + Self::Blind => Default::default(), + } + } + /// Calculates the root hash of the trie. /// /// This will update any remaining dirty nodes before computing the root hash. @@ -481,6 +549,37 @@ impl

RevealedSparseTrie

{ } } + /// Sets the fields of this `RevealedSparseTrie` to the fields of the input + /// `SparseTrieState`. + /// + /// This is meant for reusing the allocated maps contained in the `SparseTrieState`. + /// + /// Copies over any existing nodes, branch masks, and values. + pub fn use_allocated_state(&mut self, mut other: SparseTrieState) { + for (path, node) in self.nodes.drain() { + other.nodes.insert(path, node); + } + for (path, mask) in self.branch_node_tree_masks.drain() { + other.branch_node_tree_masks.insert(path, mask); + } + for (path, mask) in self.branch_node_hash_masks.drain() { + other.branch_node_hash_masks.insert(path, mask); + } + for (path, value) in self.values.drain() { + other.values.insert(path, value); + } + + self.nodes = other.nodes; + self.branch_node_tree_masks = other.branch_node_tree_masks; + self.branch_node_hash_masks = other.branch_node_hash_masks; + self.values = other.values; + } + + /// Set the provider for the trie. + pub fn set_provider(&mut self, provider: P) { + self.provider = provider; + } + /// Configures the trie to retain information about updates. /// /// If `retain_updates` is true, the trie will record branch node updates and deletions. @@ -839,6 +938,33 @@ impl

RevealedSparseTrie

{ self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped); } + /// This clears all data structures in the sparse trie, keeping the backing data structures + /// allocated. + /// + /// This is useful for reusing the trie without needing to reallocate memory. + pub fn clear(&mut self) { + self.nodes.clear(); + self.branch_node_tree_masks.clear(); + self.branch_node_hash_masks.clear(); + self.values.clear(); + self.prefix_set.clear(); + if let Some(updates) = self.updates.as_mut() { + updates.clear() + } + self.rlp_buf.clear(); + } + + /// Returns the cleared `SparseTrieState` for this `RevealedSparseTrie`. + pub fn cleared_state(mut self) -> SparseTrieState { + self.clear(); + SparseTrieState { + nodes: self.nodes, + branch_node_tree_masks: self.branch_node_tree_masks, + branch_node_hash_masks: self.branch_node_hash_masks, + values: self.values, + } + } + /// Calculates and returns the root hash of the trie. /// /// Before computing the hash, this function processes any remaining (dirty) nodes by @@ -1325,22 +1451,6 @@ pub enum LeafLookup { } impl RevealedSparseTrie

{ - /// This clears all data structures in the sparse trie, keeping the backing data structures - /// allocated. - /// - /// This is useful for reusing the trie without needing to reallocate memory. - pub fn clear(&mut self) { - self.nodes.clear(); - self.branch_node_tree_masks.clear(); - self.branch_node_hash_masks.clear(); - self.values.clear(); - self.prefix_set.clear(); - if let Some(updates) = self.updates.as_mut() { - updates.clear() - } - self.rlp_buf.clear(); - } - /// Attempts to find a leaf node at the specified path. /// /// This method traverses the trie from the root down to the given path, checking