diff --git a/.changelog/brave-dogs-fly.md b/.changelog/brave-dogs-fly.md new file mode 100644 index 0000000000..2dc52e0e30 --- /dev/null +++ b/.changelog/brave-dogs-fly.md @@ -0,0 +1,10 @@ +--- +reth-engine-primitives: patch +reth-engine-tree: patch +reth-node-core: patch +reth-trie-parallel: minor +--- + +Removed legacy proof calculation system and V2-specific configuration flags. + +Removed the legacy (non-V2) proof calculation code paths, simplified multiproof task architecture by removing the dual-mode system, and cleaned up V2-specific CLI flags (`--engine.disable-proof-v2`, `--engine.disable-trie-cache`) that are no longer needed. The codebase now exclusively uses V2 proofs with the sparse trie cache. diff --git a/Cargo.lock b/Cargo.lock index 606886d94b..be99b14e9a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8389,7 +8389,6 @@ dependencies = [ "revm-state", "schnellru", "serde_json", - "smallvec", "thiserror 2.0.18", "tokio", "tracing", @@ -10596,7 +10595,6 @@ dependencies = [ "reth-storage-errors", "reth-tasks", "reth-trie", - "reth-trie-common", "reth-trie-db", "reth-trie-sparse", "thiserror 2.0.18", diff --git a/crates/engine/primitives/src/config.rs b/crates/engine/primitives/src/config.rs index 536beb0d39..dfd1c4e006 100644 --- a/crates/engine/primitives/src/config.rs +++ b/crates/engine/primitives/src/config.rs @@ -39,11 +39,6 @@ pub const SMALL_BLOCK_MULTIPROOF_CHUNK_SIZE: usize = 30; /// Gas threshold below which the small block chunk size is used. pub const SMALL_BLOCK_GAS_THRESHOLD: u64 = 20_000_000; -/// The size of proof targets chunk to spawn in one multiproof calculation when V2 proofs are -/// enabled. This is 4x the default chunk size to take advantage of more efficient V2 proof -/// computation. -pub const DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE_V2: usize = DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE * 4; - /// Default number of reserved CPU cores for non-reth processes. /// /// This will be deducted from the thread count of main reth global threadpool. @@ -162,12 +157,8 @@ pub struct TreeConfig { storage_worker_count: usize, /// Number of account proof worker threads. account_worker_count: usize, - /// Whether to disable V2 storage proofs. - disable_proof_v2: bool, /// Whether to disable cache metrics recording (can be expensive with large cached state). disable_cache_metrics: bool, - /// Whether to disable sparse trie cache. - disable_trie_cache: bool, /// Depth for sparse trie pruning after state root computation. sparse_trie_prune_depth: usize, /// Maximum number of storage tries to retain after pruning. @@ -205,9 +196,7 @@ impl Default for TreeConfig { allow_unwind_canonical_header: false, storage_worker_count: default_storage_worker_count(), account_worker_count: default_account_worker_count(), - disable_proof_v2: false, disable_cache_metrics: false, - disable_trie_cache: false, sparse_trie_prune_depth: DEFAULT_SPARSE_TRIE_PRUNE_DEPTH, sparse_trie_max_storage_tries: DEFAULT_SPARSE_TRIE_MAX_STORAGE_TRIES, disable_sparse_trie_cache_pruning: false, @@ -241,7 +230,6 @@ impl TreeConfig { allow_unwind_canonical_header: bool, storage_worker_count: usize, account_worker_count: usize, - disable_proof_v2: bool, disable_cache_metrics: bool, sparse_trie_prune_depth: usize, sparse_trie_max_storage_tries: usize, @@ -269,9 +257,7 @@ impl TreeConfig { allow_unwind_canonical_header, storage_worker_count, account_worker_count, - disable_proof_v2, disable_cache_metrics, - disable_trie_cache: false, sparse_trie_prune_depth, sparse_trie_max_storage_tries, disable_sparse_trie_cache_pruning: false, @@ -314,16 +300,9 @@ impl TreeConfig { self.multiproof_chunk_size } - /// Return the multiproof task chunk size, using the V2 default if V2 proofs are enabled - /// and the chunk size is at the default value. + /// Return the effective multiproof task chunk size. pub const fn effective_multiproof_chunk_size(&self) -> usize { - if !self.disable_proof_v2 && - self.multiproof_chunk_size == DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE - { - DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE_V2 - } else { - self.multiproof_chunk_size - } + self.multiproof_chunk_size } /// Return the number of reserved CPU cores for non-reth processes @@ -559,17 +538,6 @@ impl TreeConfig { self } - /// Return whether V2 storage proofs are disabled. - pub const fn disable_proof_v2(&self) -> bool { - self.disable_proof_v2 - } - - /// Setter for whether to disable V2 storage proofs. - pub const fn with_disable_proof_v2(mut self, disable_proof_v2: bool) -> Self { - self.disable_proof_v2 = disable_proof_v2; - self - } - /// Returns whether cache metrics recording is disabled. pub const fn disable_cache_metrics(&self) -> bool { self.disable_cache_metrics @@ -581,17 +549,6 @@ impl TreeConfig { self } - /// Returns whether sparse trie cache is disabled. - pub const fn disable_trie_cache(&self) -> bool { - self.disable_trie_cache - } - - /// Setter for whether to disable sparse trie cache. - pub const fn with_disable_trie_cache(mut self, value: bool) -> Self { - self.disable_trie_cache = value; - self - } - /// Returns the sparse trie prune depth. pub const fn sparse_trie_prune_depth(&self) -> usize { self.sparse_trie_prune_depth diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index cfe42076ea..96771863c7 100644 --- a/crates/engine/tree/Cargo.toml +++ b/crates/engine/tree/Cargo.toml @@ -54,7 +54,6 @@ thiserror.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread", "sync", "macros"] } fixed-cache.workspace = true moka = { workspace = true, features = ["sync"] } -smallvec.workspace = true # metrics metrics.workspace = true diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs index 8ee3c6d700..145f75da97 100644 --- a/crates/engine/tree/src/tree/payload_processor/mod.rs +++ b/crates/engine/tree/src/tree/payload_processor/mod.rs @@ -2,12 +2,12 @@ use super::precompile_cache::PrecompileCacheMap; use crate::tree::{ - cached_state::{CachedStateMetrics, CachedStateProvider, ExecutionCache, SavedCache}, + cached_state::{CachedStateMetrics, ExecutionCache, SavedCache}, payload_processor::{ prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmMode, PrewarmTaskEvent}, sparse_trie::StateRootComputeOutcome, }, - sparse_trie::{SparseTrieCacheTask, SparseTrieTask, SpawnedSparseTrieTask}, + sparse_trie::SparseTrieCacheTask, CacheWaitDurations, StateProviderBuilder, TreeConfig, WaitForCaches, }; use alloy_eip7928::BlockAccessList; @@ -16,7 +16,7 @@ use alloy_evm::block::StateChangeSource; use alloy_primitives::B256; use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender}; use metrics::{Counter, Histogram}; -use multiproof::{SparseTrieUpdate, *}; +use multiproof::*; use parking_lot::RwLock; use prewarm::PrewarmMetrics; use rayon::prelude::*; @@ -30,8 +30,7 @@ use reth_evm::{ use reth_metrics::Metrics; use reth_primitives_traits::{FastInstant as Instant, NodePrimitives}; use reth_provider::{ - BlockExecutionOutput, BlockReader, DatabaseProviderROFactory, StateProvider, - StateProviderFactory, StateReader, + BlockExecutionOutput, BlockReader, DatabaseProviderROFactory, StateProviderFactory, StateReader, }; use reth_revm::{db::BundleState, state::EvmState}; use reth_tasks::{ForEachOrdered, Runtime}; @@ -249,7 +248,7 @@ where /// /// ## Sparse trie task /// - /// Responsible for calculating the state root based on the received [`SparseTrieUpdate`]. + /// Responsible for calculating the state root. /// /// This task runs until there are no further updates to process. /// @@ -284,66 +283,32 @@ where self.spawn_tx_iterator(transactions, env.transaction_count); let span = Span::current(); - let (to_sparse_trie, sparse_trie_rx) = channel(); let (to_multi_proof, from_multi_proof) = crossbeam_channel::unbounded(); - let v2_proofs_enabled = !config.disable_proof_v2(); let parent_state_root = env.parent_state_root; let transaction_count = env.transaction_count; let chunk_size = Self::adaptive_chunk_size(config, env.gas_used); let prewarm_handle = self.spawn_caching_with( env, prewarm_rx, - provider_builder.clone(), + provider_builder, Some(to_multi_proof.clone()), bal, - v2_proofs_enabled, ); // Create and spawn the storage proof task. let task_ctx = ProofTaskCtx::new(multiproof_provider_factory); let halve_workers = transaction_count <= Self::SMALL_BLOCK_PROOF_WORKER_TX_THRESHOLD; - let proof_handle = - ProofWorkerHandle::new(&self.executor, task_ctx, halve_workers, v2_proofs_enabled); - - if config.disable_trie_cache() { - let multi_proof_task = MultiProofTask::new( - proof_handle.clone(), - to_sparse_trie, - chunk_size, - to_multi_proof.clone(), - from_multi_proof.clone(), - ) - .with_v2_proofs_enabled(v2_proofs_enabled); - - // spawn multi-proof task - let parent_span = span.clone(); - let saved_cache = prewarm_handle.saved_cache.clone(); - self.executor.spawn_blocking(move || { - let _enter = parent_span.entered(); - // Build a state provider for the multiproof task - let provider = provider_builder.build().expect("failed to build provider"); - let provider = if let Some(saved_cache) = saved_cache { - let (cache, metrics, _disable_metrics) = saved_cache.split(); - Box::new(CachedStateProvider::new(provider, cache, metrics)) - as Box - } else { - Box::new(provider) - }; - multi_proof_task.run(provider); - }); - } + let proof_handle = ProofWorkerHandle::new(&self.executor, task_ctx, halve_workers); // wire the sparse trie to the state root response receiver let (state_root_tx, state_root_rx) = channel(); // Spawn the sparse trie task using any stored trie and parallel trie configuration. self.spawn_sparse_trie_task( - sparse_trie_rx, proof_handle, state_root_tx, from_multi_proof, - config, parent_state_root, chunk_size, ); @@ -373,9 +338,7 @@ where { let (prewarm_rx, execution_rx) = self.spawn_tx_iterator(transactions, env.transaction_count); - // This path doesn't use multiproof, so V2 proofs flag doesn't matter - let prewarm_handle = - self.spawn_caching_with(env, prewarm_rx, provider_builder, None, bal, false); + let prewarm_handle = self.spawn_caching_with(env, prewarm_rx, provider_builder, None, bal); PayloadHandle { to_multi_proof: None, prewarm_handle, @@ -494,7 +457,7 @@ where level = "debug", target = "engine::tree::payload_processor", skip_all, - fields(bal=%bal.is_some(), %v2_proofs_enabled) + fields(bal=%bal.is_some()) )] fn spawn_caching_with

( &self, @@ -503,7 +466,6 @@ where provider_builder: StateProviderBuilder, to_multi_proof: Option>, bal: Option>, - v2_proofs_enabled: bool, ) -> CacheTaskHandle where P: BlockReader + StateProviderFactory + StateReader + Clone + 'static, @@ -523,7 +485,6 @@ where terminate_execution: Arc::new(AtomicBool::new(false)), precompile_cache_disabled: self.precompile_cache_disabled, precompile_cache_map: self.precompile_cache_map.clone(), - v2_proofs_enabled, }; let (prewarm_task, to_prewarm_task) = PrewarmCacheTask::new( @@ -570,23 +531,19 @@ where } } - /// Spawns the [`SparseTrieTask`] for this payload processor. + /// Spawns the [`SparseTrieCacheTask`] for this payload processor. /// /// The trie is preserved when the new payload is a child of the previous one. - #[expect(clippy::too_many_arguments)] fn spawn_sparse_trie_task( &self, - sparse_trie_rx: mpsc::Receiver, proof_worker_handle: ProofWorkerHandle, state_root_tx: mpsc::Sender>, from_multi_proof: CrossbeamReceiver, - config: &TreeConfig, parent_state_root: B256, chunk_size: Option, ) { let preserved_sparse_trie = self.sparse_state_trie.clone(); let trie_metrics = self.trie_metrics.clone(); - let disable_trie_cache = config.disable_trie_cache(); let prune_depth = self.sparse_trie_prune_depth; let max_storage_tries = self.sparse_trie_max_storage_tries; let disable_cache_pruning = self.disable_sparse_trie_cache_pruning; @@ -625,23 +582,14 @@ where .with_updates(true) }); - let mut task = if disable_trie_cache { - SpawnedSparseTrieTask::Cleared(SparseTrieTask::new( - sparse_trie_rx, - proof_worker_handle, - trie_metrics.clone(), - sparse_state_trie, - )) - } else { - SpawnedSparseTrieTask::Cached(SparseTrieCacheTask::new_with_trie( - &executor, - from_multi_proof, - proof_worker_handle, - trie_metrics.clone(), - sparse_state_trie.with_skip_proof_node_filtering(true), - chunk_size, - )) - }; + let mut task = SparseTrieCacheTask::new_with_trie( + &executor, + from_multi_proof, + proof_worker_handle, + trie_metrics.clone(), + sparse_state_trie.with_skip_proof_node_filtering(true), + chunk_size, + ); let result = task.run(); // Capture the computed state_root before sending the result diff --git a/crates/engine/tree/src/tree/payload_processor/multiproof.rs b/crates/engine/tree/src/tree/payload_processor/multiproof.rs index f2efcff135..1c50a2157d 100644 --- a/crates/engine/tree/src/tree/payload_processor/multiproof.rs +++ b/crates/engine/tree/src/tree/payload_processor/multiproof.rs @@ -1,33 +1,16 @@ //! Multiproof task related functionality. -use crate::tree::payload_processor::bal::bal_to_hashed_post_state; -use alloy_eip7928::BlockAccessList; use alloy_evm::block::StateChangeSource; -use alloy_primitives::{keccak256, map::HashSet, B256}; -use crossbeam_channel::{unbounded, Receiver as CrossbeamReceiver, Sender as CrossbeamSender}; +use alloy_primitives::{keccak256, B256}; +use crossbeam_channel::Sender as CrossbeamSender; use derive_more::derive::Deref; use metrics::{Gauge, Histogram}; use reth_metrics::Metrics; -use reth_primitives_traits::FastInstant as Instant; -use reth_provider::AccountReader; use reth_revm::state::EvmState; -use reth_trie::{ - added_removed_keys::{default_added_removed_keys, MultiAddedRemovedKeys}, - proof_v2, HashedPostState, HashedStorage, MultiProofTargets, -}; -#[cfg(test)] -use reth_trie_parallel::stats::ParallelTrieTracker; -use reth_trie_parallel::{ - proof::ParallelProof, - proof_task::{ - AccountMultiproofInput, ProofResult, ProofResultContext, ProofResultMessage, - ProofWorkerHandle, - }, - targets_v2::MultiProofTargetsV2, -}; -use revm_primitives::map::{hash_map, B256Map}; -use std::{collections::BTreeMap, sync::Arc}; -use tracing::{debug, error, instrument, trace}; +use reth_trie::{HashedPostState, HashedStorage}; +use reth_trie_parallel::targets_v2::MultiProofTargetsV2; +use std::sync::Arc; +use tracing::trace; /// Source of state changes, either from EVM execution or from a Block Access List. #[derive(Clone, Copy)] @@ -53,57 +36,15 @@ impl From for Source { } } -/// Maximum number of targets to batch together for prefetch batching. -/// Prefetches are just proof requests (no state merging), so we allow a higher cap than state -/// updates -const PREFETCH_MAX_BATCH_TARGETS: usize = 512; - -/// Maximum number of prefetch messages to batch together. -/// Prevents excessive batching even with small messages. -const PREFETCH_MAX_BATCH_MESSAGES: usize = 16; - /// The default max targets, for limiting the number of account and storage proof targets to be /// fetched by a single worker. If exceeded, chunking is forced regardless of worker availability. pub(crate) const DEFAULT_MAX_TARGETS_FOR_CHUNKING: usize = 300; -/// A trie update that can be applied to sparse trie alongside the proofs for touched parts of the -/// state. -#[derive(Debug)] -pub struct SparseTrieUpdate { - /// The state update that was used to calculate the proof - pub(crate) state: HashedPostState, - /// The calculated multiproof - pub(crate) multiproof: ProofResult, -} - -impl SparseTrieUpdate { - /// Returns true if the update is empty. - pub(super) fn is_empty(&self) -> bool { - self.state.is_empty() && self.multiproof.is_empty() - } - - /// Construct update from multiproof. - #[cfg(test)] - pub(super) fn from_multiproof(multiproof: reth_trie::MultiProof) -> alloy_rlp::Result { - let stats = ParallelTrieTracker::default().finish(); - Ok(Self { - state: HashedPostState::default(), - multiproof: ProofResult::Legacy(multiproof.try_into()?, stats), - }) - } - - /// Extend update with contents of the other. - pub(super) fn extend(&mut self, other: Self) { - self.state.extend(other.state); - self.multiproof.extend(other.multiproof); - } -} - /// Messages used internally by the multi proof task. #[derive(Debug)] pub enum MultiProofMessage { /// Prefetch proof targets - PrefetchProofs(VersionedMultiProofTargets), + PrefetchProofs(MultiProofTargetsV2), /// New state update from transaction execution with its source StateUpdate(Source, EvmState), /// State update that can be applied to the sparse trie without any new proofs. @@ -122,7 +63,7 @@ pub enum MultiProofMessage { /// /// When received, the task generates a single state update from the BAL and processes it. /// No further messages are expected after receiving this variant. - BlockAccessList(Arc), + BlockAccessList(Arc), /// Signals state update stream end. /// /// This is triggered by block execution, indicating that no additional state updates are @@ -130,57 +71,6 @@ pub enum MultiProofMessage { FinishedStateUpdates, } -/// Handle to track proof calculation ordering. -#[derive(Debug, Default)] -struct ProofSequencer { - /// The next proof sequence number to be produced. - next_sequence: u64, - /// The next sequence number expected to be delivered. - next_to_deliver: u64, - /// Buffer for out-of-order proofs and corresponding state updates - pending_proofs: BTreeMap, -} - -impl ProofSequencer { - /// Gets the next sequence number and increments the counter - const fn next_sequence(&mut self) -> u64 { - let seq = self.next_sequence; - self.next_sequence += 1; - seq - } - - /// Adds a proof with the corresponding state update and returns all sequential proofs and state - /// updates if we have a continuous sequence - fn add_proof(&mut self, sequence: u64, update: SparseTrieUpdate) -> Vec { - // Optimization: fast path for in-order delivery to avoid BTreeMap overhead. - // If this is the expected sequence, return it immediately without buffering. - if sequence == self.next_to_deliver { - let mut consecutive_proofs = Vec::with_capacity(1); - consecutive_proofs.push(update); - self.next_to_deliver += 1; - - // Check if we have subsequent proofs in the pending buffer - while let Some(pending) = self.pending_proofs.remove(&self.next_to_deliver) { - consecutive_proofs.push(pending); - self.next_to_deliver += 1; - } - - return consecutive_proofs; - } - - if sequence > self.next_to_deliver { - self.pending_proofs.insert(sequence, update); - } - - Vec::new() - } - - /// Returns true if we still have pending proofs - pub(crate) fn has_pending(&self) -> bool { - !self.pending_proofs.is_empty() - } -} - /// A wrapper for the sender that signals completion when dropped. /// /// This type is intended to be used in combination with the evm executor statehook. @@ -234,303 +124,6 @@ pub(crate) fn evm_state_to_hashed_post_state(update: EvmState) -> HashedPostStat hashed_state } -/// Extends a `MultiProofTargets` with the contents of a `VersionedMultiProofTargets`, -/// regardless of which variant the latter is. -fn extend_multiproof_targets(dest: &mut MultiProofTargets, src: &VersionedMultiProofTargets) { - match src { - VersionedMultiProofTargets::Legacy(targets) => { - dest.extend_ref(targets); - } - VersionedMultiProofTargets::V2(targets) => { - // Add all account targets - for target in &targets.account_targets { - dest.entry(target.key()).or_default(); - } - - // Add all storage targets - for (hashed_address, slots) in &targets.storage_targets { - let slot_set = dest.entry(*hashed_address).or_default(); - for slot in slots { - slot_set.insert(slot.key()); - } - } - } - } -} - -/// A set of multiproof targets which can be either in the legacy or V2 representations. -#[derive(Debug)] -pub enum VersionedMultiProofTargets { - /// Legacy targets - Legacy(MultiProofTargets), - /// V2 targets - V2(MultiProofTargetsV2), -} - -impl VersionedMultiProofTargets { - /// Returns true if there are no account or storage targets. - fn is_empty(&self) -> bool { - match self { - Self::Legacy(targets) => targets.is_empty(), - Self::V2(targets) => targets.is_empty(), - } - } - - /// Returns the number of account targets in the multiproof target - fn account_targets_len(&self) -> usize { - match self { - Self::Legacy(targets) => targets.len(), - Self::V2(targets) => targets.account_targets.len(), - } - } - - /// Returns the number of storage targets in the multiproof target - fn storage_targets_len(&self) -> usize { - match self { - Self::Legacy(targets) => targets.values().map(|slots| slots.len()).sum::(), - Self::V2(targets) => { - targets.storage_targets.values().map(|slots| slots.len()).sum::() - } - } - } - - /// Returns the number of accounts in the multiproof targets. - fn len(&self) -> usize { - match self { - Self::Legacy(targets) => targets.len(), - Self::V2(targets) => targets.account_targets.len(), - } - } - - /// Returns the total storage slot count across all accounts. - fn storage_count(&self) -> usize { - match self { - Self::Legacy(targets) => targets.values().map(|slots| slots.len()).sum(), - Self::V2(targets) => targets.storage_targets.values().map(|slots| slots.len()).sum(), - } - } - - /// Returns the number of items that will be considered during chunking. - fn chunking_length(&self) -> usize { - match self { - Self::Legacy(targets) => targets.chunking_length(), - Self::V2(targets) => targets.chunking_length(), - } - } - - /// Retains the targets representing the difference with another `MultiProofTargets`. - /// Removes all targets that are already present in `other`. - fn retain_difference(&mut self, other: &MultiProofTargets) { - match self { - Self::Legacy(targets) => { - targets.retain_difference(other); - } - Self::V2(targets) => { - // Remove account targets that exist in other - targets.account_targets.retain(|target| !other.contains_key(&target.key())); - - // For each account in storage_targets, remove slots that exist in other - targets.storage_targets.retain(|hashed_address, slots| { - if let Some(other_slots) = other.get(hashed_address) { - slots.retain(|slot| !other_slots.contains(&slot.key())); - !slots.is_empty() - } else { - true - } - }); - } - } - } - - /// Extends this `VersionedMultiProofTargets` with the contents of another. - /// - /// Panics if the variants do not match. - fn extend(&mut self, other: Self) { - match (self, other) { - (Self::Legacy(dest), Self::Legacy(src)) => { - dest.extend(src); - } - (Self::V2(dest), Self::V2(src)) => { - dest.account_targets.extend(src.account_targets); - for (addr, slots) in src.storage_targets { - dest.storage_targets.entry(addr).or_default().extend(slots); - } - } - _ => panic!("Cannot extend VersionedMultiProofTargets with mismatched variants"), - } - } - - /// Chunks this `VersionedMultiProofTargets` into smaller chunks of the given size. - fn chunks(self, chunk_size: usize) -> Box> { - match self { - Self::Legacy(targets) => { - Box::new(MultiProofTargets::chunks(targets, chunk_size).map(Self::Legacy)) - } - Self::V2(targets) => Box::new(targets.chunks(chunk_size).map(Self::V2)), - } - } -} - -/// Input parameters for dispatching a multiproof calculation. -#[derive(Debug)] -struct MultiproofInput { - source: Option, - hashed_state_update: HashedPostState, - proof_targets: VersionedMultiProofTargets, - proof_sequence_number: u64, - state_root_message_sender: CrossbeamSender, - multi_added_removed_keys: Option>, -} - -impl MultiproofInput { - /// Destroys the input and sends a [`MultiProofMessage::EmptyProof`] message to the sender. - fn send_empty_proof(self) { - let _ = self.state_root_message_sender.send(MultiProofMessage::EmptyProof { - sequence_number: self.proof_sequence_number, - state: self.hashed_state_update, - }); - } -} - -/// Coordinates multiproof dispatch between `MultiProofTask` and the parallel trie workers. -/// -/// # Flow -/// 1. `MultiProofTask` asks the manager to dispatch either storage or account proof work. -/// 2. The manager builds the request, clones `proof_result_tx`, and hands everything to -/// [`ProofWorkerHandle`]. -/// 3. A worker finishes the proof and sends a [`ProofResultMessage`] through the channel included -/// in the job. -/// 4. `MultiProofTask` consumes the message from the same channel and sequences it with -/// `ProofSequencer`. -#[derive(Debug)] -pub struct MultiproofManager { - /// Handle to the proof worker pools (storage and account). - proof_worker_handle: ProofWorkerHandle, - /// Channel sender cloned into each dispatched job so workers can send back the - /// `ProofResultMessage`. - proof_result_tx: CrossbeamSender, - /// Metrics - metrics: MultiProofTaskMetrics, -} - -impl MultiproofManager { - /// Creates a new [`MultiproofManager`]. - fn new( - metrics: MultiProofTaskMetrics, - proof_worker_handle: ProofWorkerHandle, - proof_result_tx: CrossbeamSender, - ) -> Self { - // Initialize the max worker gauges with the worker pool sizes - metrics.max_storage_workers.set(proof_worker_handle.total_storage_workers() as f64); - metrics.max_account_workers.set(proof_worker_handle.total_account_workers() as f64); - - Self { metrics, proof_worker_handle, proof_result_tx } - } - - /// Dispatches a new multiproof calculation to worker pools. - fn dispatch(&self, input: MultiproofInput) { - // If there are no proof targets, we can just send an empty multiproof back immediately - if input.proof_targets.is_empty() { - trace!( - sequence_number = input.proof_sequence_number, - "No proof targets, sending empty multiproof back immediately" - ); - input.send_empty_proof(); - return; - } - - self.dispatch_multiproof(input); - } - - /// Signals that a multiproof calculation has finished. - fn on_calculation_complete(&self) { - self.metrics - .active_storage_workers_histogram - .record(self.proof_worker_handle.active_storage_workers() as f64); - self.metrics - .active_account_workers_histogram - .record(self.proof_worker_handle.active_account_workers() as f64); - self.metrics - .pending_storage_multiproofs_histogram - .record(self.proof_worker_handle.pending_storage_tasks() as f64); - self.metrics - .pending_account_multiproofs_histogram - .record(self.proof_worker_handle.pending_account_tasks() as f64); - } - - /// Dispatches a single multiproof calculation to worker pool. - fn dispatch_multiproof(&self, multiproof_input: MultiproofInput) { - let MultiproofInput { - source, - hashed_state_update, - proof_targets, - proof_sequence_number, - state_root_message_sender: _, - multi_added_removed_keys, - } = multiproof_input; - - trace!( - target: "engine::tree::payload_processor::multiproof", - proof_sequence_number, - ?proof_targets, - account_targets = proof_targets.account_targets_len(), - storage_targets = proof_targets.storage_targets_len(), - ?source, - "Dispatching multiproof to workers" - ); - - let start = Instant::now(); - - // Workers will send ProofResultMessage directly to proof_result_rx - let proof_result_sender = ProofResultContext::new( - self.proof_result_tx.clone(), - proof_sequence_number, - hashed_state_update, - start, - ); - - let input = match proof_targets { - VersionedMultiProofTargets::Legacy(proof_targets) => { - // Extend prefix sets with targets - let frozen_prefix_sets = ParallelProof::extend_prefix_sets_with_targets( - &Default::default(), - &proof_targets, - ); - - AccountMultiproofInput::Legacy { - targets: proof_targets, - prefix_sets: frozen_prefix_sets, - collect_branch_node_masks: true, - multi_added_removed_keys, - proof_result_sender, - } - } - VersionedMultiProofTargets::V2(proof_targets) => { - AccountMultiproofInput::V2 { targets: proof_targets, proof_result_sender } - } - }; - - // Dispatch account multiproof to worker pool with result sender - if let Err(e) = self.proof_worker_handle.dispatch_account_multiproof(input) { - error!(target: "engine::tree::payload_processor::multiproof", ?e, "Failed to dispatch account multiproof"); - return; - } - - self.metrics - .active_storage_workers_histogram - .record(self.proof_worker_handle.active_storage_workers() as f64); - self.metrics - .active_account_workers_histogram - .record(self.proof_worker_handle.active_account_workers() as f64); - self.metrics - .pending_storage_multiproofs_histogram - .record(self.proof_worker_handle.pending_storage_tasks() as f64); - self.metrics - .pending_account_multiproofs_histogram - .record(self.proof_worker_handle.pending_account_tasks() as f64); - } -} - #[derive(Metrics, Clone)] #[metrics(scope = "tree.root")] pub(crate) struct MultiProofTaskMetrics { @@ -590,920 +183,6 @@ pub(crate) struct MultiProofTaskMetrics { pub sparse_trie_cache_wait_duration_histogram: Histogram, } -/// Standalone task that receives a transaction state stream and updates relevant -/// data structures to calculate state root. -/// -/// ## Architecture: Dual-Channel Multiproof System -/// -/// This task orchestrates parallel proof computation using a dual-channel architecture that -/// separates control messages from proof computation results: -/// -/// ```text -/// ┌─────────────────────────────────────────────────────────────────┐ -/// │ MultiProofTask │ -/// │ Event Loop (crossbeam::select!) │ -/// └──┬──────────────────────────────────────────────────────────▲───┘ -/// │ │ -/// │ (1) Send proof request │ -/// │ via tx (control channel) │ -/// │ │ -/// ▼ │ -/// ┌──────────────────────────────────────────────────────────────┐ │ -/// │ MultiproofManager │ │ -/// │ - Deduplicates against fetched_proof_targets │ │ -/// │ - Routes to appropriate worker pool │ │ -/// └──┬───────────────────────────────────────────────────────────┘ │ -/// │ │ -/// │ (2) Dispatch to workers │ -/// │ OR send EmptyProof (fast path) │ -/// ▼ │ -/// ┌──────────────────────────────────────────────────────────────┐ │ -/// │ ProofWorkerHandle │ │ -/// │ ┌─────────────────────┐ ┌────────────────────────┐ │ │ -/// │ │ Storage Worker Pool │ │ Account Worker Pool │ │ │ -/// │ │ (spawn_blocking) │ │ (spawn_blocking) │ │ │ -/// │ └─────────────────────┘ └────────────────────────┘ │ │ -/// └──┬───────────────────────────────────────────────────────────┘ │ -/// │ │ -/// │ (3) Compute proofs in parallel │ -/// │ Send results back │ -/// │ │ -/// ▼ │ -/// ┌──────────────────────────────────────────────────────────────┐ │ -/// │ proof_result_tx (crossbeam unbounded channel) │ │ -/// │ → ProofResultMessage { multiproof, sequence_number, ... } │ │ -/// └──────────────────────────────────────────────────────────────┘ │ -/// │ -/// (4) Receive via crossbeam::select! on two channels: ───────────┘ -/// - rx: Control messages (PrefetchProofs, StateUpdate, -/// EmptyProof, FinishedStateUpdates) -/// - proof_result_rx: Computed proof results from workers -/// ``` -/// -/// ## Component Responsibilities -/// -/// - **[`MultiProofTask`]**: Event loop coordinator -/// - Receives state updates from transaction execution -/// - Deduplicates proof targets against already-fetched proofs -/// - Sequences proofs to maintain transaction ordering -/// - Feeds sequenced updates to sparse trie task -/// -/// - **[`MultiproofManager`]**: Calculation orchestrator -/// - Decides between fast path ([`EmptyProof`]) and worker dispatch -/// - Routes storage-only vs full multiproofs to appropriate workers -/// - Records metrics for monitoring -/// -/// - **[`ProofWorkerHandle`]**: Worker pool manager -/// - Maintains separate pools for storage and account proofs -/// - Dispatches work to blocking threads (CPU-intensive) -/// - Sends results directly via `proof_result_tx` (bypasses control channel) -/// -/// [`EmptyProof`]: MultiProofMessage::EmptyProof -/// [`ProofWorkerHandle`]: reth_trie_parallel::proof_task::ProofWorkerHandle -/// -/// ## Dual-Channel Design Rationale -/// -/// The system uses two separate crossbeam channels: -/// -/// 1. **Control Channel (`tx`/`rx`)**: For orchestration messages -/// - `PrefetchProofs`: Pre-fetch proofs before execution -/// - `StateUpdate`: New transaction execution results -/// - `EmptyProof`: Fast path when all targets already fetched -/// - `FinishedStateUpdates`: Signal to drain pending work -/// -/// 2. **Proof Result Channel (`proof_result_tx`/`proof_result_rx`)**: For worker results -/// - `ProofResultMessage`: Computed multiproofs from worker pools -/// - Direct path from workers to event loop (no intermediate hops) -/// - Keeps control messages separate from high-throughput proof data -/// -/// This separation enables: -/// - **Non-blocking control**: Control messages never wait behind large proof data -/// - **Backpressure management**: Each channel can apply different policies -/// - **Clear ownership**: Workers only need proof result sender, not control channel -/// -/// ## Initialization and Lifecycle -/// -/// The task initializes a blinded sparse trie and subscribes to transaction state streams. -/// As it receives transaction execution results, it fetches proofs for relevant accounts -/// from the database and reveals them to the tree, then updates relevant leaves according -/// to transaction results. This feeds updates to the sparse trie task. -/// -/// See the `run()` method documentation for detailed lifecycle flow. -#[derive(Debug)] -pub(super) struct MultiProofTask { - /// The size of proof targets chunk to spawn in one calculation. - /// If None, chunking is disabled and all targets are processed in a single proof. - chunk_size: Option, - /// Receiver for state root related messages (prefetch, state updates, finish signal). - rx: CrossbeamReceiver, - /// Sender for state root related messages. - tx: CrossbeamSender, - /// Receiver for proof results directly from workers. - proof_result_rx: CrossbeamReceiver, - /// Sender for state updates emitted by this type. - to_sparse_trie: std::sync::mpsc::Sender, - /// Proof targets that have been already fetched. - fetched_proof_targets: MultiProofTargets, - /// Tracks keys which have been added and removed throughout the entire block. - multi_added_removed_keys: MultiAddedRemovedKeys, - /// Proof sequencing handler. - proof_sequencer: ProofSequencer, - /// Manages calculation of multiproofs. - multiproof_manager: MultiproofManager, - /// multi proof task metrics - metrics: MultiProofTaskMetrics, - /// If this number is exceeded and chunking is enabled, then this will override whether or not - /// there are any active workers and force chunking across workers. This is to prevent tasks - /// which are very long from hitting a single worker. - max_targets_for_chunking: usize, - /// Whether or not V2 proof calculation is enabled. If enabled then [`MultiProofTargetsV2`] - /// will be produced by state updates. - v2_proofs_enabled: bool, -} - -impl MultiProofTask { - /// Creates a multiproof task with separate channels: control on `tx`/`rx`, proof results on - /// `proof_result_rx`. - pub(super) fn new( - proof_worker_handle: ProofWorkerHandle, - to_sparse_trie: std::sync::mpsc::Sender, - chunk_size: Option, - tx: CrossbeamSender, - rx: CrossbeamReceiver, - ) -> Self { - let (proof_result_tx, proof_result_rx) = unbounded(); - let metrics = MultiProofTaskMetrics::default(); - - Self { - chunk_size, - rx, - tx, - proof_result_rx, - to_sparse_trie, - fetched_proof_targets: Default::default(), - multi_added_removed_keys: MultiAddedRemovedKeys::new(), - proof_sequencer: ProofSequencer::default(), - multiproof_manager: MultiproofManager::new( - metrics.clone(), - proof_worker_handle, - proof_result_tx, - ), - metrics, - max_targets_for_chunking: DEFAULT_MAX_TARGETS_FOR_CHUNKING, - v2_proofs_enabled: false, - } - } - - /// Enables V2 proof target generation on state updates. - pub(super) const fn with_v2_proofs_enabled(mut self, v2_proofs_enabled: bool) -> Self { - self.v2_proofs_enabled = v2_proofs_enabled; - self - } - - /// Handles request for proof prefetch. - /// - /// Returns how many multiproof tasks were dispatched for the prefetch request. - #[instrument( - level = "debug", - target = "engine::tree::payload_processor::multiproof", - skip_all, - fields(accounts = targets.account_targets_len(), chunks = 0) - )] - fn on_prefetch_proof(&mut self, mut targets: VersionedMultiProofTargets) -> u64 { - // Remove already fetched proof targets to avoid redundant work. - targets.retain_difference(&self.fetched_proof_targets); - - if targets.is_empty() { - return 0; - } - - extend_multiproof_targets(&mut self.fetched_proof_targets, &targets); - - // For Legacy multiproofs, make sure all target accounts have an `AddedRemovedKeySet` in the - // [`MultiAddedRemovedKeys`]. Even if there are not any known removed keys for the account, - // we still want to optimistically fetch extension children for the leaf addition case. - // V2 multiproofs don't need this. - // - // Only clone the AddedRemovedKeys for accounts in the targets, not the entire accumulated - // set, to avoid O(n) cloning with many buffered blocks. - let multi_added_removed_keys = - if let VersionedMultiProofTargets::Legacy(legacy_targets) = &targets { - self.multi_added_removed_keys.touch_accounts(legacy_targets.keys().copied()); - Some(Arc::new(MultiAddedRemovedKeys { - account: self.multi_added_removed_keys.account.clone(), - storages: legacy_targets - .keys() - .filter_map(|k| { - self.multi_added_removed_keys.storages.get(k).map(|v| (*k, v.clone())) - }) - .collect(), - })) - } else { - None - }; - - self.metrics.prefetch_proof_targets_accounts_histogram.record(targets.len() as f64); - self.metrics - .prefetch_proof_targets_storages_histogram - .record(targets.storage_count() as f64); - - let chunking_len = targets.chunking_length(); - let available_account_workers = - self.multiproof_manager.proof_worker_handle.available_account_workers(); - let available_storage_workers = - self.multiproof_manager.proof_worker_handle.available_storage_workers(); - let num_chunks = dispatch_with_chunking( - targets, - chunking_len, - self.chunk_size, - self.max_targets_for_chunking, - available_account_workers, - available_storage_workers, - VersionedMultiProofTargets::chunks, - |proof_targets| { - self.multiproof_manager.dispatch(MultiproofInput { - source: None, - hashed_state_update: Default::default(), - proof_targets, - proof_sequence_number: self.proof_sequencer.next_sequence(), - state_root_message_sender: self.tx.clone(), - multi_added_removed_keys: multi_added_removed_keys.clone(), - }); - }, - ); - self.metrics.prefetch_proof_chunks_histogram.record(num_chunks as f64); - - num_chunks as u64 - } - - /// Returns true if all state updates finished and all pending proofs processed. - fn is_done(&self, metrics: &MultiproofBatchMetrics, ctx: &MultiproofBatchCtx) -> bool { - let all_proofs_processed = metrics.all_proofs_processed(); - let no_pending = !self.proof_sequencer.has_pending(); - let updates_finished = ctx.updates_finished(); - trace!( - target: "engine::tree::payload_processor::multiproof", - proofs_processed = metrics.proofs_processed, - state_update_proofs_requested = metrics.state_update_proofs_requested, - prefetch_proofs_requested = metrics.prefetch_proofs_requested, - no_pending, - updates_finished, - "Checking end condition" - ); - all_proofs_processed && no_pending && updates_finished - } - - /// Handles state updates. - /// - /// Returns how many proof dispatches were spawned (including an `EmptyProof` for already - /// fetched targets). - #[instrument( - level = "debug", - target = "engine::tree::payload_processor::multiproof", - skip(self, update), - fields(accounts = update.len(), chunks = 0) - )] - fn on_state_update(&mut self, source: Source, update: EvmState) -> u64 { - let hashed_state_update = evm_state_to_hashed_post_state(update); - self.on_hashed_state_update(source, hashed_state_update) - } - - /// Processes a hashed state update and dispatches multiproofs as needed. - /// - /// Returns the number of state updates dispatched (both `EmptyProof` and regular multiproofs). - fn on_hashed_state_update( - &mut self, - source: Source, - hashed_state_update: HashedPostState, - ) -> u64 { - // Update removed keys based on the state update. - self.multi_added_removed_keys.update_with_state(&hashed_state_update); - - // Split the state update into already fetched and not fetched according to the proof - // targets. - let (fetched_state_update, not_fetched_state_update) = hashed_state_update - .partition_by_targets(&self.fetched_proof_targets, &self.multi_added_removed_keys); - - let mut state_updates = 0; - // If there are any accounts or storage slots that we already fetched the proofs for, - // send them immediately, as they don't require dispatching any additional multiproofs. - if !fetched_state_update.is_empty() { - let _ = self.tx.send(MultiProofMessage::EmptyProof { - sequence_number: self.proof_sequencer.next_sequence(), - state: fetched_state_update, - }); - state_updates += 1; - } - - if not_fetched_state_update.is_empty() { - return state_updates; - } - - // Clone+Arc MultiAddedRemovedKeys for sharing with the dispatched multiproof tasks - let multi_added_removed_keys = Arc::new(MultiAddedRemovedKeys { - account: self.multi_added_removed_keys.account.clone(), - storages: { - let mut storages = B256Map::with_capacity_and_hasher( - not_fetched_state_update.storages.len(), - Default::default(), - ); - - for account in not_fetched_state_update - .storages - .keys() - .chain(not_fetched_state_update.accounts.keys()) - { - if let hash_map::Entry::Vacant(entry) = storages.entry(*account) { - entry.insert( - self.multi_added_removed_keys - .storages - .get(account) - .cloned() - .unwrap_or_else(default_added_removed_keys), - ); - } - } - - storages - }, - }); - - let chunking_len = not_fetched_state_update.chunking_length(); - let mut spawned_proof_targets = MultiProofTargets::default(); - let available_account_workers = - self.multiproof_manager.proof_worker_handle.available_account_workers(); - let available_storage_workers = - self.multiproof_manager.proof_worker_handle.available_storage_workers(); - - let num_chunks = dispatch_with_chunking( - not_fetched_state_update, - chunking_len, - self.chunk_size, - self.max_targets_for_chunking, - available_account_workers, - available_storage_workers, - HashedPostState::chunks, - |hashed_state_update| { - let proof_targets = get_proof_targets( - &hashed_state_update, - &self.fetched_proof_targets, - &multi_added_removed_keys, - self.v2_proofs_enabled, - ); - extend_multiproof_targets(&mut spawned_proof_targets, &proof_targets); - - self.multiproof_manager.dispatch(MultiproofInput { - source: Some(source), - hashed_state_update, - proof_targets, - proof_sequence_number: self.proof_sequencer.next_sequence(), - state_root_message_sender: self.tx.clone(), - multi_added_removed_keys: Some(multi_added_removed_keys.clone()), - }); - }, - ); - self.metrics - .state_update_proof_targets_accounts_histogram - .record(spawned_proof_targets.len() as f64); - self.metrics - .state_update_proof_targets_storages_histogram - .record(spawned_proof_targets.values().map(|slots| slots.len()).sum::() as f64); - self.metrics.state_update_proof_chunks_histogram.record(num_chunks as f64); - - self.fetched_proof_targets.extend(spawned_proof_targets); - - state_updates + num_chunks as u64 - } - - /// Handler for new proof calculated, aggregates all the existing sequential proofs. - fn on_proof( - &mut self, - sequence_number: u64, - update: SparseTrieUpdate, - ) -> Option { - let ready_proofs = self.proof_sequencer.add_proof(sequence_number, update); - - ready_proofs - .into_iter() - // Merge all ready proofs and state updates - .reduce(|mut acc_update, update| { - acc_update.extend(update); - acc_update - }) - // Return None if the resulting proof is empty - .filter(|proof| !proof.is_empty()) - } - - /// Processes a multiproof message, batching consecutive prefetch messages. - /// - /// For prefetch messages, drains queued prefetch messages and merges them into one batch before - /// processing, storing one pending message (different type or over-cap) to handle on the next - /// iteration. State updates are processed directly without batching. - /// - /// Returns `true` if done, `false` to continue. - fn process_multiproof_message

( - &mut self, - msg: MultiProofMessage, - ctx: &mut MultiproofBatchCtx, - batch_metrics: &mut MultiproofBatchMetrics, - provider: &P, - ) -> bool - where - P: AccountReader, - { - match msg { - // Prefetch proofs: batch consecutive prefetch requests up to target/message limits - MultiProofMessage::PrefetchProofs(targets) => { - trace!(target: "engine::tree::payload_processor::multiproof", "processing MultiProofMessage::PrefetchProofs"); - - if ctx.first_update_time.is_none() { - self.metrics - .first_update_wait_time_histogram - .record(ctx.start.elapsed().as_secs_f64()); - ctx.first_update_time = Some(Instant::now()); - debug!(target: "engine::tree::payload_processor::multiproof", "Started state root calculation"); - } - - let mut accumulated_count = targets.chunking_length(); - ctx.accumulated_prefetch_targets.clear(); - ctx.accumulated_prefetch_targets.push(targets); - - // Batch consecutive prefetch messages up to limits. - // EmptyProof messages are handled inline since they're very fast (~100ns) - // and shouldn't interrupt batching. - while accumulated_count < PREFETCH_MAX_BATCH_TARGETS && - ctx.accumulated_prefetch_targets.len() < PREFETCH_MAX_BATCH_MESSAGES - { - match self.rx.try_recv() { - Ok(MultiProofMessage::PrefetchProofs(next_targets)) => { - let next_count = next_targets.chunking_length(); - if accumulated_count + next_count > PREFETCH_MAX_BATCH_TARGETS { - ctx.pending_msg = - Some(MultiProofMessage::PrefetchProofs(next_targets)); - break; - } - accumulated_count += next_count; - ctx.accumulated_prefetch_targets.push(next_targets); - } - Ok(MultiProofMessage::EmptyProof { sequence_number, state }) => { - // Handle inline - very fast, don't break batching - batch_metrics.proofs_processed += 1; - if let Some(combined_update) = self.on_proof( - sequence_number, - SparseTrieUpdate { - state, - multiproof: ProofResult::empty(self.v2_proofs_enabled), - }, - ) { - let _ = self.to_sparse_trie.send(combined_update); - } - } - Ok(other_msg) => { - ctx.pending_msg = Some(other_msg); - break; - } - Err(_) => break, - } - } - - // Process all accumulated messages in a single batch - let num_batched = ctx.accumulated_prefetch_targets.len(); - self.metrics.prefetch_batch_size_histogram.record(num_batched as f64); - - // Merge all accumulated prefetch targets into a single dispatch payload. - // Use drain to preserve the buffer allocation. - let mut accumulated_iter = ctx.accumulated_prefetch_targets.drain(..); - let mut merged_targets = - accumulated_iter.next().expect("prefetch batch always has at least one entry"); - for next_targets in accumulated_iter { - merged_targets.extend(next_targets); - } - - let account_targets = merged_targets.len(); - let storage_targets = merged_targets.storage_count(); - batch_metrics.prefetch_proofs_requested += self.on_prefetch_proof(merged_targets); - trace!( - target: "engine::tree::payload_processor::multiproof", - account_targets, - storage_targets, - prefetch_proofs_requested = batch_metrics.prefetch_proofs_requested, - num_batched, - "Dispatched prefetch batch" - ); - - false - } - MultiProofMessage::StateUpdate(source, update) => { - trace!(target: "engine::tree::payload_processor::multiproof", "processing MultiProofMessage::StateUpdate"); - - if ctx.first_update_time.is_none() { - self.metrics - .first_update_wait_time_histogram - .record(ctx.start.elapsed().as_secs_f64()); - ctx.first_update_time = Some(Instant::now()); - debug!(target: "engine::tree::payload_processor::multiproof", "Started state root calculation"); - } - - let update_len = update.len(); - batch_metrics.state_update_proofs_requested += self.on_state_update(source, update); - trace!( - target: "engine::tree::payload_processor::multiproof", - ?source, - len = update_len, - state_update_proofs_requested = ?batch_metrics.state_update_proofs_requested, - "Dispatched state update" - ); - - false - } - // Process Block Access List (BAL) - complete state changes provided upfront - MultiProofMessage::BlockAccessList(bal) => { - trace!(target: "engine::tree::payload_processor::multiproof", "processing MultiProofMessage::BAL"); - - if ctx.first_update_time.is_none() { - self.metrics - .first_update_wait_time_histogram - .record(ctx.start.elapsed().as_secs_f64()); - ctx.first_update_time = Some(Instant::now()); - debug!(target: "engine::tree::payload_processor::multiproof", "Started state root calculation from BAL"); - } - - // Convert BAL to HashedPostState and process it - match bal_to_hashed_post_state(&bal, provider) { - Ok(hashed_state) => { - debug!( - target: "engine::tree::payload_processor::multiproof", - accounts = hashed_state.accounts.len(), - storages = hashed_state.storages.len(), - "Processing BAL state update" - ); - - // Use BlockAccessList as source for BAL-derived state updates - batch_metrics.state_update_proofs_requested += - self.on_hashed_state_update(Source::BlockAccessList, hashed_state); - } - Err(err) => { - error!(target: "engine::tree::payload_processor::multiproof", ?err, "Failed to convert BAL to hashed state"); - return true; - } - } - - // Mark updates as finished since BAL provides complete state - ctx.updates_finished_time = Some(Instant::now()); - - // Check if we're done (might need to wait for proofs to complete) - if self.is_done(batch_metrics, ctx) { - debug!( - target: "engine::tree::payload_processor::multiproof", - "BAL processed and all proofs complete, ending calculation" - ); - return true; - } - false - } - // Signal that no more state updates will arrive - MultiProofMessage::FinishedStateUpdates => { - trace!(target: "engine::tree::payload_processor::multiproof", "processing MultiProofMessage::FinishedStateUpdates"); - - ctx.updates_finished_time = Some(Instant::now()); - - if self.is_done(batch_metrics, ctx) { - debug!( - target: "engine::tree::payload_processor::multiproof", - "State updates finished and all proofs processed, ending calculation" - ); - return true; - } - false - } - // Handle proof result with no trie nodes (state unchanged) - MultiProofMessage::EmptyProof { sequence_number, state } => { - trace!(target: "engine::tree::payload_processor::multiproof", "processing MultiProofMessage::EmptyProof"); - - batch_metrics.proofs_processed += 1; - - if let Some(combined_update) = self.on_proof( - sequence_number, - SparseTrieUpdate { - state, - multiproof: ProofResult::empty(self.v2_proofs_enabled), - }, - ) { - let _ = self.to_sparse_trie.send(combined_update); - } - - if self.is_done(batch_metrics, ctx) { - debug!( - target: "engine::tree::payload_processor::multiproof", - "State updates finished and all proofs processed, ending calculation" - ); - return true; - } - false - } - MultiProofMessage::HashedStateUpdate(hashed_state) => { - batch_metrics.state_update_proofs_requested += - self.on_hashed_state_update(Source::BlockAccessList, hashed_state); - false - } - } - } - - /// Starts the main loop that handles all incoming messages, fetches proofs, applies them to the - /// sparse trie, updates the sparse trie, and eventually returns the state root. - /// - /// The lifecycle is the following: - /// 1. Either [`MultiProofMessage::PrefetchProofs`] or [`MultiProofMessage::StateUpdate`] is - /// received from the engine. - /// * For [`MultiProofMessage::StateUpdate`], the state update is hashed with - /// [`evm_state_to_hashed_post_state`], and then (proof targets)[`MultiProofTargets`] are - /// extracted with [`get_proof_targets`]. - /// * For both messages, proof targets are deduplicated according to `fetched_proof_targets`, - /// so that the proofs for accounts and storage slots that were already fetched are not - /// requested again. - /// 2. Using the proof targets, a new multiproof is calculated using - /// [`MultiproofManager::dispatch`]. - /// * If the list of proof targets is empty, the [`MultiProofMessage::EmptyProof`] message is - /// sent back to this task along with the original state update. - /// * Otherwise, the multiproof is dispatched to worker pools and results are sent directly - /// to this task via the `proof_result_rx` channel as [`ProofResultMessage`]. - /// 3. Either [`MultiProofMessage::EmptyProof`] (via control channel) or [`ProofResultMessage`] - /// (via proof result channel) is received. - /// * The multiproof is added to the [`ProofSequencer`]. - /// * If the proof sequencer has a contiguous sequence of multiproofs in the same order as - /// state updates arrived (i.e. transaction order), such sequence is returned. - /// 4. Once there's a sequence of contiguous multiproofs along with the proof targets and state - /// updates associated with them, a [`SparseTrieUpdate`] is generated and sent to the sparse - /// trie task. - /// 5. Steps above are repeated until this task receives a - /// [`MultiProofMessage::FinishedStateUpdates`]. - /// * Once this message is received, on every [`MultiProofMessage::EmptyProof`] and - /// [`ProofResultMessage`], we check if all proofs have been processed and if there are any - /// pending proofs in the proof sequencer left to be revealed. - /// 6. While running, consecutive [`MultiProofMessage::PrefetchProofs`] messages are batched to - /// reduce redundant work; if a different message type arrives mid-batch or a batch cap is - /// reached, it is held as `pending_msg` and processed on the next loop to preserve ordering. - /// 7. This task exits after all pending proofs are processed. - #[instrument( - level = "debug", - name = "MultiProofTask::run", - target = "engine::tree::payload_processor::multiproof", - skip_all - )] - pub(crate) fn run

(mut self, provider: P) - where - P: AccountReader, - { - let mut ctx = MultiproofBatchCtx::new(Instant::now()); - let mut batch_metrics = MultiproofBatchMetrics::default(); - - // Main event loop; select_biased! prioritizes proof results over control messages. - // Labeled so inner match arms can `break 'main` once all work is complete. - 'main: loop { - trace!(target: "engine::tree::payload_processor::multiproof", "entering main channel receiving loop"); - - if let Some(msg) = ctx.pending_msg.take() { - if self.process_multiproof_message(msg, &mut ctx, &mut batch_metrics, &provider) { - break 'main; - } - continue; - } - - // Use select_biased! to prioritize proof results over new requests. - // This prevents new work from starving completed proofs and keeps workers healthy. - crossbeam_channel::select_biased! { - recv(self.proof_result_rx) -> proof_msg => { - match proof_msg { - Ok(proof_result) => { - batch_metrics.proofs_processed += 1; - - self.metrics - .proof_calculation_duration_histogram - .record(proof_result.elapsed); - - self.multiproof_manager.on_calculation_complete(); - - // Convert ProofResultMessage to SparseTrieUpdate - match proof_result.result { - Ok(proof_result_data) => { - trace!( - target: "engine::tree::payload_processor::multiproof", - sequence = proof_result.sequence_number, - total_proofs = batch_metrics.proofs_processed, - "Processing calculated proof from worker" - ); - - let update = SparseTrieUpdate { - state: proof_result.state, - multiproof: proof_result_data, - }; - - if let Some(combined_update) = - self.on_proof(proof_result.sequence_number, update) - { - let _ = self.to_sparse_trie.send(combined_update); - } - } - Err(error) => { - error!(target: "engine::tree::payload_processor::multiproof", ?error, "proof calculation error from worker"); - return - } - } - - if self.is_done(&batch_metrics, &ctx) { - debug!( - target: "engine::tree::payload_processor::multiproof", - "State updates finished and all proofs processed, ending calculation" - ); - break 'main - } - } - Err(_) => { - error!(target: "engine::tree::payload_processor::multiproof", "Proof result channel closed unexpectedly"); - return - } - } - }, - recv(self.rx) -> message => { - let msg = match message { - Ok(m) => m, - Err(_) => { - error!(target: "engine::tree::payload_processor::multiproof", "State root related message channel closed unexpectedly"); - return - } - }; - - if self.process_multiproof_message(msg, &mut ctx, &mut batch_metrics, &provider) { - break 'main; - } - } - } - } - - debug!( - target: "engine::tree::payload_processor::multiproof", - total_updates = batch_metrics.state_update_proofs_requested, - total_proofs = batch_metrics.proofs_processed, - total_time = ?ctx.first_update_time.map(|t|t.elapsed()), - time_since_updates_finished = ?ctx.updates_finished_time.map(|t|t.elapsed()), - "All proofs processed, ending calculation" - ); - - // update total metrics on finish - self.metrics - .state_updates_received_histogram - .record(batch_metrics.state_update_proofs_requested as f64); - self.metrics.proofs_processed_histogram.record(batch_metrics.proofs_processed as f64); - if let Some(total_time) = ctx.first_update_time.map(|t| t.elapsed()) { - self.metrics.multiproof_task_total_duration_histogram.record(total_time); - } - - if let Some(updates_finished_time) = ctx.updates_finished_time { - self.metrics - .last_proof_wait_time_histogram - .record(updates_finished_time.elapsed().as_secs_f64()); - } - } -} - -/// Context for multiproof message batching loop. -/// -/// Contains processing state that persists across loop iterations. -/// -/// Used by `process_multiproof_message` to batch consecutive prefetch messages received via -/// `try_recv` for efficient processing. -struct MultiproofBatchCtx { - /// Buffers a non-matching message type encountered during batching. - /// Processed first in next iteration to preserve ordering while allowing same-type - /// messages to batch. - pending_msg: Option, - /// Timestamp when the first state update or prefetch was received. - first_update_time: Option, - /// Timestamp before the first state update or prefetch was received. - start: Instant, - /// Timestamp when state updates finished. `Some` indicates all state updates have been - /// received. - updates_finished_time: Option, - /// Reusable buffer for accumulating prefetch targets during batching. - accumulated_prefetch_targets: Vec, -} - -impl MultiproofBatchCtx { - /// Creates a new batch context with the given start time. - fn new(start: Instant) -> Self { - Self { - pending_msg: None, - first_update_time: None, - start, - updates_finished_time: None, - accumulated_prefetch_targets: Vec::with_capacity(PREFETCH_MAX_BATCH_MESSAGES), - } - } - - /// Returns `true` if all state updates have been received. - const fn updates_finished(&self) -> bool { - self.updates_finished_time.is_some() - } -} - -/// Counters for tracking proof requests and processing. -#[derive(Default)] -struct MultiproofBatchMetrics { - /// Number of proofs that have been processed. - proofs_processed: u64, - /// Number of state update proofs requested. - state_update_proofs_requested: u64, - /// Number of prefetch proofs requested. - prefetch_proofs_requested: u64, -} - -impl MultiproofBatchMetrics { - /// Returns `true` if all requested proofs have been processed. - const fn all_proofs_processed(&self) -> bool { - self.proofs_processed >= self.state_update_proofs_requested + self.prefetch_proofs_requested - } -} - -/// Returns accounts only with those storages that were not already fetched, and -/// if there are no such storages and the account itself was already fetched, the -/// account shouldn't be included. -fn get_proof_targets( - state_update: &HashedPostState, - fetched_proof_targets: &MultiProofTargets, - multi_added_removed_keys: &MultiAddedRemovedKeys, - v2_enabled: bool, -) -> VersionedMultiProofTargets { - if v2_enabled { - let mut targets = MultiProofTargetsV2::default(); - - // first collect all new accounts (not previously fetched) - for &hashed_address in state_update.accounts.keys() { - if !fetched_proof_targets.contains_key(&hashed_address) { - targets.account_targets.push(hashed_address.into()); - } - } - - // then process storage slots for all accounts in the state update - for (hashed_address, storage) in &state_update.storages { - let fetched = fetched_proof_targets.get(hashed_address); - - // If the storage is wiped, we still need to fetch the account proof. - if storage.wiped && fetched.is_none() { - targets.account_targets.push(Into::::into(*hashed_address)); - continue - } - - let changed_slots = storage - .storage - .keys() - .filter(|slot| !fetched.is_some_and(|f| f.contains(*slot))) - .map(|slot| Into::::into(*slot)) - .collect::>(); - - if !changed_slots.is_empty() { - targets.account_targets.push((*hashed_address).into()); - targets.storage_targets.insert(*hashed_address, changed_slots); - } - } - - VersionedMultiProofTargets::V2(targets) - } else { - let mut targets = MultiProofTargets::default(); - - // first collect all new accounts (not previously fetched) - for hashed_address in state_update.accounts.keys() { - if !fetched_proof_targets.contains_key(hashed_address) { - targets.insert(*hashed_address, HashSet::default()); - } - } - - // then process storage slots for all accounts in the state update - for (hashed_address, storage) in &state_update.storages { - let fetched = fetched_proof_targets.get(hashed_address); - let storage_added_removed_keys = multi_added_removed_keys.get_storage(hashed_address); - let mut changed_slots = storage - .storage - .keys() - .filter(|slot| { - !fetched.is_some_and(|f| f.contains(*slot)) || - storage_added_removed_keys.is_some_and(|k| k.is_removed(slot)) - }) - .peekable(); - - // If the storage is wiped, we still need to fetch the account proof. - if storage.wiped && fetched.is_none() { - targets.entry(*hashed_address).or_default(); - } - - if changed_slots.peek().is_some() { - targets.entry(*hashed_address).or_default().extend(changed_slots); - } - } - - VersionedMultiProofTargets::Legacy(targets) - } -} - /// Dispatches work items as a single unit or in chunks based on target size and worker /// availability. #[allow(clippy::too_many_arguments)] @@ -1539,796 +218,3 @@ where dispatch(items); 1 } - -#[cfg(test)] -mod tests { - use crate::tree::cached_state::CachedStateProvider; - - use super::*; - use alloy_eip7928::{AccountChanges, BalanceChange}; - use alloy_primitives::Address; - use reth_provider::{ - providers::OverlayStateProviderFactory, test_utils::create_test_provider_factory, - BlockNumReader, BlockReader, ChangeSetReader, DatabaseProviderFactory, LatestStateProvider, - PruneCheckpointReader, StageCheckpointReader, StateProviderBox, StorageChangeSetReader, - StorageSettingsCache, - }; - use reth_trie::MultiProof; - use reth_trie_db::ChangesetCache; - use reth_trie_parallel::proof_task::{ProofTaskCtx, ProofWorkerHandle}; - use revm_primitives::{B256, U256}; - use std::sync::{Arc, OnceLock}; - - /// Get a test runtime, creating it if necessary - fn get_test_runtime() -> &'static reth_tasks::Runtime { - static TEST_RT: OnceLock = OnceLock::new(); - TEST_RT.get_or_init(reth_tasks::Runtime::test) - } - - fn create_test_state_root_task(factory: F) -> MultiProofTask - where - F: DatabaseProviderFactory< - Provider: BlockReader - + StageCheckpointReader - + PruneCheckpointReader - + ChangeSetReader - + StorageChangeSetReader - + StorageSettingsCache - + BlockNumReader, - > + Clone - + Send - + 'static, - { - let runtime = get_test_runtime(); - let changeset_cache = ChangesetCache::new(); - let overlay_factory = OverlayStateProviderFactory::new(factory, changeset_cache); - let task_ctx = ProofTaskCtx::new(overlay_factory); - let proof_handle = ProofWorkerHandle::new(runtime, task_ctx, false, false); - let (to_sparse_trie, _receiver) = std::sync::mpsc::channel(); - let (tx, rx) = crossbeam_channel::unbounded(); - - MultiProofTask::new(proof_handle, to_sparse_trie, Some(1), tx, rx) - } - - fn create_cached_provider(factory: F) -> CachedStateProvider - where - F: DatabaseProviderFactory< - Provider: BlockReader - + StageCheckpointReader - + PruneCheckpointReader - + reth_provider::StorageSettingsCache, - > + Clone - + Send - + 'static, - { - let db_provider = factory.database_provider_ro().unwrap(); - let state_provider: StateProviderBox = Box::new(LatestStateProvider::new(db_provider)); - let cache = crate::tree::cached_state::ExecutionCache::new(1000); - CachedStateProvider::new(state_provider, cache, Default::default()) - } - - #[test] - fn test_add_proof_in_sequence() { - let mut sequencer = ProofSequencer::default(); - let proof1 = MultiProof::default(); - let proof2 = MultiProof::default(); - sequencer.next_sequence = 2; - - let ready = sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proof1).unwrap()); - assert_eq!(ready.len(), 1); - assert!(!sequencer.has_pending()); - - let ready = sequencer.add_proof(1, SparseTrieUpdate::from_multiproof(proof2).unwrap()); - assert_eq!(ready.len(), 1); - assert!(!sequencer.has_pending()); - } - - #[test] - fn test_add_proof_out_of_order() { - let mut sequencer = ProofSequencer::default(); - let proof1 = MultiProof::default(); - let proof2 = MultiProof::default(); - let proof3 = MultiProof::default(); - sequencer.next_sequence = 3; - - let ready = sequencer.add_proof(2, SparseTrieUpdate::from_multiproof(proof3).unwrap()); - assert_eq!(ready.len(), 0); - assert!(sequencer.has_pending()); - - let ready = sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proof1).unwrap()); - assert_eq!(ready.len(), 1); - assert!(sequencer.has_pending()); - - let ready = sequencer.add_proof(1, SparseTrieUpdate::from_multiproof(proof2).unwrap()); - assert_eq!(ready.len(), 2); - assert!(!sequencer.has_pending()); - } - - #[test] - fn test_add_proof_with_gaps() { - let mut sequencer = ProofSequencer::default(); - let proof1 = MultiProof::default(); - let proof3 = MultiProof::default(); - sequencer.next_sequence = 3; - - let ready = sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proof1).unwrap()); - assert_eq!(ready.len(), 1); - - let ready = sequencer.add_proof(2, SparseTrieUpdate::from_multiproof(proof3).unwrap()); - assert_eq!(ready.len(), 0); - assert!(sequencer.has_pending()); - } - - #[test] - fn test_add_proof_duplicate_sequence() { - let mut sequencer = ProofSequencer::default(); - let proof1 = MultiProof::default(); - let proof2 = MultiProof::default(); - - let ready = sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proof1).unwrap()); - assert_eq!(ready.len(), 1); - - let ready = sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proof2).unwrap()); - assert_eq!(ready.len(), 0); - assert!(!sequencer.has_pending()); - } - - #[test] - fn test_add_proof_batch_processing() { - let mut sequencer = ProofSequencer::default(); - let proofs: Vec<_> = (0..5).map(|_| MultiProof::default()).collect(); - sequencer.next_sequence = 5; - - sequencer.add_proof(4, SparseTrieUpdate::from_multiproof(proofs[4].clone()).unwrap()); - sequencer.add_proof(2, SparseTrieUpdate::from_multiproof(proofs[2].clone()).unwrap()); - sequencer.add_proof(1, SparseTrieUpdate::from_multiproof(proofs[1].clone()).unwrap()); - sequencer.add_proof(3, SparseTrieUpdate::from_multiproof(proofs[3].clone()).unwrap()); - - let ready = - sequencer.add_proof(0, SparseTrieUpdate::from_multiproof(proofs[0].clone()).unwrap()); - assert_eq!(ready.len(), 5); - assert!(!sequencer.has_pending()); - } - - fn create_get_proof_targets_state() -> HashedPostState { - let mut state = HashedPostState::default(); - - let addr1 = B256::random(); - let addr2 = B256::random(); - state.accounts.insert(addr1, Some(Default::default())); - state.accounts.insert(addr2, Some(Default::default())); - - let mut storage = HashedStorage::default(); - let slot1 = B256::random(); - let slot2 = B256::random(); - storage.storage.insert(slot1, U256::ZERO); - storage.storage.insert(slot2, U256::from(1)); - state.storages.insert(addr1, storage); - - state - } - - fn unwrap_legacy_targets(targets: VersionedMultiProofTargets) -> MultiProofTargets { - match targets { - VersionedMultiProofTargets::Legacy(targets) => targets, - VersionedMultiProofTargets::V2(_) => panic!("Expected Legacy targets"), - } - } - - #[test] - fn test_get_proof_targets_new_account_targets() { - let state = create_get_proof_targets_state(); - let fetched = MultiProofTargets::default(); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - // should return all accounts as targets since nothing was fetched before - assert_eq!(targets.len(), state.accounts.len()); - for addr in state.accounts.keys() { - assert!(targets.contains_key(addr)); - } - } - - #[test] - fn test_get_proof_targets_new_storage_targets() { - let state = create_get_proof_targets_state(); - let fetched = MultiProofTargets::default(); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - // verify storage slots are included for accounts with storage - for (addr, storage) in &state.storages { - assert!(targets.contains_key(addr)); - let target_slots = &targets[addr]; - assert_eq!(target_slots.len(), storage.storage.len()); - for slot in storage.storage.keys() { - assert!(target_slots.contains(slot)); - } - } - } - - #[test] - fn test_get_proof_targets_filter_already_fetched_accounts() { - let state = create_get_proof_targets_state(); - let mut fetched = MultiProofTargets::default(); - - // select an account that has no storage updates - let fetched_addr = state - .accounts - .keys() - .find(|&&addr| !state.storages.contains_key(&addr)) - .expect("Should have an account without storage"); - - // mark the account as already fetched - fetched.insert(*fetched_addr, HashSet::default()); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - // should not include the already fetched account since it has no storage updates - assert!(!targets.contains_key(fetched_addr)); - // other accounts should still be included - assert_eq!(targets.len(), state.accounts.len() - 1); - } - - #[test] - fn test_get_proof_targets_filter_already_fetched_storage() { - let state = create_get_proof_targets_state(); - let mut fetched = MultiProofTargets::default(); - - // mark one storage slot as already fetched - let (addr, storage) = state.storages.iter().next().unwrap(); - let mut fetched_slots = HashSet::default(); - let fetched_slot = *storage.storage.keys().next().unwrap(); - fetched_slots.insert(fetched_slot); - fetched.insert(*addr, fetched_slots); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - // should not include the already fetched storage slot - let target_slots = &targets[addr]; - assert!(!target_slots.contains(&fetched_slot)); - assert_eq!(target_slots.len(), storage.storage.len() - 1); - } - - #[test] - fn test_get_proof_targets_empty_state() { - let state = HashedPostState::default(); - let fetched = MultiProofTargets::default(); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - assert!(targets.is_empty()); - } - - #[test] - fn test_get_proof_targets_mixed_fetched_state() { - let mut state = HashedPostState::default(); - let mut fetched = MultiProofTargets::default(); - - let addr1 = B256::random(); - let addr2 = B256::random(); - let slot1 = B256::random(); - let slot2 = B256::random(); - - state.accounts.insert(addr1, Some(Default::default())); - state.accounts.insert(addr2, Some(Default::default())); - - let mut storage = HashedStorage::default(); - storage.storage.insert(slot1, U256::ZERO); - storage.storage.insert(slot2, U256::from(1)); - state.storages.insert(addr1, storage); - - let mut fetched_slots = HashSet::default(); - fetched_slots.insert(slot1); - fetched.insert(addr1, fetched_slots); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - assert!(targets.contains_key(&addr2)); - assert!(!targets[&addr1].contains(&slot1)); - assert!(targets[&addr1].contains(&slot2)); - } - - #[test] - fn test_get_proof_targets_unmodified_account_with_storage() { - let mut state = HashedPostState::default(); - let fetched = MultiProofTargets::default(); - - let addr = B256::random(); - let slot1 = B256::random(); - let slot2 = B256::random(); - - // don't add the account to state.accounts (simulating unmodified account) - // but add storage updates for this account - let mut storage = HashedStorage::default(); - storage.storage.insert(slot1, U256::from(1)); - storage.storage.insert(slot2, U256::from(2)); - state.storages.insert(addr, storage); - - assert!(!state.accounts.contains_key(&addr)); - assert!(!fetched.contains_key(&addr)); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &MultiAddedRemovedKeys::new(), - false, - )); - - // verify that we still get the storage slots for the unmodified account - assert!(targets.contains_key(&addr)); - - let target_slots = &targets[&addr]; - assert_eq!(target_slots.len(), 2); - assert!(target_slots.contains(&slot1)); - assert!(target_slots.contains(&slot2)); - } - - #[test] - fn test_get_proof_targets_with_removed_storage_keys() { - let mut state = HashedPostState::default(); - let mut fetched = MultiProofTargets::default(); - let mut multi_added_removed_keys = MultiAddedRemovedKeys::new(); - - let addr = B256::random(); - let slot1 = B256::random(); - let slot2 = B256::random(); - - // add account to state - state.accounts.insert(addr, Some(Default::default())); - - // add storage updates - let mut storage = HashedStorage::default(); - storage.storage.insert(slot1, U256::from(100)); - storage.storage.insert(slot2, U256::from(200)); - state.storages.insert(addr, storage); - - // mark slot1 as already fetched - let mut fetched_slots = HashSet::default(); - fetched_slots.insert(slot1); - fetched.insert(addr, fetched_slots); - - // update multi_added_removed_keys to mark slot1 as removed - let mut removed_state = HashedPostState::default(); - let mut removed_storage = HashedStorage::default(); - removed_storage.storage.insert(slot1, U256::ZERO); // U256::ZERO marks as removed - removed_state.storages.insert(addr, removed_storage); - multi_added_removed_keys.update_with_state(&removed_state); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &multi_added_removed_keys, - false, - )); - - // slot1 should be included despite being fetched, because it's marked as removed - assert!(targets.contains_key(&addr)); - let target_slots = &targets[&addr]; - assert_eq!(target_slots.len(), 2); - assert!(target_slots.contains(&slot1)); // included because it's removed - assert!(target_slots.contains(&slot2)); // included because it's not fetched - } - - #[test] - fn test_get_proof_targets_with_wiped_storage() { - let mut state = HashedPostState::default(); - let fetched = MultiProofTargets::default(); - let multi_added_removed_keys = MultiAddedRemovedKeys::new(); - - let addr = B256::random(); - let slot1 = B256::random(); - - // add account to state - state.accounts.insert(addr, Some(Default::default())); - - // add wiped storage - let mut storage = HashedStorage::new(true); - storage.storage.insert(slot1, U256::from(100)); - state.storages.insert(addr, storage); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &multi_added_removed_keys, - false, - )); - - // account should be included because storage is wiped and account wasn't fetched - assert!(targets.contains_key(&addr)); - let target_slots = &targets[&addr]; - assert_eq!(target_slots.len(), 1); - assert!(target_slots.contains(&slot1)); - } - - #[test] - fn test_get_proof_targets_removed_keys_not_in_state_update() { - let mut state = HashedPostState::default(); - let mut fetched = MultiProofTargets::default(); - let mut multi_added_removed_keys = MultiAddedRemovedKeys::new(); - - let addr = B256::random(); - let slot1 = B256::random(); - let slot2 = B256::random(); - let slot3 = B256::random(); - - // add account to state - state.accounts.insert(addr, Some(Default::default())); - - // add storage updates for slot1 and slot2 only - let mut storage = HashedStorage::default(); - storage.storage.insert(slot1, U256::from(100)); - storage.storage.insert(slot2, U256::from(200)); - state.storages.insert(addr, storage); - - // mark all slots as already fetched - let mut fetched_slots = HashSet::default(); - fetched_slots.insert(slot1); - fetched_slots.insert(slot2); - fetched_slots.insert(slot3); // slot3 is fetched but not in state update - fetched.insert(addr, fetched_slots); - - // mark slot3 as removed (even though it's not in the state update) - let mut removed_state = HashedPostState::default(); - let mut removed_storage = HashedStorage::default(); - removed_storage.storage.insert(slot3, U256::ZERO); - removed_state.storages.insert(addr, removed_storage); - multi_added_removed_keys.update_with_state(&removed_state); - - let targets = unwrap_legacy_targets(get_proof_targets( - &state, - &fetched, - &multi_added_removed_keys, - false, - )); - - // only slots in the state update can be included, so slot3 should not appear - assert!(!targets.contains_key(&addr)); - } - - /// Verifies that consecutive prefetch proof messages are batched together. - #[test] - fn test_prefetch_proofs_batching() { - let test_provider_factory = create_test_provider_factory(); - let mut task = create_test_state_root_task(test_provider_factory); - - // send multiple messages - let addr1 = B256::random(); - let addr2 = B256::random(); - let addr3 = B256::random(); - - let mut targets1 = MultiProofTargets::default(); - targets1.insert(addr1, HashSet::default()); - - let mut targets2 = MultiProofTargets::default(); - targets2.insert(addr2, HashSet::default()); - - let mut targets3 = MultiProofTargets::default(); - targets3.insert(addr3, HashSet::default()); - - let tx = task.tx.clone(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(targets1))) - .unwrap(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(targets2))) - .unwrap(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(targets3))) - .unwrap(); - - let proofs_requested = - if let Ok(MultiProofMessage::PrefetchProofs(targets)) = task.rx.recv() { - // simulate the batching logic - let mut merged_targets = targets; - let mut num_batched = 1; - while let Ok(MultiProofMessage::PrefetchProofs(next_targets)) = task.rx.try_recv() { - merged_targets.extend(next_targets); - num_batched += 1; - } - - assert_eq!(num_batched, 3); - assert_eq!(merged_targets.len(), 3); - let legacy_targets = unwrap_legacy_targets(merged_targets); - assert!(legacy_targets.contains_key(&addr1)); - assert!(legacy_targets.contains_key(&addr2)); - assert!(legacy_targets.contains_key(&addr3)); - - task.on_prefetch_proof(VersionedMultiProofTargets::Legacy(legacy_targets)) - } else { - panic!("Expected PrefetchProofs message"); - }; - - assert!(proofs_requested >= 1); - } - - /// Verifies that different message types arriving mid-batch are not lost and preserve order. - #[test] - fn test_batching_preserves_ordering_with_different_message_type() { - use alloy_evm::block::StateChangeSource; - use revm_state::Account; - - let test_provider_factory = create_test_provider_factory(); - let task = create_test_state_root_task(test_provider_factory); - - let addr1 = B256::random(); - let addr2 = B256::random(); - let addr3 = B256::random(); - let state_addr1 = alloy_primitives::Address::random(); - let state_addr2 = alloy_primitives::Address::random(); - - // Create PrefetchProofs targets - let mut targets1 = MultiProofTargets::default(); - targets1.insert(addr1, HashSet::default()); - - let mut targets2 = MultiProofTargets::default(); - targets2.insert(addr2, HashSet::default()); - - let mut targets3 = MultiProofTargets::default(); - targets3.insert(addr3, HashSet::default()); - - // Create StateUpdate 1 - let mut state_update1 = EvmState::default(); - state_update1.insert( - state_addr1, - Account { - info: revm_state::AccountInfo { - balance: U256::from(100), - nonce: 1, - code_hash: Default::default(), - code: Default::default(), - account_id: None, - }, - original_info: Box::new(revm_state::AccountInfo::default()), - transaction_id: Default::default(), - storage: Default::default(), - status: revm_state::AccountStatus::Touched, - }, - ); - - // Create StateUpdate 2 - let mut state_update2 = EvmState::default(); - state_update2.insert( - state_addr2, - Account { - info: revm_state::AccountInfo { - balance: U256::from(200), - nonce: 2, - code_hash: Default::default(), - code: Default::default(), - account_id: None, - }, - original_info: Box::new(revm_state::AccountInfo::default()), - transaction_id: Default::default(), - storage: Default::default(), - status: revm_state::AccountStatus::Touched, - }, - ); - - let source = StateChangeSource::Transaction(42); - - // Queue: [PrefetchProofs1, PrefetchProofs2, StateUpdate1, StateUpdate2, PrefetchProofs3] - let tx = task.tx.clone(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(targets1))) - .unwrap(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(targets2))) - .unwrap(); - tx.send(MultiProofMessage::StateUpdate(source.into(), state_update1)).unwrap(); - tx.send(MultiProofMessage::StateUpdate(source.into(), state_update2)).unwrap(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy( - targets3.clone(), - ))) - .unwrap(); - - // Step 1: Receive and batch PrefetchProofs (should get targets1 + targets2) - let mut pending_msg: Option = None; - if let Ok(MultiProofMessage::PrefetchProofs(targets)) = task.rx.recv() { - let mut merged_targets = targets; - let mut num_batched = 1; - - loop { - match task.rx.try_recv() { - Ok(MultiProofMessage::PrefetchProofs(next_targets)) => { - merged_targets.extend(next_targets); - num_batched += 1; - } - Ok(other_msg) => { - // Store locally to preserve ordering (the fix) - pending_msg = Some(other_msg); - break; - } - Err(_) => break, - } - } - - // Should have batched exactly 2 PrefetchProofs (not 3!) - assert_eq!(num_batched, 2, "Should batch only until different message type"); - assert_eq!(merged_targets.len(), 2); - let legacy_targets = unwrap_legacy_targets(merged_targets); - assert!(legacy_targets.contains_key(&addr1)); - assert!(legacy_targets.contains_key(&addr2)); - assert!(!legacy_targets.contains_key(&addr3), "addr3 should NOT be in first batch"); - } else { - panic!("Expected PrefetchProofs message"); - } - - // Step 2: The pending message should be StateUpdate1 (preserved ordering) - match pending_msg { - Some(MultiProofMessage::StateUpdate(_src, update)) => { - assert!(update.contains_key(&state_addr1), "Should be first StateUpdate"); - } - _ => panic!("StateUpdate1 was lost or reordered! The ordering fix is broken."), - } - - // Step 3: Next in channel should be StateUpdate2 - match task.rx.try_recv() { - Ok(MultiProofMessage::StateUpdate(_src, update)) => { - assert!(update.contains_key(&state_addr2), "Should be second StateUpdate"); - } - _ => panic!("StateUpdate2 was lost!"), - } - - // Step 4: Next in channel should be PrefetchProofs3 - match task.rx.try_recv() { - Ok(MultiProofMessage::PrefetchProofs(targets)) => { - assert_eq!(targets.len(), 1); - let legacy_targets = unwrap_legacy_targets(targets); - assert!(legacy_targets.contains_key(&addr3)); - } - _ => panic!("PrefetchProofs3 was lost!"), - } - } - - /// Verifies that a pending message is processed before the next loop iteration (ordering). - #[test] - fn test_pending_message_processed_before_next_iteration() { - use alloy_evm::block::StateChangeSource; - use revm_state::Account; - - let test_provider_factory = create_test_provider_factory(); - let test_provider = create_cached_provider(test_provider_factory.clone()); - let mut task = create_test_state_root_task(test_provider_factory); - - // Queue: Prefetch1, StateUpdate, Prefetch2 - let prefetch_addr1 = B256::random(); - let prefetch_addr2 = B256::random(); - let mut prefetch1 = MultiProofTargets::default(); - prefetch1.insert(prefetch_addr1, HashSet::default()); - let mut prefetch2 = MultiProofTargets::default(); - prefetch2.insert(prefetch_addr2, HashSet::default()); - - let state_addr = alloy_primitives::Address::random(); - let mut state_update = EvmState::default(); - state_update.insert( - state_addr, - Account { - info: revm_state::AccountInfo { - balance: U256::from(42), - nonce: 1, - code_hash: Default::default(), - code: Default::default(), - account_id: None, - }, - original_info: Box::new(revm_state::AccountInfo::default()), - transaction_id: Default::default(), - storage: Default::default(), - status: revm_state::AccountStatus::Touched, - }, - ); - - let source = StateChangeSource::Transaction(99); - - let tx = task.tx.clone(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy(prefetch1))) - .unwrap(); - tx.send(MultiProofMessage::StateUpdate(source.into(), state_update)).unwrap(); - tx.send(MultiProofMessage::PrefetchProofs(VersionedMultiProofTargets::Legacy( - prefetch2.clone(), - ))) - .unwrap(); - - let mut ctx = MultiproofBatchCtx::new(Instant::now()); - let mut batch_metrics = MultiproofBatchMetrics::default(); - - // First message: Prefetch1 batches; StateUpdate becomes pending. - let first = task.rx.recv().unwrap(); - assert!(matches!(first, MultiProofMessage::PrefetchProofs(_))); - assert!(!task.process_multiproof_message( - first, - &mut ctx, - &mut batch_metrics, - &test_provider - )); - let pending = ctx.pending_msg.take().expect("pending message captured"); - assert!(matches!(pending, MultiProofMessage::StateUpdate(_, _))); - - // Pending message should be handled before the next select loop. - // StateUpdate is processed directly without batching. - assert!(!task.process_multiproof_message( - pending, - &mut ctx, - &mut batch_metrics, - &test_provider - )); - - // Since StateUpdate doesn't batch, Prefetch2 remains in the channel (not in pending_msg). - assert!(ctx.pending_msg.is_none()); - - // Prefetch2 should still be in the channel. - match task.rx.try_recv() { - Ok(MultiProofMessage::PrefetchProofs(targets)) => { - assert_eq!(targets.len(), 1); - let legacy_targets = unwrap_legacy_targets(targets); - assert!(legacy_targets.contains_key(&prefetch_addr2)); - } - other => panic!("Expected PrefetchProofs2 in channel, got {:?}", other), - } - } - - /// Verifies that BAL messages are processed correctly and generate state updates. - #[test] - fn test_bal_message_processing() { - let test_provider_factory = create_test_provider_factory(); - let test_provider = create_cached_provider(test_provider_factory.clone()); - let mut task = create_test_state_root_task(test_provider_factory); - - // Create a simple BAL with one account change - let account_address = Address::random(); - let account_changes = AccountChanges { - address: account_address, - balance_changes: vec![BalanceChange::new(0, U256::from(1000))], - nonce_changes: vec![], - code_changes: vec![], - storage_changes: vec![], - storage_reads: vec![], - }; - - let bal = Arc::new(vec![account_changes]); - - let mut ctx = MultiproofBatchCtx::new(Instant::now()); - let mut batch_metrics = MultiproofBatchMetrics::default(); - - let should_finish = task.process_multiproof_message( - MultiProofMessage::BlockAccessList(bal), - &mut ctx, - &mut batch_metrics, - &test_provider, - ); - - // BAL should mark updates as finished - assert!(ctx.updates_finished_time.is_some()); - - // Should have dispatched state update proofs - assert!(batch_metrics.state_update_proofs_requested > 0); - - // Should need to wait for the results of those proofs to arrive - assert!(!should_finish, "Should continue waiting for proofs"); - } -} diff --git a/crates/engine/tree/src/tree/payload_processor/preserved_sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/preserved_sparse_trie.rs index 6a147691af..03677e7b44 100644 --- a/crates/engine/tree/src/tree/payload_processor/preserved_sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/preserved_sparse_trie.rs @@ -12,7 +12,7 @@ pub(super) type SparseTrie = SparseStateTrie; /// Shared handle to a preserved sparse trie that can be reused across payload validations. /// /// This is stored in [`PayloadProcessor`](super::PayloadProcessor) and cloned to pass to -/// [`SparseTrieTask`](super::sparse_trie::SparseTrieTask) for trie reuse. +/// [`SparseTrieCacheTask`](super::sparse_trie::SparseTrieCacheTask) for trie reuse. #[derive(Debug, Default, Clone)] pub(super) struct SharedPreservedSparseTrie(Arc>>); diff --git a/crates/engine/tree/src/tree/payload_processor/prewarm.rs b/crates/engine/tree/src/tree/payload_processor/prewarm.rs index e89b6e1dcd..e465fcd85f 100644 --- a/crates/engine/tree/src/tree/payload_processor/prewarm.rs +++ b/crates/engine/tree/src/tree/payload_processor/prewarm.rs @@ -13,11 +13,7 @@ use crate::tree::{ cached_state::{CachedStateProvider, SavedCache}, - payload_processor::{ - bal, - multiproof::{MultiProofMessage, VersionedMultiProofTargets}, - PayloadExecutionCache, - }, + payload_processor::{bal, multiproof::MultiProofMessage, PayloadExecutionCache}, precompile_cache::{CachedPrecompile, PrecompileCacheMap}, ExecutionEnv, StateProviderBuilder, }; @@ -25,7 +21,7 @@ use alloy_consensus::transaction::TxHashRef; use alloy_eip7928::BlockAccessList; use alloy_eips::eip4895::Withdrawal; use alloy_evm::Database; -use alloy_primitives::{keccak256, map::B256Set, StorageKey, B256}; +use alloy_primitives::{keccak256, StorageKey, B256}; use crossbeam_channel::{Receiver as CrossbeamReceiver, Sender as CrossbeamSender}; use metrics::{Counter, Gauge, Histogram}; use rayon::prelude::*; @@ -38,7 +34,7 @@ use reth_provider::{ }; use reth_revm::{database::StateProviderDatabase, state::EvmState}; use reth_tasks::Runtime; -use reth_trie::MultiProofTargets; +use reth_trie_parallel::targets_v2::MultiProofTargetsV2; use std::sync::{ atomic::{AtomicBool, Ordering}, mpsc::{self, channel, Receiver, Sender, SyncSender}, @@ -187,9 +183,9 @@ where && let Some(withdrawals) = &ctx.env.withdrawals && !withdrawals.is_empty() { - let targets = - multiproof_targets_from_withdrawals(withdrawals, ctx.v2_proofs_enabled); - let _ = to_multi_proof.send(MultiProofMessage::PrefetchProofs(targets)); + let targets = multiproof_targets_from_withdrawals(withdrawals); + let _ = to_multi_proof + .send(MultiProofMessage::PrefetchProofs(targets)); } // drop sender and wait for all tasks to finish @@ -456,8 +452,6 @@ where pub precompile_cache_disabled: bool, /// The precompile cache map. pub precompile_cache_map: PrecompileCacheMap>, - /// Whether V2 proof calculation is enabled. - pub v2_proofs_enabled: bool, } impl PrewarmContext @@ -466,12 +460,9 @@ where P: BlockReader + StateProviderFactory + StateReader + Clone + 'static, Evm: ConfigureEvm + 'static, { - /// Splits this context into an evm, an evm config, metrics, the atomic bool for terminating - /// execution, and whether V2 proofs are enabled. + /// Splits this context into an evm, metrics, and the atomic bool for terminating execution. #[instrument(level = "debug", target = "engine::tree::payload_processor::prewarm", skip_all)] - fn evm_for_ctx( - self, - ) -> Option<(EvmFor, PrewarmMetrics, Arc, bool)> { + fn evm_for_ctx(self) -> Option<(EvmFor, PrewarmMetrics, Arc)> { let Self { env, evm_config, @@ -481,7 +472,6 @@ where terminate_execution, precompile_cache_disabled, precompile_cache_map, - v2_proofs_enabled, } = self; let mut state_provider = match provider.build() { @@ -532,7 +522,7 @@ where }); } - Some((evm, metrics, terminate_execution, v2_proofs_enabled)) + Some((evm, metrics, terminate_execution)) } /// Accepts a [`CrossbeamReceiver`] of transactions and a handle to prewarm task. Executes @@ -553,10 +543,7 @@ where ) where Tx: ExecutableTxFor, { - let Some((mut evm, metrics, terminate_execution, v2_proofs_enabled)) = self.evm_for_ctx() - else { - return - }; + let Some((mut evm, metrics, terminate_execution)) = self.evm_for_ctx() else { return }; while let Ok(IndexedTransaction { index, tx }) = txs.recv() { let _enter = debug_span!( @@ -601,8 +588,7 @@ where // Only send outcome for transactions after the first txn // as the main execution will be just as fast if index > 0 { - let (targets, storage_targets) = - multiproof_targets_from_state(res.state, v2_proofs_enabled); + let (targets, storage_targets) = multiproof_targets_from_state(res.state); metrics.prefetch_storage_targets.record(storage_targets as f64); if let Some(to_multi_proof) = &to_multi_proof { let _ = to_multi_proof.send(MultiProofMessage::PrefetchProofs(targets)); @@ -694,59 +680,10 @@ where } } -/// Returns a set of [`VersionedMultiProofTargets`] and the total amount of storage targets, based -/// on the given state. -fn multiproof_targets_from_state( - state: EvmState, - v2_enabled: bool, -) -> (VersionedMultiProofTargets, usize) { - if v2_enabled { - multiproof_targets_v2_from_state(state) - } else { - multiproof_targets_legacy_from_state(state) - } -} - -/// Returns legacy [`MultiProofTargets`] and the total amount of storage targets, based on the +/// Returns a set of [`MultiProofTargetsV2`] and the total amount of storage targets, based on the /// given state. -fn multiproof_targets_legacy_from_state(state: EvmState) -> (VersionedMultiProofTargets, usize) { - let mut targets = MultiProofTargets::with_capacity(state.len()); - let mut storage_targets = 0; - for (addr, account) in state { - // if the account was not touched, or if the account was selfdestructed, do not - // fetch proofs for it - // - // Since selfdestruct can only happen in the same transaction, we can skip - // prefetching proofs for selfdestructed accounts - // - // See: https://eips.ethereum.org/EIPS/eip-6780 - if !account.is_touched() || account.is_selfdestructed() { - continue - } - - let mut storage_set = - B256Set::with_capacity_and_hasher(account.storage.len(), Default::default()); - for (key, slot) in account.storage { - // do nothing if unchanged - if !slot.is_changed() { - continue - } - - storage_set.insert(keccak256(B256::new(key.to_be_bytes()))); - } - - storage_targets += storage_set.len(); - targets.insert(keccak256(addr), storage_set); - } - - (VersionedMultiProofTargets::Legacy(targets), storage_targets) -} - -/// Returns V2 [`reth_trie_parallel::targets_v2::MultiProofTargetsV2`] and the total amount of -/// storage targets, based on the given state. -fn multiproof_targets_v2_from_state(state: EvmState) -> (VersionedMultiProofTargets, usize) { +fn multiproof_targets_from_state(state: EvmState) -> (MultiProofTargetsV2, usize) { use reth_trie::proof_v2; - use reth_trie_parallel::targets_v2::MultiProofTargetsV2; let mut targets = MultiProofTargetsV2::default(); let mut storage_target_count = 0; @@ -782,27 +719,17 @@ fn multiproof_targets_v2_from_state(state: EvmState) -> (VersionedMultiProofTarg } } - (VersionedMultiProofTargets::V2(targets), storage_target_count) + (targets, storage_target_count) } -/// Returns [`VersionedMultiProofTargets`] for withdrawal addresses. +/// Returns [`MultiProofTargetsV2`] for withdrawal addresses. /// /// Withdrawals only modify account balances (no storage), so the targets contain /// only account-level entries with empty storage sets. -fn multiproof_targets_from_withdrawals( - withdrawals: &[Withdrawal], - v2_enabled: bool, -) -> VersionedMultiProofTargets { - use reth_trie_parallel::targets_v2::MultiProofTargetsV2; - if v2_enabled { - VersionedMultiProofTargets::V2(MultiProofTargetsV2 { - account_targets: withdrawals.iter().map(|w| keccak256(w.address).into()).collect(), - ..Default::default() - }) - } else { - VersionedMultiProofTargets::Legacy( - withdrawals.iter().map(|w| (keccak256(w.address), Default::default())).collect(), - ) +fn multiproof_targets_from_withdrawals(withdrawals: &[Withdrawal]) -> MultiProofTargetsV2 { + MultiProofTargetsV2 { + account_targets: withdrawals.iter().map(|w| keccak256(w.address).into()).collect(), + ..Default::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 7e4290594b..b9eb6f5d9b 100644 --- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs +++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs @@ -3,9 +3,9 @@ use crate::tree::{ multiproof::{ dispatch_with_chunking, evm_state_to_hashed_post_state, MultiProofMessage, - VersionedMultiProofTargets, DEFAULT_MAX_TARGETS_FOR_CHUNKING, + DEFAULT_MAX_TARGETS_FOR_CHUNKING, }, - payload_processor::multiproof::{MultiProofTaskMetrics, SparseTrieUpdate}, + payload_processor::multiproof::MultiProofTaskMetrics, }; use alloy_primitives::B256; use alloy_rlp::{Decodable, Encodable}; @@ -14,13 +14,12 @@ use rayon::iter::ParallelIterator; use reth_primitives_traits::{Account, FastInstant as Instant, ParallelBridgeBuffered}; use reth_tasks::Runtime; use reth_trie::{ - proof_v2::Target, updates::TrieUpdates, DecodedMultiProofV2, HashedPostState, Nibbles, - TrieAccount, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, + proof_v2::Target, updates::TrieUpdates, DecodedMultiProofV2, HashedPostState, TrieAccount, + EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE, }; use reth_trie_parallel::{ proof_task::{ - AccountMultiproofInput, ProofResult, ProofResultContext, ProofResultMessage, - ProofWorkerHandle, + AccountMultiproofInput, ProofResultContext, ProofResultMessage, ProofWorkerHandle, }, root::ParallelStateRootError, targets_v2::MultiProofTargetsV2, @@ -28,194 +27,11 @@ use reth_trie_parallel::{ #[cfg(feature = "trie-debug")] use reth_trie_sparse::debug_recorder::TrieDebugRecorder; use reth_trie_sparse::{ - errors::{SparseStateTrieResult, SparseTrieErrorKind, SparseTrieResult}, - provider::{TrieNodeProvider, TrieNodeProviderFactory}, - DeferredDrops, LeafUpdate, ParallelSparseTrie, SparseStateTrie, SparseTrie, + errors::SparseTrieResult, DeferredDrops, LeafUpdate, ParallelSparseTrie, SparseStateTrie, + SparseTrie, }; use revm_primitives::{hash_map::Entry, B256Map}; -use smallvec::SmallVec; -use std::{sync::mpsc, time::Duration}; -use tracing::{debug, debug_span, error, instrument, trace}; - -#[expect(clippy::large_enum_variant)] -pub(super) enum SpawnedSparseTrieTask -where - BPF: TrieNodeProviderFactory + Send + Sync, - BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync, - BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync, - A: SparseTrie + Send + Sync + Default, - S: SparseTrie + Send + Sync + Default + Clone, -{ - Cleared(SparseTrieTask), - Cached(SparseTrieCacheTask), -} - -impl SpawnedSparseTrieTask -where - BPF: TrieNodeProviderFactory + Send + Sync + Clone, - BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync, - BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync, - A: SparseTrie + Send + Sync + Default, - S: SparseTrie + Send + Sync + Default + Clone, -{ - pub(super) fn run(&mut self) -> Result { - match self { - Self::Cleared(task) => task.run(), - Self::Cached(task) => task.run(), - } - } - - pub(super) fn into_trie_for_reuse( - self, - prune_depth: usize, - max_storage_tries: usize, - max_nodes_capacity: usize, - max_values_capacity: usize, - disable_pruning: bool, - ) -> (SparseStateTrie, DeferredDrops) { - match self { - Self::Cleared(task) => task.into_cleared_trie(max_nodes_capacity, max_values_capacity), - Self::Cached(task) => task.into_trie_for_reuse( - prune_depth, - max_storage_tries, - max_nodes_capacity, - max_values_capacity, - disable_pruning, - ), - } - } - - pub(super) fn into_cleared_trie( - self, - max_nodes_capacity: usize, - max_values_capacity: usize, - ) -> (SparseStateTrie, DeferredDrops) { - match self { - Self::Cleared(task) => task.into_cleared_trie(max_nodes_capacity, max_values_capacity), - Self::Cached(task) => task.into_cleared_trie(max_nodes_capacity, max_values_capacity), - } - } -} - -/// A task responsible for populating the sparse trie. -pub(super) struct SparseTrieTask -where - BPF: TrieNodeProviderFactory + Send + Sync, - BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync, - BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync, -{ - /// Receives updates from the state root task. - pub(super) updates: mpsc::Receiver, - /// `SparseStateTrie` used for computing the state root. - pub(super) trie: SparseStateTrie, - pub(super) metrics: MultiProofTaskMetrics, - /// Trie node provider factory. - blinded_provider_factory: BPF, -} - -impl SparseTrieTask -where - BPF: TrieNodeProviderFactory + Send + Sync + Clone, - BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync, - BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync, - A: SparseTrie + Send + Sync + Default, - S: SparseTrie + Send + Sync + Default + Clone, -{ - /// Creates a new sparse trie task with the given trie. - pub(super) const fn new( - updates: mpsc::Receiver, - blinded_provider_factory: BPF, - metrics: MultiProofTaskMetrics, - trie: SparseStateTrie, - ) -> Self { - Self { updates, metrics, trie, blinded_provider_factory } - } - - /// Runs the sparse trie task to completion, computing the state root. - /// - /// Receives [`SparseTrieUpdate`]s until the channel is closed, applying each update - /// to the trie. Once all updates are processed, computes and returns the final state root. - #[instrument( - name = "SparseTrieTask::run", - level = "debug", - target = "engine::tree::payload_processor::sparse_trie", - skip_all - )] - pub(super) fn run(&mut self) -> Result { - let now = Instant::now(); - - let mut num_iterations = 0; - - while let Ok(mut update) = self.updates.recv() { - num_iterations += 1; - let mut num_updates = 1; - let _enter = - debug_span!(target: "engine::tree::payload_processor::sparse_trie", "drain updates") - .entered(); - while let Ok(next) = self.updates.try_recv() { - update.extend(next); - num_updates += 1; - } - drop(_enter); - - debug!( - target: "engine::root", - num_updates, - account_proofs = update.multiproof.account_proofs_len(), - storage_proofs = update.multiproof.storage_proofs_len(), - "Updating sparse trie" - ); - - let elapsed = - update_sparse_trie(&mut self.trie, update, &self.blinded_provider_factory) - .map_err(|e| { - ParallelStateRootError::Other(format!( - "could not calculate state root: {e:?}" - )) - })?; - self.metrics.sparse_trie_update_duration_histogram.record(elapsed); - trace!(target: "engine::root", ?elapsed, num_iterations, "Root calculation completed"); - } - - debug!(target: "engine::root", num_iterations, "All proofs processed, ending calculation"); - - let start = Instant::now(); - let (state_root, trie_updates) = - self.trie.root_with_updates(&self.blinded_provider_factory).map_err(|e| { - ParallelStateRootError::Other(format!("could not calculate state root: {e:?}")) - })?; - - #[cfg(feature = "trie-debug")] - let debug_recorders = self.trie.take_debug_recorders(); - - let end = Instant::now(); - self.metrics.sparse_trie_final_update_duration_histogram.record(end.duration_since(start)); - self.metrics.sparse_trie_total_duration_histogram.record(end.duration_since(now)); - - Ok(StateRootComputeOutcome { - state_root, - trie_updates, - #[cfg(feature = "trie-debug")] - debug_recorders, - }) - } - - /// Clears and shrinks the trie, discarding all state. - /// - /// Use this when the payload was invalid or cancelled - we don't want to preserve - /// potentially invalid trie state, but we keep the allocations for reuse. - pub(super) fn into_cleared_trie( - self, - max_nodes_capacity: usize, - max_values_capacity: usize, - ) -> (SparseStateTrie, DeferredDrops) { - let Self { mut trie, .. } = self; - trie.clear(); - trie.shrink_to(max_nodes_capacity, max_values_capacity); - let deferred = trie.take_deferred_drops(); - (trie, deferred) - } -} +use tracing::{debug, debug_span, error, instrument}; /// Maximum number of pending/prewarm updates that we accumulate in memory before actually applying. const MAX_PENDING_UPDATES: usize = 100; @@ -436,14 +252,10 @@ where let Ok(result) = message else { unreachable!("we own the sender half") }; - let ProofResult::V2(mut result) = result.result? else { - unreachable!("sparse trie as cache must only be used with multiproof v2"); - }; + let mut result = result.result?; while let Ok(next) = self.proof_result_rx.try_recv() { - let ProofResult::V2(res) = next.result? else { - unreachable!("sparse trie as cache must only be used with multiproof v2"); - }; + let res = next.result?; result.extend(res); } @@ -516,11 +328,7 @@ where target = "engine::tree::payload_processor::sparse_trie", skip_all )] - fn on_prewarm_targets(&mut self, targets: VersionedMultiProofTargets) { - let VersionedMultiProofTargets::V2(targets) = targets else { - unreachable!("sparse trie as cache must only be used with V2 multiproof targets"); - }; - + fn on_prewarm_targets(&mut self, targets: MultiProofTargetsV2) { for target in targets.account_targets { // Only touch accounts that are not yet present in the updates set. self.new_account_updates.entry(target.key()).or_insert(LeafUpdate::Touched); @@ -873,7 +681,7 @@ where MultiProofTargetsV2::chunks, |proof_targets| { if let Err(e) = self.proof_worker_handle.dispatch_account_multiproof( - AccountMultiproofInput::V2 { + AccountMultiproofInput { targets: proof_targets, proof_result_sender: ProofResultContext::new( self.proof_result_tx.clone(), @@ -896,7 +704,7 @@ enum SparseTrieTaskMessage { /// A hashed state update ready to be processed. HashedState(HashedPostState), /// Prefetch proof targets (passed through directly). - PrefetchProofs(VersionedMultiProofTargets), + PrefetchProofs(MultiProofTargetsV2), /// Signals that all state updates have been received. FinishedStateUpdates, } @@ -915,160 +723,6 @@ pub struct StateRootComputeOutcome { pub debug_recorders: Vec<(Option, TrieDebugRecorder)>, } -/// Updates the sparse trie with the given proofs and state, and returns the elapsed time. -#[instrument(level = "debug", target = "engine::tree::payload_processor::sparse_trie", skip_all)] -pub(crate) fn update_sparse_trie( - trie: &mut SparseStateTrie, - SparseTrieUpdate { mut state, multiproof }: SparseTrieUpdate, - blinded_provider_factory: &BPF, -) -> SparseStateTrieResult -where - BPF: TrieNodeProviderFactory + Send + Sync, - BPF::AccountNodeProvider: TrieNodeProvider + Send + Sync, - BPF::StorageNodeProvider: TrieNodeProvider + Send + Sync, - A: SparseTrie + Send + Sync + Default, - S: SparseTrie + Send + Sync + Default + Clone, -{ - trace!(target: "engine::root::sparse", "Updating sparse trie"); - let started_at = Instant::now(); - - // Reveal new accounts and storage slots. - match multiproof { - ProofResult::Legacy(decoded, _) => { - trie.reveal_decoded_multiproof(decoded)?; - } - ProofResult::V2(decoded_v2) => { - trie.reveal_decoded_multiproof_v2(decoded_v2)?; - } - } - let reveal_multiproof_elapsed = started_at.elapsed(); - trace!( - target: "engine::root::sparse", - ?reveal_multiproof_elapsed, - "Done revealing multiproof" - ); - - // Update storage slots with new values and calculate storage roots. - let span = tracing::Span::current(); - let results: Vec<_> = state - .storages - .into_iter() - .map(|(address, storage)| (address, storage, trie.take_storage_trie(&address))) - .par_bridge_buffered() - .map(|(address, storage, storage_trie)| { - let _enter = - debug_span!(target: "engine::tree::payload_processor::sparse_trie", parent: &span, "storage trie", ?address) - .entered(); - - trace!(target: "engine::tree::payload_processor::sparse_trie", "Updating storage"); - let storage_provider = blinded_provider_factory.storage_node_provider(address); - let mut storage_trie = storage_trie.ok_or(SparseTrieErrorKind::Blind)?; - - if storage.wiped { - trace!(target: "engine::tree::payload_processor::sparse_trie", "Wiping storage"); - storage_trie.wipe()?; - } - - // Defer leaf removals until after updates/additions, so that we don't delete an - // intermediate branch node during a removal and then re-add that branch back during a - // later leaf addition. This is an optimization, but also a requirement inherited from - // multiproof generating, which can't know the order that leaf operations happen in. - let mut removed_slots = SmallVec::<[Nibbles; 8]>::new(); - - for (slot, value) in storage.storage { - let slot_nibbles = Nibbles::unpack(slot); - - if value.is_zero() { - removed_slots.push(slot_nibbles); - continue; - } - - trace!(target: "engine::tree::payload_processor::sparse_trie", ?slot_nibbles, "Updating storage slot"); - storage_trie.update_leaf( - slot_nibbles, - alloy_rlp::encode_fixed_size(&value).to_vec(), - &storage_provider, - )?; - } - - for slot_nibbles in removed_slots { - trace!(target: "engine::root::sparse", ?slot_nibbles, "Removing storage slot"); - storage_trie.remove_leaf(&slot_nibbles, &storage_provider)?; - } - - storage_trie.root(); - - SparseStateTrieResult::Ok((address, storage_trie)) - }) - .collect(); - - // Defer leaf removals until after updates/additions, so that we don't delete an intermediate - // branch node during a removal and then re-add that branch back during a later leaf addition. - // This is an optimization, but also a requirement inherited from multiproof generating, which - // can't know the order that leaf operations happen in. - let mut removed_accounts = Vec::new(); - - // Update account storage roots - let _enter = - tracing::debug_span!(target: "engine::tree::payload_processor::sparse_trie", "account trie") - .entered(); - for result in results { - let (address, storage_trie) = result?; - trie.insert_storage_trie(address, storage_trie); - - if let Some(account) = state.accounts.remove(&address) { - // If the account itself has an update, remove it from the state update and update in - // one go instead of doing it down below. - trace!(target: "engine::root::sparse", ?address, "Updating account and its storage root"); - if !trie.update_account( - address, - account.unwrap_or_default(), - blinded_provider_factory, - )? { - removed_accounts.push(address); - } - } else if trie.is_account_revealed(address) { - // Otherwise, if the account is revealed, only update its storage root. - trace!(target: "engine::root::sparse", ?address, "Updating account storage root"); - if !trie.update_account_storage_root(address, blinded_provider_factory)? { - removed_accounts.push(address); - } - } - } - - // Update accounts - for (address, account) in state.accounts { - trace!(target: "engine::root::sparse", ?address, "Updating account"); - if !trie.update_account(address, account.unwrap_or_default(), blinded_provider_factory)? { - removed_accounts.push(address); - } - } - - // Remove accounts - for address in removed_accounts { - trace!(target: "engine::root::sparse", ?address, "Removing account"); - let nibbles = Nibbles::unpack(address); - trie.remove_account_leaf(&nibbles, blinded_provider_factory)?; - } - - let elapsed_before = started_at.elapsed(); - trace!( - target: "engine::root::sparse", - "Calculating subtries" - ); - trie.calculate_subtries(); - - let elapsed = started_at.elapsed(); - let below_level_elapsed = elapsed - elapsed_before; - trace!( - target: "engine::root::sparse", - ?below_level_elapsed, - "Intermediate nodes calculated" - ); - - Ok(elapsed) -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/node/core/src/args/engine.rs b/crates/node/core/src/args/engine.rs index 7964b0ff20..5dede46a4b 100644 --- a/crates/node/core/src/args/engine.rs +++ b/crates/node/core/src/args/engine.rs @@ -39,9 +39,7 @@ pub struct DefaultEngineValues { storage_worker_count: Option, account_worker_count: Option, prewarming_threads: Option, - disable_proof_v2: bool, cache_metrics_disabled: bool, - disable_trie_cache: bool, sparse_trie_prune_depth: usize, sparse_trie_max_storage_tries: usize, disable_sparse_trie_cache_pruning: bool, @@ -176,24 +174,12 @@ impl DefaultEngineValues { self } - /// Set whether to disable proof V2 by default - pub const fn with_disable_proof_v2(mut self, v: bool) -> Self { - self.disable_proof_v2 = v; - self - } - /// Set whether to disable cache metrics by default pub const fn with_cache_metrics_disabled(mut self, v: bool) -> Self { self.cache_metrics_disabled = v; self } - /// Set whether to disable sparse trie cache by default - pub const fn with_disable_trie_cache(mut self, v: bool) -> Self { - self.disable_trie_cache = v; - self - } - /// Set the sparse trie prune depth by default pub const fn with_sparse_trie_prune_depth(mut self, v: usize) -> Self { self.sparse_trie_prune_depth = v; @@ -241,9 +227,7 @@ impl Default for DefaultEngineValues { storage_worker_count: None, account_worker_count: None, prewarming_threads: None, - disable_proof_v2: false, cache_metrics_disabled: false, - disable_trie_cache: false, sparse_trie_prune_depth: DEFAULT_SPARSE_TRIE_PRUNE_DEPTH, sparse_trie_max_storage_tries: DEFAULT_SPARSE_TRIE_MAX_STORAGE_TRIES, disable_sparse_trie_cache_pruning: false, @@ -373,18 +357,10 @@ pub struct EngineArgs { #[arg(long = "engine.prewarming-threads", default_value = Resettable::from(DefaultEngineValues::get_global().prewarming_threads.map(|v| v.to_string().into())))] pub prewarming_threads: Option, - /// Disable V2 storage proofs for state root calculations - #[arg(long = "engine.disable-proof-v2", default_value_t = DefaultEngineValues::get_global().disable_proof_v2)] - pub disable_proof_v2: bool, - /// Disable cache metrics recording, which can take up to 50ms with large cached state. #[arg(long = "engine.disable-cache-metrics", default_value_t = DefaultEngineValues::get_global().cache_metrics_disabled)] pub cache_metrics_disabled: bool, - /// Disable sparse trie cache. - #[arg(long = "engine.disable-trie-cache", default_value_t = DefaultEngineValues::get_global().disable_trie_cache, conflicts_with = "disable_proof_v2")] - pub disable_trie_cache: bool, - /// Sparse trie prune depth. #[arg(long = "engine.sparse-trie-prune-depth", default_value_t = DefaultEngineValues::get_global().sparse_trie_prune_depth)] pub sparse_trie_prune_depth: usize, @@ -438,9 +414,7 @@ impl Default for EngineArgs { storage_worker_count, account_worker_count, prewarming_threads, - disable_proof_v2, cache_metrics_disabled, - disable_trie_cache, sparse_trie_prune_depth, sparse_trie_max_storage_tries, disable_sparse_trie_cache_pruning, @@ -470,9 +444,7 @@ impl Default for EngineArgs { storage_worker_count, account_worker_count, prewarming_threads, - disable_proof_v2, cache_metrics_disabled, - disable_trie_cache, sparse_trie_prune_depth, sparse_trie_max_storage_tries, disable_sparse_trie_cache_pruning, @@ -506,9 +478,7 @@ impl EngineArgs { .with_unwind_canonical_header(self.allow_unwind_canonical_header) .with_storage_worker_count_opt(self.storage_worker_count) .with_account_worker_count_opt(self.account_worker_count) - .with_disable_proof_v2(self.disable_proof_v2) .without_cache_metrics(self.cache_metrics_disabled) - .with_disable_trie_cache(self.disable_trie_cache) .with_sparse_trie_prune_depth(self.sparse_trie_prune_depth) .with_sparse_trie_max_storage_tries(self.sparse_trie_max_storage_tries) .with_disable_sparse_trie_cache_pruning(self.disable_sparse_trie_cache_pruning) @@ -562,9 +532,7 @@ mod tests { storage_worker_count: Some(16), account_worker_count: Some(8), prewarming_threads: Some(4), - disable_proof_v2: false, cache_metrics_disabled: true, - disable_trie_cache: true, sparse_trie_prune_depth: 10, sparse_trie_max_storage_tries: 100, disable_sparse_trie_cache_pruning: true, @@ -601,7 +569,6 @@ mod tests { "--engine.prewarming-threads", "4", "--engine.disable-cache-metrics", - "--engine.disable-trie-cache", "--engine.sparse-trie-prune-depth", "10", "--engine.sparse-trie-max-storage-tries", diff --git a/crates/trie/parallel/Cargo.toml b/crates/trie/parallel/Cargo.toml index 678f29eb3c..87f4de92ab 100644 --- a/crates/trie/parallel/Cargo.toml +++ b/crates/trie/parallel/Cargo.toml @@ -17,7 +17,6 @@ reth-primitives-traits = { workspace = true, features = ["dashmap", "std"] } reth-execution-errors.workspace = true reth-provider.workspace = true reth-storage-errors.workspace = true -reth-trie-common.workspace = true reth-trie-sparse = { workspace = true, features = ["std"] } reth-tasks = { workspace = true, features = ["rayon"] } reth-trie.workspace = true @@ -60,7 +59,6 @@ metrics = ["reth-metrics", "dep:metrics", "reth-trie/metrics", "reth-trie-sparse test-utils = [ "reth-primitives-traits/test-utils", "reth-provider/test-utils", - "reth-trie-common/test-utils", "reth-trie-db/test-utils", "reth-trie-sparse/test-utils", "reth-trie/test-utils", diff --git a/crates/trie/parallel/src/lib.rs b/crates/trie/parallel/src/lib.rs index cba9d9440e..8b6a53f142 100644 --- a/crates/trie/parallel/src/lib.rs +++ b/crates/trie/parallel/src/lib.rs @@ -18,8 +18,6 @@ pub mod stats; pub mod root; /// Implementation of parallel proof computation. -pub mod proof; - pub mod proof_task; /// Async value encoder for V2 proofs. diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs deleted file mode 100644 index dbbf4f7356..0000000000 --- a/crates/trie/parallel/src/proof.rs +++ /dev/null @@ -1,283 +0,0 @@ -use crate::{ - metrics::ParallelTrieMetrics, - proof_task::{AccountMultiproofInput, ProofResult, ProofResultContext, ProofWorkerHandle}, - root::ParallelStateRootError, - StorageRootTargets, -}; -use crossbeam_channel::unbounded as crossbeam_unbounded; -use reth_primitives_traits::FastInstant as Instant; -use reth_trie::{ - prefix_set::{PrefixSetMut, TriePrefixSets, TriePrefixSetsMut}, - DecodedMultiProof, HashedPostState, MultiProofTargets, Nibbles, -}; -use reth_trie_common::added_removed_keys::MultiAddedRemovedKeys; -use std::sync::Arc; -use tracing::trace; - -/// Parallel proof calculator. -/// -/// This can collect proof for many targets in parallel, spawning a task for each hashed address -/// that has proof targets. -#[derive(Debug)] -pub struct ParallelProof { - /// The collection of prefix sets for the computation. - pub prefix_sets: Arc, - /// Flag indicating whether to include branch node masks in the proof. - collect_branch_node_masks: bool, - /// Provided by the user to give the necessary context to retain extra proofs. - multi_added_removed_keys: Option>, - /// Handle to the proof worker pools. - proof_worker_handle: ProofWorkerHandle, - /// Whether to use V2 storage proofs. - v2_proofs_enabled: bool, - #[cfg(feature = "metrics")] - metrics: ParallelTrieMetrics, -} - -impl ParallelProof { - /// Create new state proof generator. - pub fn new( - prefix_sets: Arc, - proof_worker_handle: ProofWorkerHandle, - ) -> Self { - Self { - prefix_sets, - collect_branch_node_masks: false, - multi_added_removed_keys: None, - proof_worker_handle, - v2_proofs_enabled: false, - #[cfg(feature = "metrics")] - metrics: ParallelTrieMetrics::new_with_labels(&[("type", "proof")]), - } - } - - /// Set whether to use V2 storage proofs. - pub const fn with_v2_proofs_enabled(mut self, v2_proofs_enabled: bool) -> Self { - self.v2_proofs_enabled = v2_proofs_enabled; - self - } - - /// Set the flag indicating whether to include branch node masks in the proof. - pub const fn with_branch_node_masks(mut self, branch_node_masks: bool) -> Self { - self.collect_branch_node_masks = branch_node_masks; - self - } - - /// Configure the `ParallelProof` with a [`MultiAddedRemovedKeys`], allowing for retaining - /// extra proofs needed to add and remove leaf nodes from the tries. - pub fn with_multi_added_removed_keys( - mut self, - multi_added_removed_keys: Option>, - ) -> Self { - self.multi_added_removed_keys = multi_added_removed_keys; - self - } - - /// Extends prefix sets with the given multiproof targets and returns the frozen result. - /// - /// This is a helper function used to prepare prefix sets before computing multiproofs. - /// Returns frozen (immutable) prefix sets ready for use in proof computation. - pub fn extend_prefix_sets_with_targets( - base_prefix_sets: &TriePrefixSetsMut, - targets: &MultiProofTargets, - ) -> TriePrefixSets { - let mut extended = base_prefix_sets.clone(); - extended.extend(TriePrefixSetsMut { - account_prefix_set: PrefixSetMut::from(targets.keys().copied().map(Nibbles::unpack)), - storage_prefix_sets: targets - .iter() - .filter(|&(_hashed_address, slots)| !slots.is_empty()) - .map(|(hashed_address, slots)| { - (*hashed_address, PrefixSetMut::from(slots.iter().map(Nibbles::unpack))) - }) - .collect(), - destroyed_accounts: Default::default(), - }); - extended.freeze() - } - - /// Generate a state multiproof according to specified targets. - pub fn decoded_multiproof( - self, - targets: MultiProofTargets, - ) -> Result { - // Extend prefix sets with targets - let prefix_sets = Self::extend_prefix_sets_with_targets(&self.prefix_sets, &targets); - - let storage_root_targets_len = StorageRootTargets::count( - &prefix_sets.account_prefix_set, - &prefix_sets.storage_prefix_sets, - ); - - trace!( - target: "trie::parallel_proof", - total_targets = storage_root_targets_len, - "Starting parallel proof generation" - ); - - // Queue account multiproof request to account worker pool - // Create channel for receiving ProofResultMessage - let (result_tx, result_rx) = crossbeam_unbounded(); - let account_multiproof_start_time = Instant::now(); - - let input = AccountMultiproofInput::Legacy { - targets, - prefix_sets, - collect_branch_node_masks: self.collect_branch_node_masks, - multi_added_removed_keys: self.multi_added_removed_keys.clone(), - proof_result_sender: ProofResultContext::new( - result_tx, - 0, - HashedPostState::default(), - account_multiproof_start_time, - ), - }; - - self.proof_worker_handle - .dispatch_account_multiproof(input) - .map_err(|e| ParallelStateRootError::Other(e.to_string()))?; - - // Wait for account multiproof result from worker - let proof_result_msg = result_rx.recv().map_err(|_| { - ParallelStateRootError::Other( - "Account multiproof channel dropped: worker died or pool shutdown".to_string(), - ) - })?; - - let ProofResult::Legacy(multiproof, stats) = proof_result_msg.result? else { - panic!("AccountMultiproofInput::Legacy was submitted, expected legacy result") - }; - - #[cfg(feature = "metrics")] - self.metrics.record(stats); - - trace!( - target: "trie::parallel_proof", - total_targets = storage_root_targets_len, - duration = ?stats.duration(), - branches_added = stats.branches_added(), - leaves_added = stats.leaves_added(), - missed_leaves = stats.missed_leaves(), - precomputed_storage_roots = stats.precomputed_storage_roots(), - "Calculated decoded proof", - ); - - Ok(multiproof) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::proof_task::{ProofTaskCtx, ProofWorkerHandle}; - use alloy_primitives::{ - keccak256, - map::{B256Set, DefaultHashBuilder, HashMap}, - Address, B256, U256, - }; - use rand::Rng; - use reth_primitives_traits::{Account, StorageEntry}; - use reth_provider::{test_utils::create_test_provider_factory, HashingWriter}; - use reth_trie::proof::Proof; - use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory}; - - #[test] - fn random_parallel_proof() { - let factory = create_test_provider_factory(); - - let mut rng = rand::rng(); - let state = (0..100) - .map(|_| { - let address = Address::random(); - let account = - Account { balance: U256::from(rng.random::()), ..Default::default() }; - let mut storage = HashMap::::default(); - let has_storage = rng.random_bool(0.7); - if has_storage { - for _ in 0..100 { - storage.insert( - B256::from(U256::from(rng.random::())), - U256::from(rng.random::()), - ); - } - } - (address, (account, storage)) - }) - .collect::>(); - - { - let provider_rw = factory.provider_rw().unwrap(); - provider_rw - .insert_account_for_hashing( - state.iter().map(|(address, (account, _))| (*address, Some(*account))), - ) - .unwrap(); - provider_rw - .insert_storage_for_hashing(state.iter().map(|(address, (_, storage))| { - ( - *address, - storage - .iter() - .map(|(slot, value)| StorageEntry { key: *slot, value: *value }), - ) - })) - .unwrap(); - provider_rw.commit().unwrap(); - } - - let mut targets = MultiProofTargets::default(); - for (address, (_, storage)) in state.iter().take(10) { - let hashed_address = keccak256(*address); - let mut target_slots = B256Set::default(); - - for (slot, _) in storage.iter().take(5) { - target_slots.insert(*slot); - } - - if !target_slots.is_empty() { - targets.insert(hashed_address, target_slots); - } - } - - let provider_rw = factory.provider_rw().unwrap(); - let trie_cursor_factory = DatabaseTrieCursorFactory::new(provider_rw.tx_ref()); - let hashed_cursor_factory = DatabaseHashedCursorFactory::new(provider_rw.tx_ref()); - - let changeset_cache = reth_trie_db::ChangesetCache::new(); - let factory = - reth_provider::providers::OverlayStateProviderFactory::new(factory, changeset_cache); - let task_ctx = ProofTaskCtx::new(factory); - let runtime = reth_tasks::Runtime::test(); - let proof_worker_handle = ProofWorkerHandle::new(&runtime, task_ctx, false, false); - - let parallel_result = ParallelProof::new(Default::default(), proof_worker_handle.clone()) - .decoded_multiproof(targets.clone()) - .unwrap(); - - let sequential_result_raw = Proof::new(trie_cursor_factory, hashed_cursor_factory) - .multiproof(targets.clone()) - .unwrap(); // targets might be consumed by parallel_result - let sequential_result_decoded: DecodedMultiProof = sequential_result_raw - .try_into() - .expect("Failed to decode sequential_result for test comparison"); - - // to help narrow down what is wrong - first compare account subtries - assert_eq!(parallel_result.account_subtree, sequential_result_decoded.account_subtree); - - // then compare length of all storage subtries - assert_eq!(parallel_result.storages.len(), sequential_result_decoded.storages.len()); - - // then compare each storage subtrie - for (hashed_address, storage_proof) in ¶llel_result.storages { - let sequential_storage_proof = - sequential_result_decoded.storages.get(hashed_address).unwrap(); - assert_eq!(storage_proof, sequential_storage_proof); - } - - // then compare the entire thing for any mask differences - assert_eq!(parallel_result, sequential_result_decoded); - - // Workers shut down automatically when handle is dropped - drop(proof_worker_handle); - } -} diff --git a/crates/trie/parallel/src/proof_task.rs b/crates/trie/parallel/src/proof_task.rs index 3143c6b445..34eb32699a 100644 --- a/crates/trie/parallel/src/proof_task.rs +++ b/crates/trie/parallel/src/proof_task.rs @@ -6,66 +6,50 @@ //! - **Worker Pools**: Pre-spawned workers with dedicated database transactions //! - Storage pool: Handles storage proofs and blinded storage node requests //! - Account pool: Handles account multiproofs and blinded account node requests -//! - **Direct Channel Access**: [`ProofWorkerHandle`] provides type-safe queue methods with direct +//! - **Direct Channel Access**: `ProofWorkerHandle` provides type-safe queue methods with direct //! access to worker channels, eliminating routing overhead //! - **Automatic Shutdown**: Workers terminate gracefully when all handles are dropped //! //! # Message Flow //! -//! 1. `MultiProofTask` prepares a storage or account job and hands it to [`ProofWorkerHandle`]. The -//! job carries a [`ProofResultContext`] so the worker knows how to send the result back. -//! 2. A worker receives the job, runs the proof, and sends a [`ProofResultMessage`] through the -//! provided [`ProofResultSender`]. -//! 3. `MultiProofTask` receives the message, uses `sequence_number` to keep proofs in order, and +//! 1. The multiproof task prepares a storage or account job and hands it to `ProofWorkerHandle`. +//! The job carries a `ProofResultContext` so the worker knows how to send the result back. +//! 2. A worker receives the job, runs the proof, and sends a `ProofResultMessage` through the +//! provided `ProofResultSender`. +//! 3. The multiproof task receives the message, uses `sequence_number` to keep proofs in order, and //! proceeds with its state-root logic. //! -//! Each job gets its own direct channel so results go straight back to `MultiProofTask`. That keeps -//! ordering decisions in one place and lets workers run independently. +//! Each job gets its own direct channel so results go straight back to the multiproof task. That +//! keeps ordering decisions in one place and lets workers run independently. //! //! ```text -//! MultiProofTask -> MultiproofManager -> ProofWorkerHandle -> Storage/Account Worker -//! ^ | -//! | v -//! ProofResultMessage <-------- ProofResultSender --- +//! SparseTrieCacheTask -> ProofWorkerHandle -> Storage/Account Worker +//! ^ | +//! | v +//! ProofResultMessage <-- ProofResultSender //! ``` use crate::{ root::ParallelStateRootError, - stats::{ParallelTrieStats, ParallelTrieTracker}, targets_v2::MultiProofTargetsV2, value_encoder::{AsyncAccountValueEncoder, ValueEncoderStats}, - StorageRootTargets, }; use alloy_primitives::{ map::{B256Map, B256Set}, B256, }; -use alloy_rlp::{BufMut, Encodable}; use crossbeam_channel::{unbounded, Receiver as CrossbeamReceiver, Sender as CrossbeamSender}; use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind, StateProofError}; -use reth_primitives_traits::{ - dashmap::{self, DashMap}, - FastInstant as Instant, -}; +use reth_primitives_traits::{dashmap::DashMap, FastInstant as Instant}; use reth_provider::{DatabaseProviderROFactory, ProviderError, ProviderResult}; use reth_storage_errors::db::DatabaseError; use reth_tasks::Runtime; use reth_trie::{ - hashed_cursor::{HashedCursorFactory, HashedCursorMetricsCache, InstrumentedHashedCursor}, - node_iter::{TrieElement, TrieNodeIter}, - prefix_set::TriePrefixSets, - proof::{ProofBlindedAccountProvider, ProofBlindedStorageProvider, StorageProof}, + hashed_cursor::HashedCursorFactory, + proof::{ProofBlindedAccountProvider, ProofBlindedStorageProvider}, proof_v2, - trie_cursor::{InstrumentedTrieCursor, TrieCursorFactory, TrieCursorMetricsCache}, - walker::TrieWalker, - DecodedMultiProof, DecodedMultiProofV2, DecodedStorageMultiProof, HashBuilder, HashedPostState, - MultiProofTargets, Nibbles, ProofTrieNodeV2, TRIE_ACCOUNT_RLP_MAX_SIZE, -}; -use reth_trie_common::{ - added_removed_keys::MultiAddedRemovedKeys, - prefix_set::{PrefixSet, PrefixSetMut}, - proof::{DecodedProofNodes, ProofRetainer}, - BranchNodeMasks, BranchNodeMasksMap, + trie_cursor::TrieCursorFactory, + DecodedMultiProofV2, HashedPostState, Nibbles, ProofTrieNodeV2, }; use reth_trie_sparse::provider::{RevealedNode, TrieNodeProvider, TrieNodeProviderFactory}; use std::{ @@ -124,8 +108,6 @@ pub struct ProofWorkerHandle { storage_worker_count: usize, /// Total number of account workers spawned account_worker_count: usize, - /// Whether V2 storage proofs are enabled - v2_proofs_enabled: bool, } impl ProofWorkerHandle { @@ -138,7 +120,6 @@ impl ProofWorkerHandle { /// - `runtime`: The centralized runtime used to spawn blocking worker tasks /// - `task_ctx`: Shared context with database view and prefix sets /// - `halve_workers`: Whether to halve the worker pool size (for small blocks) - /// - `v2_proofs_enabled`: Whether to enable V2 storage proofs #[instrument( name = "ProofWorkerHandle::new", level = "debug", @@ -149,7 +130,6 @@ impl ProofWorkerHandle { runtime: &Runtime, task_ctx: ProofTaskCtx, halve_workers: bool, - v2_proofs_enabled: bool, ) -> Self where Factory: DatabaseProviderROFactory @@ -177,7 +157,6 @@ impl ProofWorkerHandle { storage_worker_count, account_worker_count, halve_workers, - ?v2_proofs_enabled, "Spawning proof worker pools" ); @@ -209,8 +188,7 @@ impl ProofWorkerHandle { metrics, #[cfg(feature = "metrics")] cursor_metrics, - ) - .with_v2_proofs(v2_proofs_enabled); + ); if let Err(error) = worker.run() { error!( target: "trie::proof_task", @@ -248,8 +226,7 @@ impl ProofWorkerHandle { metrics, #[cfg(feature = "metrics")] cursor_metrics, - ) - .with_v2_proofs(v2_proofs_enabled); + ); if let Err(error) = worker.run() { error!( target: "trie::proof_task", @@ -268,15 +245,9 @@ impl ProofWorkerHandle { account_available_workers, storage_worker_count, account_worker_count, - v2_proofs_enabled, } } - /// Returns whether V2 storage proofs are enabled for this worker pool. - pub const fn v2_proofs_enabled(&self) -> bool { - self.v2_proofs_enabled - } - /// Returns how many storage workers are currently available/idle. pub fn available_storage_workers(&self) -> usize { self.storage_available_workers.load(Ordering::Relaxed) @@ -329,7 +300,7 @@ impl ProofWorkerHandle { input: StorageProofInput, proof_result_sender: CrossbeamSender, ) -> Result<(), ProviderError> { - let hashed_address = input.hashed_address(); + let hashed_address = input.hashed_address; self.storage_work_tx .send(StorageWorkerJob::StorageProof { input, proof_result_sender }) .map_err(|err| { @@ -447,69 +418,6 @@ impl ProofTaskTx where Provider: TrieCursorFactory + HashedCursorFactory, { - /// Compute storage proof. - /// - /// Used by storage workers in the worker pool to compute storage proofs. - #[inline] - fn compute_legacy_storage_proof( - &self, - input: StorageProofInput, - trie_cursor_metrics: &mut TrieCursorMetricsCache, - hashed_cursor_metrics: &mut HashedCursorMetricsCache, - ) -> Result { - // Consume the input so we can move large collections (e.g. target slots) without cloning. - let StorageProofInput::Legacy { - hashed_address, - prefix_set, - target_slots, - with_branch_node_masks, - multi_added_removed_keys, - } = input - else { - panic!("compute_legacy_storage_proof only accepts StorageProofInput::Legacy") - }; - - // Get or create added/removed keys context - let multi_added_removed_keys = - multi_added_removed_keys.unwrap_or_else(|| Arc::new(MultiAddedRemovedKeys::new())); - let added_removed_keys = multi_added_removed_keys.get_storage(&hashed_address); - - let span = debug_span!( - target: "trie::proof_task", - "Storage proof calculation", - target_slots = ?target_slots.len(), - ); - let _span_guard = span.enter(); - - let proof_start = Instant::now(); - - // Compute raw storage multiproof - let raw_proof_result = - StorageProof::new_hashed(&self.provider, &self.provider, hashed_address) - .with_prefix_set_mut(PrefixSetMut::from(prefix_set.iter().copied())) - .with_branch_node_masks(with_branch_node_masks) - .with_added_removed_keys(added_removed_keys) - .with_trie_cursor_metrics(trie_cursor_metrics) - .with_hashed_cursor_metrics(hashed_cursor_metrics) - .storage_multiproof(target_slots); - trie_cursor_metrics.record_span("trie_cursor"); - hashed_cursor_metrics.record_span("hashed_cursor"); - - // Decode proof into DecodedStorageMultiProof - let decoded_result = - raw_proof_result.and_then(|raw_proof| raw_proof.try_into().map_err(Into::into))?; - - trace!( - target: "trie::proof_task", - hashed_address = ?hashed_address, - proof_time_us = proof_start.elapsed().as_micros(), - worker_id = self.id, - "Completed storage proof calculation" - ); - - Ok(StorageProofResult::Legacy { proof: decoded_result }) - } - fn compute_v2_storage_proof( &self, input: StorageProofInput, @@ -518,9 +426,7 @@ where ::StorageCursor<'_>, >, ) -> Result { - let StorageProofInput::V2 { hashed_address, mut targets } = input else { - panic!("compute_v2_storage_proof only accepts StorageProofInput::V2") - }; + let StorageProofInput { hashed_address, mut targets } = input; let span = debug_span!( target: "trie::proof_task", @@ -550,7 +456,7 @@ where "Completed V2 storage proof calculation" ); - Ok(StorageProofResult::V2 { proof, root }) + Ok(StorageProofResult { proof, root }) } /// Process a blinded storage node request. @@ -615,69 +521,6 @@ impl TrieNodeProvider for ProofTaskTrieNodeProvider { } } -/// Result of a multiproof calculation. -#[derive(Debug)] -pub enum ProofResult { - /// Legacy multiproof calculation result. - Legacy(DecodedMultiProof, ParallelTrieStats), - /// V2 multiproof calculation result. - V2(DecodedMultiProofV2), -} - -impl ProofResult { - /// Creates an empty [`ProofResult`] of the appropriate variant based on `v2_enabled`. - /// - /// Use this when constructing empty proofs (e.g., for state updates where all targets - /// were already fetched) to ensure consistency with the proof version being used. - pub fn empty(v2_enabled: bool) -> Self { - if v2_enabled { - Self::V2(DecodedMultiProofV2::default()) - } else { - let stats = ParallelTrieTracker::default().finish(); - Self::Legacy(DecodedMultiProof::default(), stats) - } - } - - /// Returns true if the result contains no proofs - pub fn is_empty(&self) -> bool { - match self { - Self::Legacy(proof, _) => proof.is_empty(), - Self::V2(proof) => proof.is_empty(), - } - } - - /// Extends the receiver with the value of the given results. - /// - /// # Panics - /// - /// This method panics if the two [`ProofResult`]s are not the same variant. - pub fn extend(&mut self, other: Self) { - match (self, other) { - (Self::Legacy(proof, _), Self::Legacy(other, _)) => proof.extend(other), - (Self::V2(proof), Self::V2(other)) => proof.extend(other), - _ => panic!("mismatched ProofResults, cannot extend one with the other"), - } - } - - /// Returns the number of account proofs. - pub fn account_proofs_len(&self) -> usize { - match self { - Self::Legacy(proof, _) => proof.account_subtree.len(), - Self::V2(proof) => proof.account_proofs.len(), - } - } - - /// Returns the total number of storage proofs - pub fn storage_proofs_len(&self) -> usize { - match self { - Self::Legacy(proof, _) => { - proof.storages.values().map(|p| p.subtree.len()).sum::() - } - Self::V2(proof) => proof.storage_proofs.values().map(|p| p.len()).sum::(), - } - } -} - /// Channel used by worker threads to deliver `ProofResultMessage` items back to /// `MultiProofTask`. /// @@ -692,8 +535,8 @@ pub type ProofResultSender = CrossbeamSender; pub struct ProofResultMessage { /// Sequence number for ordering proofs pub sequence_number: u64, - /// The proof calculation result (either account multiproof or storage proof) - pub result: Result, + /// The proof calculation result + pub result: Result, /// Time taken for the entire proof calculation (from dispatch to completion) pub elapsed: Duration, /// Original state update that triggered this proof @@ -730,26 +573,17 @@ impl ProofResultContext { /// The results of a storage proof calculation. #[derive(Debug)] -pub(crate) enum StorageProofResult { - Legacy { - /// The storage multiproof - proof: DecodedStorageMultiProof, - }, - V2 { - /// The calculated V2 proof nodes - proof: Vec, - /// The storage root calculated by the V2 proof - root: Option, - }, +pub(crate) struct StorageProofResult { + /// The calculated V2 proof nodes + pub proof: Vec, + /// The storage root calculated by the V2 proof + pub root: Option, } impl StorageProofResult { /// Returns the calculated root of the trie, if one can be calculated from the proof. const fn root(&self) -> Option { - match self { - Self::Legacy { proof } => Some(proof.root), - Self::V2 { root, .. } => *root, - } + self.root } } @@ -757,6 +591,7 @@ impl StorageProofResult { #[derive(Debug)] pub struct StorageProofResultMessage { /// The hashed address this storage proof belongs to + #[allow(dead_code)] pub(crate) hashed_address: B256, /// The storage proof calculation result pub(crate) result: Result, @@ -804,8 +639,6 @@ struct StorageProofWorker { /// Cursor metrics for this worker #[cfg(feature = "metrics")] cursor_metrics: ProofTaskCursorMetrics, - /// Set to true if V2 proofs are enabled. - v2_enabled: bool, } impl StorageProofWorker @@ -832,16 +665,9 @@ where metrics, #[cfg(feature = "metrics")] cursor_metrics, - v2_enabled: false, } } - /// Changes whether or not V2 proofs are enabled. - const fn with_v2_proofs(mut self, v2_enabled: bool) -> Self { - self.v2_enabled = v2_enabled; - self - } - /// Runs the worker loop, processing jobs until the channel closes. /// /// # Lifecycle @@ -873,13 +699,10 @@ where let mut storage_proofs_processed = 0u64; let mut storage_nodes_processed = 0u64; let mut cursor_metrics_cache = ProofTaskCursorMetricsCache::default(); - let mut v2_calculator = if self.v2_enabled { - let trie_cursor = proof_tx.provider.storage_trie_cursor(B256::ZERO)?; - let hashed_cursor = proof_tx.provider.hashed_storage_cursor(B256::ZERO)?; - Some(proof_v2::StorageProofCalculator::new_storage(trie_cursor, hashed_cursor)) - } else { - None - }; + let trie_cursor = proof_tx.provider.storage_trie_cursor(B256::ZERO)?; + let hashed_cursor = proof_tx.provider.hashed_storage_cursor(B256::ZERO)?; + let mut v2_calculator = + proof_v2::StorageProofCalculator::new_storage(trie_cursor, hashed_cursor); // Initially mark this worker as available. self.available_workers.fetch_add(1, Ordering::Relaxed); @@ -897,11 +720,10 @@ where StorageWorkerJob::StorageProof { input, proof_result_sender } => { self.process_storage_proof( &proof_tx, - v2_calculator.as_mut(), + &mut v2_calculator, input, proof_result_sender, &mut storage_proofs_processed, - &mut cursor_metrics_cache, ); } @@ -946,53 +768,28 @@ where fn process_storage_proof( &self, proof_tx: &ProofTaskTx, - v2_calculator: Option< - &mut proof_v2::StorageProofCalculator< - ::StorageTrieCursor<'_>, - ::StorageCursor<'_>, - >, + v2_calculator: &mut proof_v2::StorageProofCalculator< + ::StorageTrieCursor<'_>, + ::StorageCursor<'_>, >, input: StorageProofInput, proof_result_sender: CrossbeamSender, storage_proofs_processed: &mut u64, - cursor_metrics_cache: &mut ProofTaskCursorMetricsCache, ) where Provider: TrieCursorFactory + HashedCursorFactory, { - let mut trie_cursor_metrics = TrieCursorMetricsCache::default(); - let mut hashed_cursor_metrics = HashedCursorMetricsCache::default(); - let hashed_address = input.hashed_address(); + let hashed_address = input.hashed_address; let proof_start = Instant::now(); - let result = match &input { - StorageProofInput::Legacy { hashed_address, prefix_set, target_slots, .. } => { - trace!( - target: "trie::proof_task", - worker_id = self.worker_id, - hashed_address = ?hashed_address, - prefix_set_len = prefix_set.len(), - target_slots_len = target_slots.len(), - "Processing storage proof" - ); + trace!( + target: "trie::proof_task", + worker_id = self.worker_id, + hashed_address = ?hashed_address, + targets_len = input.targets.len(), + "Processing V2 storage proof" + ); - proof_tx.compute_legacy_storage_proof( - input, - &mut trie_cursor_metrics, - &mut hashed_cursor_metrics, - ) - } - StorageProofInput::V2 { hashed_address, targets } => { - trace!( - target: "trie::proof_task", - worker_id = self.worker_id, - hashed_address = ?hashed_address, - targets_len = targets.len(), - "Processing V2 storage proof" - ); - proof_tx - .compute_v2_storage_proof(input, v2_calculator.expect("v2 calculator provided")) - } - }; + let result = proof_tx.compute_v2_storage_proof(input, v2_calculator); let proof_elapsed = proof_start.elapsed(); *storage_proofs_processed += 1; @@ -1019,25 +816,9 @@ where hashed_address = ?hashed_address, proof_time_us = proof_elapsed.as_micros(), total_processed = storage_proofs_processed, - trie_cursor_duration_us = trie_cursor_metrics.total_duration.as_micros(), - hashed_cursor_duration_us = hashed_cursor_metrics.total_duration.as_micros(), - ?trie_cursor_metrics, - ?hashed_cursor_metrics, ?root, "Storage proof completed" ); - - #[cfg(feature = "metrics")] - { - // Accumulate per-proof metrics into the worker's cache - let per_proof_cache = ProofTaskCursorMetricsCache { - account_trie_cursor: TrieCursorMetricsCache::default(), - account_hashed_cursor: HashedCursorMetricsCache::default(), - storage_trie_cursor: trie_cursor_metrics, - storage_hashed_cursor: hashed_cursor_metrics, - }; - cursor_metrics_cache.extend(&per_proof_cache); - } } /// Processes a blinded storage node lookup request. @@ -1111,8 +892,6 @@ struct AccountProofWorker { /// Cursor metrics for this worker #[cfg(feature = "metrics")] cursor_metrics: ProofTaskCursorMetrics, - /// Set to true if V2 proofs are enabled. - v2_enabled: bool, } impl AccountProofWorker @@ -1142,16 +921,9 @@ where metrics, #[cfg(feature = "metrics")] cursor_metrics, - v2_enabled: false, } } - /// Changes whether or not V2 proofs are enabled. - const fn with_v2_proofs(mut self, v2_enabled: bool) -> Self { - self.v2_enabled = v2_enabled; - self - } - /// Runs the worker loop, processing jobs until the channel closes. /// /// # Lifecycle @@ -1184,30 +956,25 @@ where // Create both account and storage calculators for V2 proofs. // The storage calculator is wrapped in Rc> for sharing with value encoders. - let (mut v2_account_calculator, v2_storage_calculator) = if self.v2_enabled { - let account_trie_cursor = provider.account_trie_cursor()?; - let account_hashed_cursor = provider.hashed_account_cursor()?; + let account_trie_cursor = provider.account_trie_cursor()?; + let account_hashed_cursor = provider.hashed_account_cursor()?; - let storage_trie_cursor = provider.storage_trie_cursor(B256::ZERO)?; - let storage_hashed_cursor = provider.hashed_storage_cursor(B256::ZERO)?; + let storage_trie_cursor = provider.storage_trie_cursor(B256::ZERO)?; + let storage_hashed_cursor = provider.hashed_storage_cursor(B256::ZERO)?; - ( - Some(proof_v2::ProofCalculator::< - _, - _, - AsyncAccountValueEncoder< - ::StorageTrieCursor<'_>, - ::StorageCursor<'_>, - >, - >::new(account_trie_cursor, account_hashed_cursor)), - Some(Rc::new(RefCell::new(proof_v2::StorageProofCalculator::new_storage( - storage_trie_cursor, - storage_hashed_cursor, - )))), - ) - } else { - (None, None) - }; + let mut v2_account_calculator = proof_v2::ProofCalculator::< + _, + _, + AsyncAccountValueEncoder< + ::StorageTrieCursor<'_>, + ::StorageCursor<'_>, + >, + >::new(account_trie_cursor, account_hashed_cursor); + let v2_storage_calculator = + Rc::new(RefCell::new(proof_v2::StorageProofCalculator::new_storage( + storage_trie_cursor, + storage_hashed_cursor, + ))); // Count this worker as available only after successful initialization. self.available_workers.fetch_add(1, Ordering::Relaxed); @@ -1224,9 +991,8 @@ where match job { AccountWorkerJob::AccountMultiproof { input } => { - let value_encoder_stats = self.process_account_multiproof( - &provider, - v2_account_calculator.as_mut(), + let value_encoder_stats = self.process_account_multiproof::( + &mut v2_account_calculator, v2_storage_calculator.clone(), *input, &mut account_proofs_processed, @@ -1273,78 +1039,12 @@ where Ok(()) } - fn compute_legacy_account_multiproof( - &self, - provider: &Provider, - targets: MultiProofTargets, - mut prefix_sets: TriePrefixSets, - collect_branch_node_masks: bool, - multi_added_removed_keys: Option>, - proof_cursor_metrics: &mut ProofTaskCursorMetricsCache, - ) -> Result<(ProofResult, Duration), ParallelStateRootError> - where - Provider: TrieCursorFactory + HashedCursorFactory, - { - let span = debug_span!( - target: "trie::proof_task", - "Account multiproof calculation", - targets = targets.len(), - num_slots = targets.values().map(|slots| slots.len()).sum::(), - ); - let _span_guard = span.enter(); - - trace!( - target: "trie::proof_task", - "Processing account multiproof" - ); - - let mut tracker = ParallelTrieTracker::default(); - - let mut storage_prefix_sets = std::mem::take(&mut prefix_sets.storage_prefix_sets); - - let storage_root_targets_len = - StorageRootTargets::count(&prefix_sets.account_prefix_set, &storage_prefix_sets); - - tracker.set_precomputed_storage_roots(storage_root_targets_len as u64); - - let storage_proof_receivers = dispatch_storage_proofs( - &self.storage_work_tx, - &targets, - &mut storage_prefix_sets, - collect_branch_node_masks, - multi_added_removed_keys.as_ref(), - )?; - - let account_prefix_set = std::mem::take(&mut prefix_sets.account_prefix_set); - - let ctx = AccountMultiproofParams { - targets: &targets, - prefix_set: account_prefix_set, - collect_branch_node_masks, - multi_added_removed_keys: multi_added_removed_keys.as_ref(), - storage_proof_receivers, - cached_storage_roots: &self.cached_storage_roots, - }; - - let mut storage_wait_time = Duration::ZERO; - let result = build_account_multiproof_with_storage_roots( - provider, - ctx, - &mut tracker, - proof_cursor_metrics, - &mut storage_wait_time, - )?; - - let stats = tracker.finish(); - Ok((ProofResult::Legacy(result, stats), storage_wait_time)) - } - fn compute_v2_account_multiproof<'a, Provider>( &self, v2_account_calculator: &mut V2AccountProofCalculator<'a, Provider>, v2_storage_calculator: Rc>>, targets: MultiProofTargetsV2, - ) -> Result<(ProofResult, ValueEncoderStats), ParallelStateRootError> + ) -> Result<(DecodedMultiProofV2, ValueEncoderStats), ParallelStateRootError> where Provider: TrieCursorFactory + HashedCursorFactory + 'a, { @@ -1376,7 +1076,7 @@ where let proof = DecodedMultiProofV2 { account_proofs, storage_proofs }; - Ok((ProofResult::V2(proof), value_encoder_stats)) + Ok((proof, value_encoder_stats)) } /// Processes an account multiproof request. @@ -1384,9 +1084,8 @@ where /// Returns stats from the value encoder used during proof computation. fn process_account_multiproof<'a, Provider>( &self, - provider: &Provider, - v2_account_calculator: Option<&mut V2AccountProofCalculator<'a, Provider>>, - v2_storage_calculator: Option>>>, + v2_account_calculator: &mut V2AccountProofCalculator<'a, Provider>, + v2_storage_calculator: Rc>>, input: AccountMultiproofInput, account_proofs_processed: &mut u64, cursor_metrics_cache: &mut ProofTaskCursorMetricsCache, @@ -1394,45 +1093,17 @@ where where Provider: TrieCursorFactory + HashedCursorFactory + 'a, { - let mut proof_cursor_metrics = ProofTaskCursorMetricsCache::default(); + let proof_cursor_metrics = ProofTaskCursorMetricsCache::default(); let proof_start = Instant::now(); - let (proof_result_sender, result, value_encoder_stats) = match input { - AccountMultiproofInput::Legacy { - targets, - prefix_sets, - collect_branch_node_masks, - multi_added_removed_keys, - proof_result_sender, - } => { - let (result, value_encoder_stats) = match self.compute_legacy_account_multiproof( - provider, - targets, - prefix_sets, - collect_branch_node_masks, - multi_added_removed_keys, - &mut proof_cursor_metrics, - ) { - Ok((proof, wait_time)) => ( - Ok(proof), - ValueEncoderStats { storage_wait_time: wait_time, ..Default::default() }, - ), - Err(e) => (Err(e), ValueEncoderStats::default()), - }; - (proof_result_sender, result, value_encoder_stats) - } - AccountMultiproofInput::V2 { targets, proof_result_sender } => { - let (result, value_encoder_stats) = match self - .compute_v2_account_multiproof::( - v2_account_calculator.expect("v2 account calculator provided"), - v2_storage_calculator.expect("v2 storage calculator provided"), - targets, - ) { - Ok((proof, stats)) => (Ok(proof), stats), - Err(e) => (Err(e), ValueEncoderStats::default()), - }; - (proof_result_sender, result, value_encoder_stats) - } + let AccountMultiproofInput { targets, proof_result_sender } = input; + let (result, value_encoder_stats) = match self.compute_v2_account_multiproof::( + v2_account_calculator, + v2_storage_calculator, + targets, + ) { + Ok((proof, stats)) => (Ok(proof), stats), + Err(e) => (Err(e), ValueEncoderStats::default()), }; let ProofResultContext { @@ -1537,245 +1208,6 @@ where } } -/// Builds an account multiproof by consuming storage proof receivers lazily during trie walk. -/// -/// This is a helper function used by account workers to build the account subtree proof -/// while storage proofs are still being computed. Receivers are consumed only when needed, -/// enabling interleaved parallelism between account trie traversal and storage proof computation. -/// -/// Returns a `DecodedMultiProof` containing the account subtree and storage proofs. -/// Also accumulates the time spent waiting for storage proofs into `storage_wait_time`. -fn build_account_multiproof_with_storage_roots

( - provider: &P, - ctx: AccountMultiproofParams<'_>, - tracker: &mut ParallelTrieTracker, - proof_cursor_metrics: &mut ProofTaskCursorMetricsCache, - storage_wait_time: &mut Duration, -) -> Result -where - P: TrieCursorFactory + HashedCursorFactory, -{ - let accounts_added_removed_keys = - ctx.multi_added_removed_keys.as_ref().map(|keys| keys.get_accounts()); - - // Wrap account trie cursor with instrumented cursor - let account_trie_cursor = provider.account_trie_cursor().map_err(ProviderError::Database)?; - let account_trie_cursor = InstrumentedTrieCursor::new( - account_trie_cursor, - &mut proof_cursor_metrics.account_trie_cursor, - ); - - // Create the walker. - let walker = TrieWalker::<_>::state_trie(account_trie_cursor, ctx.prefix_set) - .with_added_removed_keys(accounts_added_removed_keys) - .with_deletions_retained(true); - - // Create a hash builder to rebuild the root node since it is not available in the database. - let retainer = ctx - .targets - .keys() - .map(Nibbles::unpack) - .collect::() - .with_added_removed_keys(accounts_added_removed_keys); - let mut hash_builder = HashBuilder::default() - .with_proof_retainer(retainer) - .with_updates(ctx.collect_branch_node_masks); - - // Initialize storage multiproofs map with pre-allocated capacity. - // Proofs will be inserted as they're consumed from receivers during trie walk. - let mut collected_decoded_storages: B256Map = - B256Map::with_capacity_and_hasher(ctx.targets.len(), Default::default()); - let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE); - - // Wrap account hashed cursor with instrumented cursor - let account_hashed_cursor = - provider.hashed_account_cursor().map_err(ProviderError::Database)?; - let account_hashed_cursor = InstrumentedHashedCursor::new( - account_hashed_cursor, - &mut proof_cursor_metrics.account_hashed_cursor, - ); - - let mut account_node_iter = TrieNodeIter::state_trie(walker, account_hashed_cursor); - - let mut storage_proof_receivers = ctx.storage_proof_receivers; - - while let Some(account_node) = account_node_iter.try_next().map_err(ProviderError::Database)? { - match account_node { - TrieElement::Branch(node) => { - hash_builder.add_branch(node.key, node.value, node.children_are_in_trie); - } - TrieElement::Leaf(hashed_address, account) => { - let root = match storage_proof_receivers.remove(&hashed_address) { - Some(receiver) => { - let _guard = debug_span!( - target: "trie::proof_task", - "Waiting for storage proof", - ); - // Block on this specific storage proof receiver - enables interleaved - // parallelism - let wait_start = Instant::now(); - let proof_msg = receiver.recv().map_err(|_| { - ParallelStateRootError::StorageRoot( - reth_execution_errors::StorageRootError::Database( - DatabaseError::Other(format!( - "Storage proof channel closed for {hashed_address}" - )), - ), - ) - })?; - *storage_wait_time += wait_start.elapsed(); - - drop(_guard); - - // Extract storage proof from the result - debug_assert_eq!( - proof_msg.hashed_address, hashed_address, - "storage worker must return same address" - ); - let StorageProofResult::Legacy { proof } = proof_msg.result? else { - unreachable!("v2 result in legacy worker") - }; - let root = proof.root; - collected_decoded_storages.insert(hashed_address, proof); - root - } - // Since we do not store all intermediate nodes in the database, there might - // be a possibility of re-adding a non-modified leaf to the hash builder. - None => { - tracker.inc_missed_leaves(); - - match ctx.cached_storage_roots.entry(hashed_address) { - dashmap::Entry::Occupied(occ) => *occ.get(), - dashmap::Entry::Vacant(vac) => { - let root = - StorageProof::new_hashed(provider, provider, hashed_address) - .with_prefix_set_mut(Default::default()) - .with_trie_cursor_metrics( - &mut proof_cursor_metrics.storage_trie_cursor, - ) - .with_hashed_cursor_metrics( - &mut proof_cursor_metrics.storage_hashed_cursor, - ) - .storage_multiproof( - ctx.targets - .get(&hashed_address) - .cloned() - .unwrap_or_default(), - ) - .map_err(|e| { - ParallelStateRootError::StorageRoot( - reth_execution_errors::StorageRootError::Database( - DatabaseError::Other(e.to_string()), - ), - ) - })? - .root; - - vac.insert(root); - root - } - } - } - }; - - // Encode account - account_rlp.clear(); - let account = account.into_trie_account(root); - account.encode(&mut account_rlp as &mut dyn BufMut); - - hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp); - } - } - } - - let _ = hash_builder.root(); - - let account_subtree_raw_nodes = hash_builder.take_proof_nodes(); - let decoded_account_subtree = DecodedProofNodes::try_from(account_subtree_raw_nodes)?; - - let branch_node_masks = if ctx.collect_branch_node_masks { - let updated_branch_nodes = hash_builder.updated_branch_nodes.unwrap_or_default(); - updated_branch_nodes - .into_iter() - .map(|(path, node)| { - (path, BranchNodeMasks { hash_mask: node.hash_mask, tree_mask: node.tree_mask }) - }) - .collect() - } else { - BranchNodeMasksMap::default() - }; - - // Consume remaining storage proof receivers for accounts not encountered during trie walk. - // Done last to allow storage workers more time to complete while we finalized the account trie. - for (hashed_address, receiver) in storage_proof_receivers { - let wait_start = Instant::now(); - if let Ok(proof_msg) = receiver.recv() { - *storage_wait_time += wait_start.elapsed(); - let StorageProofResult::Legacy { proof } = proof_msg.result? else { - unreachable!("v2 result in legacy worker") - }; - collected_decoded_storages.insert(hashed_address, proof); - } - } - - Ok(DecodedMultiProof { - account_subtree: decoded_account_subtree, - branch_node_masks, - storages: collected_decoded_storages, - }) -} -/// Queues storage proofs for all accounts in the targets and returns receivers. -/// -/// This function queues all storage proof tasks to the worker pool but returns immediately -/// with receivers, allowing the account trie walk to proceed in parallel with storage proof -/// computation. This enables interleaved parallelism for better performance. -/// -/// Propagates errors up if queuing fails. Receivers must be consumed by the caller. -fn dispatch_storage_proofs( - storage_work_tx: &CrossbeamSender, - targets: &MultiProofTargets, - storage_prefix_sets: &mut B256Map, - with_branch_node_masks: bool, - multi_added_removed_keys: Option<&Arc>, -) -> Result>, ParallelStateRootError> { - let mut storage_proof_receivers = - B256Map::with_capacity_and_hasher(targets.len(), Default::default()); - - let mut sorted_targets: Vec<_> = targets.iter().collect(); - sorted_targets.sort_unstable_by_key(|(addr, _)| *addr); - - // Dispatch all storage proofs to worker pool - for (hashed_address, target_slots) in sorted_targets { - // Create channel for receiving ProofResultMessage - let (result_tx, result_rx) = crossbeam_channel::unbounded(); - - // Create computation input based on V2 flag - let prefix_set = storage_prefix_sets.remove(hashed_address).unwrap_or_default(); - let input = StorageProofInput::legacy( - *hashed_address, - prefix_set, - target_slots.clone(), - with_branch_node_masks, - multi_added_removed_keys.cloned(), - ); - - // Always dispatch a storage proof so we obtain the storage root even when no slots are - // requested. - storage_work_tx - .send(StorageWorkerJob::StorageProof { input, proof_result_sender: result_tx }) - .map_err(|_| { - ParallelStateRootError::Other(format!( - "Failed to queue storage proof for {}: storage worker pool unavailable", - hashed_address - )) - })?; - - storage_proof_receivers.insert(*hashed_address, result_rx); - } - - Ok(storage_proof_receivers) -} - /// Queues V2 storage proofs for all accounts in the targets and returns receivers. /// /// This function queues all storage proof tasks to the worker pool but returns immediately @@ -1854,116 +1286,36 @@ fn dispatch_v2_storage_proofs( /// Input parameters for storage proof computation. #[derive(Debug)] -pub enum StorageProofInput { - /// Legacy storage proof variant - Legacy { - /// The hashed address for which the proof is calculated. - hashed_address: B256, - /// The prefix set for the proof calculation. - prefix_set: PrefixSet, - /// The target slots for the proof calculation. - target_slots: B256Set, - /// Whether or not to collect branch node masks - with_branch_node_masks: bool, - /// Provided by the user to give the necessary context to retain extra proofs. - multi_added_removed_keys: Option>, - }, - /// V2 storage proof variant - V2 { - /// The hashed address for which the proof is calculated. - hashed_address: B256, - /// The set of proof targets - targets: Vec, - }, +pub struct StorageProofInput { + /// The hashed address for which the proof is calculated. + pub hashed_address: B256, + /// The set of proof targets + pub targets: Vec, } impl StorageProofInput { - /// Creates a legacy [`StorageProofInput`] with the given hashed address, prefix set, and target - /// slots. - pub const fn legacy( - hashed_address: B256, - prefix_set: PrefixSet, - target_slots: B256Set, - with_branch_node_masks: bool, - multi_added_removed_keys: Option>, - ) -> Self { - Self::Legacy { - hashed_address, - prefix_set, - target_slots, - with_branch_node_masks, - multi_added_removed_keys, - } - } - /// Creates a new [`StorageProofInput`] with the given hashed address and target slots. pub const fn new(hashed_address: B256, targets: Vec) -> Self { - Self::V2 { hashed_address, targets } - } - - /// Returns the targeted hashed address. - pub const fn hashed_address(&self) -> B256 { - match self { - Self::Legacy { hashed_address, .. } | Self::V2 { hashed_address, .. } => { - *hashed_address - } - } + Self { hashed_address, targets } } } /// Input parameters for account multiproof computation. #[derive(Debug)] -pub enum AccountMultiproofInput { - /// Legacy account multiproof proof variant - Legacy { - /// The targets for which to compute the multiproof. - targets: MultiProofTargets, - /// The prefix sets for the proof calculation. - prefix_sets: TriePrefixSets, - /// Whether or not to collect branch node masks. - collect_branch_node_masks: bool, - /// Provided by the user to give the necessary context to retain extra proofs. - multi_added_removed_keys: Option>, - /// Context for sending the proof result. - proof_result_sender: ProofResultContext, - }, - /// V2 account multiproof variant - V2 { - /// The targets for which to compute the multiproof. - targets: MultiProofTargetsV2, - /// Context for sending the proof result. - proof_result_sender: ProofResultContext, - }, +pub struct AccountMultiproofInput { + /// The targets for which to compute the multiproof. + pub targets: MultiProofTargetsV2, + /// Context for sending the proof result. + pub proof_result_sender: ProofResultContext, } impl AccountMultiproofInput { /// Returns the [`ProofResultContext`] for this input, consuming the input. fn into_proof_result_sender(self) -> ProofResultContext { - match self { - Self::Legacy { proof_result_sender, .. } | Self::V2 { proof_result_sender, .. } => { - proof_result_sender - } - } + self.proof_result_sender } } -/// Parameters for building an account multiproof with pre-computed storage roots. -struct AccountMultiproofParams<'a> { - /// The targets for which to compute the multiproof. - targets: &'a MultiProofTargets, - /// The prefix set for the account trie walk. - prefix_set: PrefixSet, - /// Whether or not to collect branch node masks. - collect_branch_node_masks: bool, - /// Provided by the user to give the necessary context to retain extra proofs. - multi_added_removed_keys: Option<&'a Arc>, - /// Receivers for storage proofs being computed in parallel. - storage_proof_receivers: B256Map>, - /// Cached storage roots. This will be used to read storage roots for missed leaves, as well as - /// to write calculated storage roots. - cached_storage_roots: &'a DashMap, -} - /// Internal message for account workers. #[derive(Debug)] enum AccountWorkerJob { @@ -2002,7 +1354,7 @@ mod tests { let ctx = test_ctx(factory); let runtime = reth_tasks::Runtime::test(); - let proof_handle = ProofWorkerHandle::new(&runtime, ctx, false, false); + let proof_handle = ProofWorkerHandle::new(&runtime, ctx, false); // Verify handle can be cloned let _cloned_handle = proof_handle.clone(); diff --git a/crates/trie/parallel/src/value_encoder.rs b/crates/trie/parallel/src/value_encoder.rs index 5e63ffddf8..14912ed99a 100644 --- a/crates/trie/parallel/src/value_encoder.rs +++ b/crates/trie/parallel/src/value_encoder.rs @@ -1,4 +1,4 @@ -use crate::proof_task::{StorageProofResult, StorageProofResultMessage}; +use crate::proof_task::StorageProofResultMessage; use alloy_primitives::{map::B256Map, B256}; use alloy_rlp::Encodable; use core::cell::RefCell; @@ -109,11 +109,7 @@ impl Drop for AsyncAccountDeferredValueEncoder { stats.borrow_mut().storage_wait_time += wait_start.elapsed(); - let StorageProofResult::V2 { proof, .. } = result else { - panic!("StorageProofResult is not V2: {result:?}") - }; - - storage_proof_results.borrow_mut().insert(*hashed_address, proof); + storage_proof_results.borrow_mut().insert(*hashed_address, result.proof); Ok(()) })() } else { @@ -159,13 +155,9 @@ where .result?; stats.borrow_mut().storage_wait_time += wait_start.elapsed(); - let StorageProofResult::V2 { root, proof } = result else { - panic!("StorageProofResult is not V2: {result:?}") - }; + storage_proof_results.borrow_mut().insert(hashed_address, result.proof); - storage_proof_results.borrow_mut().insert(hashed_address, proof); - - let root = match root { + let root = match result.root { Some(root) => root, None => { // In `compute_v2_account_multiproof` we ensure that all dispatched storage @@ -290,11 +282,7 @@ impl AsyncAccountValueEncoder { .result?; stats.storage_wait_time += wait_start.elapsed(); - let StorageProofResult::V2 { proof, .. } = result else { - panic!("StorageProofResult is not V2: {result:?}") - }; - - storage_proof_results.insert(*hashed_address, proof); + storage_proof_results.insert(*hashed_address, result.proof); } Ok((storage_proof_results, stats)) diff --git a/docs/vocs/docs/pages/cli/reth/node.mdx b/docs/vocs/docs/pages/cli/reth/node.mdx index 114bfa36bc..3ec6de8a0b 100644 --- a/docs/vocs/docs/pages/cli/reth/node.mdx +++ b/docs/vocs/docs/pages/cli/reth/node.mdx @@ -984,15 +984,9 @@ Engine: --engine.prewarming-threads Configure the number of prewarming threads. If not specified, defaults to available parallelism - --engine.disable-proof-v2 - Disable V2 storage proofs for state root calculations - --engine.disable-cache-metrics Disable cache metrics recording, which can take up to 50ms with large cached state - --engine.disable-trie-cache - Disable sparse trie cache - --engine.sparse-trie-prune-depth Sparse trie prune depth