refactor(trie): remove proof task manager (#18934)

Co-authored-by: Brian Picciano <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Alexey Shekhirin <[email protected]>

Author: 4 people

crates/engine/primitives/src/config.rs

@@ -9,11 +9,14 @@ pub const DEFAULT_MEMORY_BLOCK_BUFFER_TARGET: u64 = 0;
 /// 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 = 32;
+
 /// 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(2, 64)).unwrap_or(8)
+        std::thread::available_parallelism().map_or(8, |n| n.get() * 2).min(MIN_WORKER_COUNT)
     }
     #[cfg(not(feature = "std"))]
     {
@@ -491,8 +494,8 @@ impl TreeConfig {
     }
 
     /// 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;
+    pub fn with_storage_worker_count(mut self, storage_worker_count: usize) -> Self {
+        self.storage_worker_count = storage_worker_count.max(MIN_WORKER_COUNT);
         self
     }
 
@@ -502,8 +505,8 @@ impl TreeConfig {
     }
 
     /// 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;
+    pub 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/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::{ProofTaskCtx, ProofWorkerHandle},
     root::ParallelStateRootError,
 };
 use reth_trie_sparse::{
@@ -167,8 +167,7 @@ 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.
+    /// 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,
@@ -204,14 +203,14 @@ where
         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 = match ProofTaskManager::new(
+        let proof_handle = match ProofWorkerHandle::new(
             self.executor.handle().clone(),
             consistent_view,
             task_ctx,
             storage_worker_count,
             account_worker_count,
         ) {
-            Ok(task) => task,
+            Ok(handle) => handle,
             Err(error) => {
                 return Err((error, transactions, env, provider_builder));
             }
@@ -223,7 +222,7 @@ where
         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()),
@@ -252,19 +251,7 @@ 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);
-
-        // 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"
-                );
-            }
-        });
+        self.spawn_sparse_trie_task(sparse_trie_rx, proof_handle, state_root_tx);
 
         Ok(PayloadHandle {
             to_multi_proof,
@@ -406,7 +393,7 @@ where
     fn spawn_sparse_trie_task<BPF>(
         &self,
         sparse_trie_rx: mpsc::Receiver<SparseTrieUpdate>,
-        proof_task_handle: BPF,
+        proof_worker_handle: BPF,
         state_root_tx: mpsc::Sender<Result<StateRootComputeOutcome, ParallelStateRootError>>,
     ) where
         BPF: TrieNodeProviderFactory + Clone + Send + Sync + 'static,
@@ -436,7 +423,7 @@ where
         let task =
             SparseTrieTask::<_, ConfiguredSparseTrie, ConfiguredSparseTrie>::new_with_cleared_trie(
                 sparse_trie_rx,
-                proof_task_handle,
+                proof_worker_handle,
                 self.trie_metrics.clone(),
                 sparse_state_trie,
             );

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

@@ -20,7 +20,7 @@ use reth_trie::{
 };
 use reth_trie_parallel::{
     proof::ParallelProof,
-    proof_task::{AccountMultiproofInput, ProofTaskKind, ProofTaskManagerHandle},
+    proof_task::{AccountMultiproofInput, ProofWorkerHandle},
     root::ParallelStateRootError,
 };
 use std::{
@@ -346,11 +346,8 @@ pub struct MultiproofManager {
     pending: VecDeque<PendingMultiproofTask>,
     /// Executor for tasks
     executor: WorkloadExecutor,
-    /// Handle to the proof task manager used for creating `ParallelProof` instances for storage
-    /// proofs.
-    storage_proof_task_handle: ProofTaskManagerHandle,
-    /// Handle to the proof task manager used for account multiproofs.
-    account_proof_task_handle: ProofTaskManagerHandle,
+    /// Handle to the proof worker pools (storage and account).
+    proof_worker_handle: ProofWorkerHandle,
     /// Cached storage proof roots for missed leaves; this maps
     /// hashed (missed) addresses to their storage proof roots.
     ///
@@ -372,8 +369,7 @@ impl MultiproofManager {
     fn new(
         executor: WorkloadExecutor,
         metrics: MultiProofTaskMetrics,
-        storage_proof_task_handle: ProofTaskManagerHandle,
-        account_proof_task_handle: ProofTaskManagerHandle,
+        proof_worker_handle: ProofWorkerHandle,
         max_concurrent: usize,
     ) -> Self {
         Self {
@@ -382,8 +378,7 @@ impl MultiproofManager {
             executor,
             inflight: 0,
             metrics,
-            storage_proof_task_handle,
-            account_proof_task_handle,
+            proof_worker_handle,
             missed_leaves_storage_roots: Default::default(),
         }
     }
@@ -452,7 +447,7 @@ impl MultiproofManager {
             multi_added_removed_keys,
         } = storage_multiproof_input;
 
-        let storage_proof_task_handle = self.storage_proof_task_handle.clone();
+        let storage_proof_worker_handle = self.proof_worker_handle.clone();
         let missed_leaves_storage_roots = self.missed_leaves_storage_roots.clone();
 
         self.executor.spawn_blocking(move || {
@@ -471,7 +466,7 @@ impl MultiproofManager {
                 config.state_sorted,
                 config.prefix_sets,
                 missed_leaves_storage_roots,
-                storage_proof_task_handle,
+                storage_proof_worker_handle,
             )
             .with_branch_node_masks(true)
             .with_multi_added_removed_keys(Some(multi_added_removed_keys))
@@ -524,7 +519,7 @@ impl MultiproofManager {
             state_root_message_sender,
             multi_added_removed_keys,
         } = multiproof_input;
-        let account_proof_task_handle = self.account_proof_task_handle.clone();
+        let account_proof_worker_handle = self.proof_worker_handle.clone();
         let missed_leaves_storage_roots = self.missed_leaves_storage_roots.clone();
 
         self.executor.spawn_blocking(move || {
@@ -556,15 +551,10 @@ impl MultiproofManager {
                 missed_leaves_storage_roots,
             };
 
-            let (sender, receiver) = channel();
             let proof_result: Result<DecodedMultiProof, ParallelStateRootError> = (|| {
-                account_proof_task_handle
-                    .queue_task(ProofTaskKind::AccountMultiproof(input, sender))
-                    .map_err(|_| {
-                        ParallelStateRootError::Other(
-                            "Failed to queue account multiproof to worker pool".into(),
-                        )
-                    })?;
+                let receiver = account_proof_worker_handle
+                    .queue_account_multiproof(input)
+                    .map_err(|e| ParallelStateRootError::Other(e.to_string()))?;
 
                 receiver
                     .recv()
@@ -693,7 +683,7 @@ impl MultiProofTask {
     pub(super) fn new(
         config: MultiProofConfig,
         executor: WorkloadExecutor,
-        proof_task_handle: ProofTaskManagerHandle,
+        proof_worker_handle: ProofWorkerHandle,
         to_sparse_trie: Sender<SparseTrieUpdate>,
         max_concurrency: usize,
         chunk_size: Option<usize>,
@@ -713,8 +703,7 @@ impl MultiProofTask {
             multiproof_manager: MultiproofManager::new(
                 executor,
                 metrics.clone(),
-                proof_task_handle.clone(), // handle for storage proof workers
-                proof_task_handle,         // handle for account proof workers
+                proof_worker_handle,
                 max_concurrency,
             ),
             metrics,
@@ -1223,7 +1212,7 @@ mod tests {
         DatabaseProviderFactory,
     };
     use reth_trie::{MultiProof, TrieInput};
-    use reth_trie_parallel::proof_task::{ProofTaskCtx, ProofTaskManager};
+    use reth_trie_parallel::proof_task::{ProofTaskCtx, ProofWorkerHandle};
     use revm_primitives::{B256, U256};
 
     fn create_test_state_root_task<F>(factory: F) -> MultiProofTask
@@ -1238,12 +1227,12 @@ mod tests {
             config.prefix_sets.clone(),
         );
         let consistent_view = ConsistentDbView::new(factory, None);
-        let proof_task =
-            ProofTaskManager::new(executor.handle().clone(), consistent_view, task_ctx, 1, 1)
-                .expect("Failed to create ProofTaskManager");
+        let proof_handle =
+            ProofWorkerHandle::new(executor.handle().clone(), consistent_view, task_ctx, 1, 1)
+                .expect("Failed to spawn proof workers");
         let channel = channel();
 
-        MultiProofTask::new(config, executor, proof_task.handle(), channel.0, 1, None)
+        MultiProofTask::new(config, executor, proof_handle, channel.0, 1, None)
     }
 
     #[test]

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

@@ -892,13 +892,12 @@ where
                             (handle, StateRootStrategy::StateRootTask)
                         }
                         Err((error, txs, env, provider_builder)) => {
-                            // Failed to initialize proof task manager, fallback to parallel state
-                            // root
+                            // Failed to spawn proof workers, 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"
+                                "Failed to spawn proof workers, falling back to parallel state root"
                             );
                             (
                                 self.payload_processor.spawn_cache_exclusive(