perf(tree): worker pooling for storage in multiproof generation

crates/engine/primitives/src/config.rs

@@ -6,9 +6,22 @@ pub const DEFAULT_PERSISTENCE_THRESHOLD: u64 = 2;
 /// How close to the canonical head we persist blocks.
 pub const DEFAULT_MEMORY_BLOCK_BUFFER_TARGET: u64 = 0;
 
-/// Default maximum concurrency for proof tasks
+/// Default maximum concurrency for on-demand proof tasks (blinded nodes)
 pub const DEFAULT_MAX_PROOF_TASK_CONCURRENCY: u64 = 256;
 
+/// Returns the default number of storage worker threads based on available parallelism.
+/// Defaults to half of available parallelism, clamped between 2 and 8.
+fn default_storage_worker_count() -> usize {
+    #[cfg(feature = "std")]
+    {
+        std::thread::available_parallelism().map(|n| (n.get() / 2).clamp(2, 8)).unwrap_or(4)
+    }
+    #[cfg(not(feature = "std"))]
+    {
+        4
+    }
+}
+
 /// The size of proof targets chunk to spawn in one multiproof calculation.
 pub const DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE: usize = 10;
 
@@ -109,6 +122,8 @@ pub struct TreeConfig {
     prewarm_max_concurrency: usize,
     /// Whether to unwind canonical header to ancestor during forkchoice updates.
     allow_unwind_canonical_header: bool,
+    /// Number of storage proof worker threads.
+    storage_worker_count: usize,
 }
 
 impl Default for TreeConfig {
@@ -135,6 +150,7 @@ impl Default for TreeConfig {
             always_process_payload_attributes_on_canonical_head: false,
             prewarm_max_concurrency: DEFAULT_PREWARM_MAX_CONCURRENCY,
             allow_unwind_canonical_header: false,
+            storage_worker_count: default_storage_worker_count(),
         }
     }
 }
