Cleanup

crates/engine/tree/benches/state_root_task.rs

@@ -24,9 +24,11 @@ use reth_provider::{
     test_utils::{create_test_provider_factory_with_chain_spec, MockNodeTypesWithDB},
     AccountReader, ChainSpecProvider, HashingWriter, ProviderFactory,
 };
+use reth_trie_parallel::proof_task::ProofWorkerHandle;
 use revm_primitives::{HashMap, U256};
 use revm_state::{Account as RevmAccount, AccountInfo, AccountStatus, EvmState, EvmStorageSlot};
 use std::{hint::black_box, sync::Arc};
+use tokio::runtime::Runtime;
 
 #[derive(Debug, Clone)]
 struct BenchParams {
@@ -215,8 +217,12 @@ fn bench_state_root(c: &mut Criterion) {
                         let state_updates = create_bench_state_updates(params);
                         setup_provider(&factory, &state_updates).expect("failed to setup provider");
 
+                        let rt = Runtime::new().unwrap();
+                        let proof_worker_handle = ProofWorkerHandle::new(rt.handle().clone(), 1, 1);
+
                         let payload_processor = PayloadProcessor::new(
                             WorkloadExecutor::default(),
+                            proof_worker_handle,
                             EthEvmConfig::new(factory.chain_spec()),
                             &TreeConfig::default(),
                             PrecompileCacheMap::default(),

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

@@ -123,8 +123,7 @@ where
 impl<N, Evm, ProofProviderFactory> PayloadProcessor<Evm, ProofProviderFactory>
 where
     N: NodePrimitives,
-    Evm: ConfigureEvm<Primitives = N> + 'static,
-    ProofProviderFactory: Clone + Send,
+    Evm: ConfigureEvm<Primitives = N>,
 {
     /// Creates a new payload processor.
     pub fn new(
@@ -149,7 +148,13 @@ where
             prewarm_max_concurrency: config.prewarm_max_concurrency(),
         }
     }
+}
 
+impl<N, Evm, ProofProviderFactory> PayloadProcessor<Evm, ProofProviderFactory>
+where
+    N: NodePrimitives,
+    Evm: ConfigureEvm<Primitives = N> + 'static,
+{
     /// Spawns all background tasks and returns a handle connected to the tasks.
     ///
     /// - Transaction prewarming task
@@ -181,7 +186,7 @@ where
     ///
     ///
     /// This returns a handle to await the final state root and to interact with the tasks (e.g.
-    /// canceling).
+    /// canceling)
     #[allow(clippy::type_complexity)]
     #[instrument(
         level = "debug",
@@ -199,6 +204,7 @@ where
     ) -> PayloadHandle<WithTxEnv<TxEnvFor<Evm>, I::Tx>, I::Error>
     where
         P: BlockReader + StateProviderFactory + StateReader + Clone + 'static,
+        ProofProviderFactory: Clone + Send,
     {
         let span = tracing::Span::current();
         let (to_sparse_trie, sparse_trie_rx) = channel();
@@ -684,9 +690,11 @@ mod tests {
     };
     use reth_testing_utils::generators;
     use reth_trie::{test_utils::state_root, HashedPostState};
+    use reth_trie_parallel::proof_task::ProofWorkerHandle;
     use revm_primitives::{Address, HashMap, B256, KECCAK_EMPTY, U256};
     use revm_state::{AccountInfo, AccountStatus, EvmState, EvmStorageSlot};
     use std::sync::Arc;
+    use tokio::runtime::Runtime;
 
     fn make_saved_cache(hash: B256) -> SavedCache {
         let execution_cache = ExecutionCacheBuilder::default().build_caches(1_000);
@@ -859,8 +867,13 @@ mod tests {
             }
         }
 
+        // Create proof worker handle for the test
+        let rt = Runtime::new().unwrap();
+        let proof_worker_handle = ProofWorkerHandle::new(rt.handle().clone(), 1, 1);
+
         let mut payload_processor = PayloadProcessor::new(
             WorkloadExecutor::default(),
+            proof_worker_handle,
             EthEvmConfig::new(factory.chain_spec()),
             &TreeConfig::default(),
             PrecompileCacheMap::default(),

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

@@ -1287,7 +1287,8 @@ mod tests {
     use alloy_primitives::map::B256Set;
     use reth_provider::{
         providers::OverlayStateProviderFactory, test_utils::create_test_provider_factory,
-        BlockReader, PruneCheckpointReader, StageCheckpointReader, TrieReader,
+        BlockReader, DatabaseProviderFactory, PruneCheckpointReader, StageCheckpointReader,
+        TrieReader,
     };
     use reth_trie::MultiProof;
     use reth_trie_parallel::proof_task::ProofWorkerHandle;
@@ -1306,11 +1307,19 @@ mod tests {
             .clone()
     }
 
-    fn create_test_state_root_task<F>(factory: F) -> MultiProofTask {
+    fn create_test_state_root_task<F>(factory: F) -> MultiProofTask
+    where
+        F: DatabaseProviderFactory<
+                Provider: BlockReader + TrieReader + StageCheckpointReader + PruneCheckpointReader,
+            > + Clone
+            + Send
+            + 'static,
+    {
         let rt_handle = get_test_runtime_handle();
         let overlay_factory = OverlayStateProviderFactory::new(factory);
-        let proof_handle = ProofWorkerHandle::new(rt_handle, 1, 1);
-        let proof_dispatcher = proof_handle.into_dispatcher(overlay_factory);
+        let proof_handle: ProofWorkerHandle<OverlayStateProviderFactory<F>> =
+            ProofWorkerHandle::new(rt_handle, 1, 1);
+        let proof_dispatcher = proof_handle.new_dispatcher(overlay_factory);
         let (to_sparse_trie, _receiver) = std::sync::mpsc::channel();
 
         MultiProofTask::new(proof_dispatcher, to_sparse_trie, Some(1))

crates/trie/parallel/src/proof.rs

@@ -329,8 +329,8 @@ mod tests {
         let rt = Runtime::new().unwrap();
 
         let factory = reth_provider::providers::OverlayStateProviderFactory::new(factory);
-        let proof_worker_dispatcher =
-            ProofWorkerHandle::new(rt.handle().clone(), 1, 1).into_dispatcher(factory);
+        let proof_worker_handle = ProofWorkerHandle::new(rt.handle().clone(), 1, 1);
+        let proof_worker_dispatcher = proof_worker_handle.new_dispatcher(factory);
 
         let parallel_result = ParallelProof::new(
             Default::default(),