refactor(trie): remove proof task manager

crates/engine/primitives/src/config.rs

@@ -6,9 +6,36 @@ 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;
 
+/// Minimum number of workers we allow configuring explicitly.
+pub const MIN_WORKER_COUNT: usize = 2;
+
+/// Returns the default number of storage worker threads based on available parallelism.
+fn default_storage_worker_count() -> usize {
+    #[cfg(feature = "std")]
+    {
+        std::thread::available_parallelism()
+            .map(|n| (n.get() * 2).clamp(MIN_WORKER_COUNT, 64))
+            .unwrap_or(8)
+    }
+    #[cfg(not(feature = "std"))]
+    {
+        8
+    }
+}
+
+/// Returns the default number of account worker threads optimized for I/O-bound coordination.
+///
+/// Account workers primarily coordinate storage proof collection and account trie traversal.
+/// They spend significant time blocked on `receiver.recv()` calls waiting for storage proofs,
+/// so we use higher concurrency (1.5x storage workers) to maximize throughput and overlap.
+/// While storage workers are CPU-bound, account workers are I/O-bound coordinators.
+fn default_account_worker_count() -> usize {
+    ((default_storage_worker_count() * 3) / 2).max(MIN_WORKER_COUNT)
+}
+
 /// The size of proof targets chunk to spawn in one multiproof calculation.
 pub const DEFAULT_MULTIPROOF_TASK_CHUNK_SIZE: usize = 10;
 
@@ -109,6 +136,10 @@ 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,
+    /// Number of account proof worker threads.
+    account_worker_count: usize,
 }
 
 impl Default for TreeConfig {
@@ -135,6 +166,8 @@ 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(),
+            account_worker_count: default_account_worker_count(),
         }
     }
 }
@@ -164,7 +197,10 @@ impl TreeConfig {
         always_process_payload_attributes_on_canonical_head: bool,
         prewarm_max_concurrency: usize,
         allow_unwind_canonical_header: bool,
+        storage_worker_count: usize,
+        account_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 +223,8 @@ impl TreeConfig {
             always_process_payload_attributes_on_canonical_head,
             prewarm_max_concurrency,
             allow_unwind_canonical_header,
+            storage_worker_count,
+            account_worker_count,
         }
     }
 
@@ -394,6 +432,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 +491,26 @@ 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.max(MIN_WORKER_COUNT);
+        self
+    }
+
+    /// Return the number of account proof worker threads.
+    pub const fn account_worker_count(&self) -> usize {
+        self.account_worker_count
+    }
+
+    /// Setter for the number of account proof worker threads.
+    pub const fn with_account_worker_count(mut self, account_worker_count: usize) -> Self {
+        self.account_worker_count = account_worker_count.max(MIN_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

@@ -32,7 +32,7 @@ use reth_provider::{
 use reth_revm::{db::BundleState, state::EvmState};
 use reth_trie::TrieInput;
 use reth_trie_parallel::{
-    proof_task::{ProofTaskCtx, ProofTaskManager},
+    proof_task::{spawn_proof_workers, ProofTaskCtx},
     root::ParallelStateRootError,
 };
 use reth_trie_sparse::{
@@ -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,9 @@ 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 proof worker spawning fails.
+    #[allow(clippy::type_complexity)]
     pub fn spawn<P, I: ExecutableTxIterator<Evm>>(
         &mut self,
         env: ExecutionEnv<Evm>,
@@ -174,7 +177,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
@@ -185,8 +191,7 @@ where
     {
         let (to_sparse_trie, sparse_trie_rx) = channel();
         // spawn multiproof task, save the trie input
-        let (trie_input, state_root_config) =
-            MultiProofConfig::new_from_input(consistent_view, trie_input);
+        let (trie_input, state_root_config) = MultiProofConfig::from_input(trie_input);
         self.trie_input = Some(trie_input);
 
         // Create and spawn the storage proof task
@@ -195,21 +200,29 @@ where
             state_root_config.state_sorted.clone(),
             state_root_config.prefix_sets.clone(),
         );
+        let storage_worker_count = config.storage_worker_count();
+        let account_worker_count = config.account_worker_count();
         let max_proof_task_concurrency = config.max_proof_task_concurrency() as usize;
-        let proof_task = ProofTaskManager::new(
+        let proof_handle = match spawn_proof_workers(
             self.executor.handle().clone(),
-            state_root_config.consistent_view.clone(),
+            consistent_view,
             task_ctx,
-            max_proof_task_concurrency,
-        );
+            storage_worker_count,
+            account_worker_count,
+        ) {
+            Ok(handle) => handle,
+            Err(error) => {
+                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.
         let max_multi_proof_task_concurrency = max_proof_task_concurrency / 2;
         let multi_proof_task = MultiProofTask::new(
             state_root_config,
             self.executor.clone(),
-            proof_task.handle(),
+            proof_handle.clone(),
             to_sparse_trie,
             max_multi_proof_task_concurrency,
             config.multiproof_chunking_enabled().then_some(config.multiproof_chunk_size()),
@@ -238,26 +251,14 @@ where
         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_task.handle(), state_root_tx);
+        self.spawn_sparse_trie_task(sparse_trie_rx, proof_handle, state_root_tx);
 
-        // spawn the proof task
-        self.executor.spawn_blocking(move || {
-            if let Err(err) = proof_task.run() {
-                // At least log if there is an error at any point
-                tracing::error!(
-                    target: "engine::root",
-                    ?err,
-                    "Storage proof task returned an error"
-                );
-            }
-        });
-
-        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.
@@ -857,14 +858,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();