@@ -164,7 +180,9 @@ impl TreeConfig {
         always_process_payload_attributes_on_canonical_head: bool,
         prewarm_max_concurrency: usize,
         allow_unwind_canonical_header: bool,
+        storage_worker_count: usize,
     ) -> Self {
+        assert!(max_proof_task_concurrency > 0, "max_proof_task_concurrency must be at least 1");
         Self {
             persistence_threshold,
             memory_block_buffer_target,
@@ -187,6 +205,7 @@ impl TreeConfig {
             always_process_payload_attributes_on_canonical_head,
             prewarm_max_concurrency,
             allow_unwind_canonical_header,
+            storage_worker_count,
         }
     }
 
@@ -394,6 +413,7 @@ impl TreeConfig {
         mut self,
         max_proof_task_concurrency: u64,
     ) -> Self {
+        assert!(max_proof_task_concurrency > 0, "max_proof_task_concurrency must be at least 1");
         self.max_proof_task_concurrency = max_proof_task_concurrency;
         self
     }
@@ -452,4 +472,15 @@ impl TreeConfig {
     pub const fn prewarm_max_concurrency(&self) -> usize {
         self.prewarm_max_concurrency
     }
+
+    /// Return the number of storage proof worker threads.
+    pub const fn storage_worker_count(&self) -> usize {
+        self.storage_worker_count
+    }
+
+    /// Setter for the number of storage proof worker threads.
+    pub const fn with_storage_worker_count(mut self, storage_worker_count: usize) -> Self {
+        self.storage_worker_count = storage_worker_count;
+        self
+    }
 }

crates/engine/tree/benches/state_root_task.rs

@@ -228,16 +228,22 @@ fn bench_state_root(c: &mut Criterion) {
                     },
                     |(genesis_hash, mut payload_processor, provider, state_updates)| {
                         black_box({
-                            let mut handle = payload_processor.spawn(
-                                Default::default(),
-                                core::iter::empty::<
-                                    Result<Recovered<TransactionSigned>, core::convert::Infallible>,
-                                >(),
-                                StateProviderBuilder::new(provider.clone(), genesis_hash, None),
-                                ConsistentDbView::new_with_latest_tip(provider).unwrap(),
-                                TrieInput::default(),
-                                &TreeConfig::default(),
-                            );
+                            let mut handle = payload_processor
+                                .spawn(
+                                    Default::default(),
+                                    core::iter::empty::<
+                                        Result<
+                                            Recovered<TransactionSigned>,
+                                            core::convert::Infallible,
+                                        >,
+                                    >(),
+                                    StateProviderBuilder::new(provider.clone(), genesis_hash, None),
+                                    ConsistentDbView::new_with_latest_tip(provider).unwrap(),
+                                    TrieInput::default(),
+                                    &TreeConfig::default(),
+                                )
+                                .map_err(|(err, ..)| err)
+                                .expect("failed to spawn payload processor");
 
                             let mut state_hook = handle.state_hook();
 

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

@@ -45,7 +45,7 @@ use std::sync::{
     mpsc::{self, channel, Sender},
     Arc,
 };
-use tracing::{debug, instrument};
+use tracing::{debug, instrument, warn};
 
 mod configured_sparse_trie;
 pub mod executor;
@@ -166,6 +166,10 @@ where
     ///
     /// This returns a handle to await the final state root and to interact with the tasks (e.g.
     /// canceling)
+    ///
+    /// Returns an error with the original transactions iterator if the proof task manager fails to
+    /// initialize.
+    #[allow(clippy::type_complexity)]
     pub fn spawn<P, I: ExecutableTxIterator<Evm>>(
         &mut self,
         env: ExecutionEnv<Evm>,
@@ -174,7 +178,10 @@ where
         consistent_view: ConsistentDbView<P>,
         trie_input: TrieInput,
         config: &TreeConfig,
-    ) -> PayloadHandle<WithTxEnv<TxEnvFor<Evm>, I::Tx>, I::Error>
+    ) -> Result<
+        PayloadHandle<WithTxEnv<TxEnvFor<Evm>, I::Tx>, I::Error>,
+        (reth_provider::ProviderError, I, ExecutionEnv<Evm>, StateProviderBuilder<N, P>),
+    >
     where
         P: DatabaseProviderFactory<Provider: BlockReader>
             + BlockReader
@@ -196,12 +203,25 @@ where
             state_root_config.prefix_sets.clone(),
         );
         let max_proof_task_concurrency = config.max_proof_task_concurrency() as usize;
-        let proof_task = ProofTaskManager::new(
+        let storage_worker_count = config.storage_worker_count();
+        let proof_task = match ProofTaskManager::new(
             self.executor.handle().clone(),
             state_root_config.consistent_view.clone(),
             task_ctx,
             max_proof_task_concurrency,
-        );
+            storage_worker_count,
+        ) {
+            Ok(task) => task,
+            Err(error) => {
+                // Fall back to parallel state root if proof task manager fails to initialize
+                tracing::error!(
+                    target: "engine::tree",
+                    ?error,
+                    "Failed to initialize proof task manager, falling back to parallel state root"
+                );
+                return Err((error, transactions, env, provider_builder));
+            }
+        };
 
         // We set it to half of the proof task concurrency, because often for each multiproof we
         // spawn one Tokio task for the account proof, and one Tokio task for the storage proof.
@@ -252,12 +272,12 @@ where
             }
         });
 
-        PayloadHandle {
+        Ok(PayloadHandle {
             to_multi_proof,
             prewarm_handle,
             state_root: Some(state_root_rx),
             transactions: execution_rx,
-        }
+        })
     }
 
     /// Spawns a task that exclusively handles cache prewarming for transaction execution.
@@ -466,6 +486,11 @@ impl<Tx, Err> PayloadHandle<Tx, Err> {
             .map_err(|_| ParallelStateRootError::Other("sparse trie task dropped".to_string()))?
     }
 
+    /// Returns `true` if the handle is connected to a background state root task.
+    pub const fn supports_state_root(&self) -> bool {
+        self.state_root.is_some()
+    }
+
     /// Returns a state hook to be used to send state updates to this task.
     ///
     /// If a multiproof task is spawned the hook will notify it about new states.
