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>
This commit is contained in:
Dan Cline
2025-06-16 15:17:00 +02:00
committed by GitHub
parent a8522e6a25
commit 3e0960cb11
6 changed files with 221 additions and 36 deletions

View File

@@ -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(),

View File

@@ -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");

View File

@@ -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(),

View File

@@ -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.

View File

@@ -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`].

View File

@@ -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<Nibbles, SparseNode>,
/// When a branch is set, the corresponding child subtree is stored in the database.
branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
/// When a bit is set, the corresponding child is stored as a hash in the database.
branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
/// Map from leaf key paths to their values.
values: HashMap<Nibbles, Vec<u8>>,
}
/// 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<P = DefaultBlindedProvider> {
/// 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<P = DefaultBlindedProvider> {
impl<P> fmt::Debug for SparseTrie<P> {
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<P> SparseTrie<P> {
masks: TrieMasks,
retain_updates: bool,
) -> SparseTrieResult<&mut RevealedSparseTrie<P>> {
// 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<SparseTrieState> {
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<P> SparseTrie<P> {
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<P> RevealedSparseTrie<P> {
}
}
/// 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<P> RevealedSparseTrie<P> {
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<P: BlindedProvider> RevealedSparseTrie<P> {
/// 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