refactor(engine): streamline proof task management with new handle

- Replaced `ProofTaskManager` with a new `new_proof_task_handle` function for improved task management.
- Updated related code to utilize the new proof task handle, enhancing clarity and reducing complexity.
- Adjusted tests to reflect changes in proof task initialization and handling.

Author: yongkangc

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::{new_proof_task_handle, ProofTaskCtx},
     root::ParallelStateRootError,
 };
 use reth_trie_sparse::{
@@ -196,23 +196,21 @@ where
             state_root_config.prefix_sets.clone(),
         );
         let max_proof_task_concurrency = config.max_proof_task_concurrency() as usize;
-        let storage_worker_count = config.storage_proof_workers();
-        let proof_task = ProofTaskManager::new(
+        let proof_task_handle = new_proof_task_handle(
             self.executor.handle().clone(),
             state_root_config.consistent_view.clone(),
             task_ctx,
             max_proof_task_concurrency,
-            storage_worker_count,
         )
-        .expect("Failed to create ProofTaskManager with storage workers");
+        .expect("Failed to create proof task handle");
 
         // 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_task_handle.clone(),
             to_sparse_trie,
             max_multi_proof_task_concurrency,
             config.multiproof_chunking_enabled().then_some(config.multiproof_chunk_size()),
@@ -241,19 +239,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_task_handle, state_root_tx);
 
         PayloadHandle {
             to_multi_proof,

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

@@ -1204,7 +1204,7 @@ mod tests {
     use alloy_primitives::map::B256Set;
     use reth_provider::{providers::ConsistentDbView, test_utils::create_test_provider_factory};
     use reth_trie::{MultiProof, TrieInput};
-    use reth_trie_parallel::proof_task::{ProofTaskCtx, ProofTaskManager};
+    use reth_trie_parallel::proof_task::{new_proof_task_handle, ProofTaskCtx};
     use revm_primitives::{B256, U256};
     use std::sync::Arc;
 
@@ -1231,17 +1231,16 @@ mod tests {
             config.state_sorted.clone(),
             config.prefix_sets.clone(),
         );
-        let proof_task = ProofTaskManager::new(
+        let proof_task_handle = new_proof_task_handle(
             executor.handle().clone(),
             config.consistent_view.clone(),
             task_ctx,
-            1,
-            1, // storage_worker_count for test
+            1, // max_concurrency for test
         )
-        .expect("Failed to create ProofTaskManager for multiproof test");
+        .expect("Failed to create proof task handle for multiproof test");
         let channel = channel();
 
-        MultiProofTask::new(config, executor, proof_task.handle(), channel.0, 1, None)
+        MultiProofTask::new(config, executor, proof_task_handle, channel.0, 1, None)
     }
 
     #[test]

crates/trie/parallel/src/proof.rs

@@ -14,8 +14,7 @@ use dashmap::DashMap;
 use itertools::Itertools;
 use reth_execution_errors::StorageRootError;
 use reth_provider::{
-    providers::ConsistentDbView, BlockReader, DBProvider, DatabaseProviderFactory, FactoryTx,
-    ProviderError,
+    providers::ConsistentDbView, BlockReader, DBProvider, DatabaseProviderFactory, ProviderError,
 };
 use reth_storage_errors::db::DatabaseError;
 use reth_trie::{
@@ -34,7 +33,7 @@ use reth_trie_common::{
     proof::{DecodedProofNodes, ProofRetainer},
 };
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
-use std::sync::{mpsc::Receiver, Arc};
+use std::sync::Arc;
 use tracing::trace;
 
 /// Parallel proof calculator.
@@ -59,7 +58,7 @@ pub struct ParallelProof<Factory: DatabaseProviderFactory> {
     /// Provided by the user to give the necessary context to retain extra proofs.
     multi_added_removed_keys: Option<Arc<MultiAddedRemovedKeys>>,
     /// Handle to the storage proof task.
-    storage_proof_task_handle: ProofTaskManagerHandle<FactoryTx<Factory>>,
+    storage_proof_task_handle: ProofTaskManagerHandle<<Factory::Provider as DBProvider>::Tx>,
     /// Cached storage proof roots for missed leaves; this maps
     /// hashed (missed) addresses to their storage proof roots.
     missed_leaves_storage_roots: Arc<DashMap<B256, B256>>,
@@ -75,7 +74,7 @@ impl<Factory: DatabaseProviderFactory> ParallelProof<Factory> {
         state_sorted: Arc<HashedPostStateSorted>,
         prefix_sets: Arc<TriePrefixSetsMut>,
         missed_leaves_storage_roots: Arc<DashMap<B256, B256>>,
-        storage_proof_task_handle: ProofTaskManagerHandle<FactoryTx<Factory>>,
+        storage_proof_task_handle: ProofTaskManagerHandle<<Factory::Provider as DBProvider>::Tx>,
     ) -> Self {
         Self {
             view,
@@ -118,18 +117,17 @@ where
         hashed_address: B256,
         prefix_set: PrefixSet,
         target_slots: B256Set,
-    ) -> Receiver<Result<DecodedStorageMultiProof, ParallelStateRootError>> {
+    ) -> crossbeam_channel::Receiver<Result<DecodedStorageMultiProof, ParallelStateRootError>> {
         let input = StorageProofInput::new(
             hashed_address,
             prefix_set,
-            Arc::new(target_slots),
+            target_slots,
             self.collect_branch_node_masks,
             self.multi_added_removed_keys.clone(),
         );
 
-        let (sender, receiver) = std::sync::mpsc::channel();
-        let _ =
-            self.storage_proof_task_handle.queue_task(ProofTaskKind::StorageProof(input, sender));
+        let (sender, receiver) = crossbeam_channel::unbounded();
+        self.storage_proof_task_handle.queue_task(ProofTaskKind::StorageProof(input, sender));
         receiver
     }
 
@@ -368,7 +366,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::proof_task::{ProofTaskCtx, ProofTaskManager};
+    use crate::proof_task::{new_proof_task_handle, ProofTaskCtx};
     use alloy_primitives::{
         keccak256,
         map::{B256Set, DefaultHashBuilder},
@@ -447,19 +445,13 @@ mod tests {
 
         let task_ctx =
             ProofTaskCtx::new(Default::default(), Default::default(), Default::default());
-        let proof_task = ProofTaskManager::new(
+        let proof_task_handle = new_proof_task_handle(
             rt.handle().clone(),
             consistent_view.clone(),
             task_ctx,
-            1,
-            1, // storage_worker_count for test
+            1, // max_concurrency for test
         )
         .expect("Failed to create proof task");
-        let proof_task_handle = proof_task.handle();
-
-        // keep the join handle around to make sure it does not return any errors
-        // after we compute the state root
-        let join_handle = rt.spawn_blocking(move || proof_task.run());
 
         let parallel_result = ParallelProof::new(
             consistent_view,
@@ -495,9 +487,8 @@ mod tests {
         // then compare the entire thing for any mask differences
         assert_eq!(parallel_result, sequential_result_decoded);
 
-        // drop the handle to terminate the task and then block on the proof task handle to make
-        // sure it does not return any errors
+        // Drop the handle to release transaction pool resources
+        // Note: No manager loop to join in the new design - handle manages lifecycle via Drop
         drop(proof_task_handle);
-        rt.block_on(join_handle).unwrap().expect("The proof task should not return an error");
     }
 }