@@ -857,14 +882,20 @@ mod tests {
             PrecompileCacheMap::default(),
         );
         let provider = BlockchainProvider::new(factory).unwrap();
-        let mut handle = payload_processor.spawn(
-            Default::default(),
-            core::iter::empty::<Result<Recovered<TransactionSigned>, core::convert::Infallible>>(),
-            StateProviderBuilder::new(provider.clone(), genesis_hash, None),
-            ConsistentDbView::new_with_latest_tip(provider).unwrap(),
-            TrieInput::from_state(hashed_state),
-            &TreeConfig::default(),
-        );
+        let mut handle =
+            payload_processor
+                .spawn(
+                    Default::default(),
+                    core::iter::empty::<
+                        Result<Recovered<TransactionSigned>, core::convert::Infallible>,
+                    >(),
+                    StateProviderBuilder::new(provider.clone(), genesis_hash, None),
+                    ConsistentDbView::new_with_latest_tip(provider).unwrap(),
+                    TrieInput::from_state(hashed_state),
+                    &TreeConfig::default(),
+                )
+                .map_err(|(err, ..)| err)
+                .expect("failed to spawn payload processor");
 
         let mut state_hook = handle.state_hook();
 

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

@@ -1236,7 +1236,9 @@ mod tests {
             config.consistent_view.clone(),
             task_ctx,
             1,
-        );
+            1,
+        )
+        .expect("Failed to create ProofTaskManager");
         let channel = channel();
 
         MultiProofTask::new(config, executor, proof_task.handle(), channel.0, 1, None)

crates/engine/tree/src/tree/payload_validator.rs

@@ -877,17 +877,37 @@ where
                 // too expensive because it requires walking all paths in every proof.
                 let spawn_start = Instant::now();
                 let (handle, strategy) = if trie_input.prefix_sets.is_empty() {
-                    (
-                        self.payload_processor.spawn(
-                            env,
-                            txs,
-                            provider_builder,
-                            consistent_view,
-                            trie_input,
-                            &self.config,
-                        ),
-                        StateRootStrategy::StateRootTask,
-                    )
+                    match self.payload_processor.spawn(
+                        env,
+                        txs,
+                        provider_builder,
+                        consistent_view,
+                        trie_input,
+                        &self.config,
+                    ) {
+                        Ok(handle) => {
+                            // Successfully spawned with state root task support
+                            (handle, StateRootStrategy::StateRootTask)
+                        }
+                        Err((error, txs, env, provider_builder)) => {
+                            // Failed to initialize proof task manager, fallback to parallel state
+                            // root
+                            error!(
+                                target: "engine::tree",
+                                block=?block_num_hash,
+                                ?error,
+                                "Failed to initialize proof task manager, falling back to parallel state root"
+                            );
+                            (
+                                self.payload_processor.spawn_cache_exclusive(
+                                    env,
+                                    txs,
+                                    provider_builder,
+                                ),
+                                StateRootStrategy::Parallel,
+                            )
+                        }
+                    }
                 // if prefix sets are not empty, we spawn a task that exclusively handles cache
                 // prewarming for transaction execution
                 } else {

crates/trie/parallel/Cargo.toml

@@ -36,6 +36,7 @@ derive_more.workspace = true
 rayon.workspace = true
 itertools.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread"] }
+crossbeam-channel.workspace = true
 
 # `metrics` feature
 reth-metrics = { workspace = true, optional = true }

crates/trie/parallel/src/proof.rs

@@ -448,7 +448,8 @@ mod tests {
         let task_ctx =
             ProofTaskCtx::new(Default::default(), Default::default(), Default::default());
         let proof_task =
-            ProofTaskManager::new(rt.handle().clone(), consistent_view.clone(), task_ctx, 1);
+            ProofTaskManager::new(rt.handle().clone(), consistent_view.clone(), task_ctx, 1, 1)
+                .unwrap();
         let proof_task_handle = proof_task.handle();
 
         // keep the join handle around to make sure it does not return any errors