fix(engine): do not create another cache for multiproof task

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

@@ -1,6 +1,5 @@
 //! BAL (Block Access List, EIP-7928) related functionality.
 
-use crate::tree::cached_state::CachedStateProvider;
 use alloy_consensus::constants::KECCAK_EMPTY;
 use alloy_eip7928::BlockAccessList;
 use alloy_primitives::{keccak256, Address, StorageKey, U256};
@@ -118,7 +117,7 @@ impl<'a> Iterator for BALSlotIter<'a> {
 /// of modified accounts and storage slots.
 pub(crate) fn bal_to_hashed_post_state<P>(
     bal: &BlockAccessList,
-    provider: &CachedStateProvider<P>,
+    provider: P,
 ) -> Result<HashedPostState, ProviderError>
 where
     P: AccountReader,
@@ -198,17 +197,12 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::tree::cached_state::{ExecutionCache, ExecutionCacheBuilder};
     use alloy_eip7928::{
         AccountChanges, BalanceChange, CodeChange, NonceChange, SlotChanges, StorageChange,
     };
     use alloy_primitives::{Address, Bytes, StorageKey, B256};
     use reth_revm::test_utils::StateProviderTest;
 
-    fn new_cache() -> ExecutionCache {
-        ExecutionCacheBuilder::default().build_caches(1000)
-    }
-
     #[test]
     fn test_bal_to_hashed_post_state_basic() {
         let provider = StateProviderTest::default();
@@ -224,7 +218,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         assert_eq!(result.accounts.len(), 1);
@@ -259,7 +252,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);
@@ -289,7 +281,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);
@@ -317,7 +308,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);
@@ -352,7 +342,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);
@@ -385,7 +374,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);
@@ -426,7 +414,6 @@ mod tests {
         };
 
         let bal = vec![account_changes];
-        let provider = CachedStateProvider::new(provider, new_cache(), Default::default());
         let result = bal_to_hashed_post_state(&bal, &provider).unwrap();
 
         let hashed_address = keccak256(address);

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

@@ -30,7 +30,9 @@ use reth_evm::{
 };
 use reth_execution_types::ExecutionOutcome;
 use reth_primitives_traits::NodePrimitives;
-use reth_provider::{BlockReader, DatabaseProviderROFactory, StateProviderFactory, StateReader};
+use reth_provider::{
+    BlockReader, DatabaseProviderROFactory, StateProvider, StateProviderFactory, StateReader,
+};
 use reth_revm::{db::BundleState, state::EvmState};
 use reth_trie::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory};
 use reth_trie_parallel::{
@@ -44,6 +46,7 @@ use reth_trie_sparse::{
 use reth_trie_sparse_parallel::{ParallelSparseTrie, ParallelismThresholds};
 use std::{
     collections::BTreeMap,
+    ops::Not,
     sync::{
         atomic::AtomicBool,
         mpsc::{self, channel},
@@ -230,8 +233,6 @@ where
             + Send
             + 'static,
     {
-        let parent_hash = env.parent_hash;
-
         // start preparing transactions immediately
         let (prewarm_rx, execution_rx, transaction_count_hint) =
             self.spawn_tx_iterator(transactions);
@@ -293,17 +294,19 @@ where
 
         // spawn multi-proof task
         let parent_span = span.clone();
-        self.executor.spawn_blocking({
-            let saved_cache = self.cache_for(parent_hash);
-            let cache = saved_cache.cache().clone();
-            let cache_metrics = saved_cache.metrics().clone();
-            move || {
-                let _enter = parent_span.entered();
-                // Build a state provider for the multiproof task
-                let provider = provider_builder.build().expect("failed to build provider");
-                let provider = CachedStateProvider::new(provider, cache, cache_metrics);
-                multi_proof_task.run(provider);
-            }
+        let saved_cache = prewarm_handle.saved_cache.clone();
+        self.executor.spawn_blocking(move || {
+            let _enter = parent_span.entered();
+            // Build a state provider for the multiproof task
+            let provider = provider_builder.build().expect("failed to build provider");
+            let provider = if let Some(saved_cache) = saved_cache {
+                let (cache, metrics) = saved_cache.split();
+                Box::new(CachedStateProvider::new(provider, cache, metrics))
+                    as Box<dyn StateProvider>
+            } else {
+                Box::new(provider)
+            };
+            multi_proof_task.run(provider);
         });
 
         // wire the sparse trie to the state root response receiver
@@ -422,20 +425,13 @@ where
             transactions = mpsc::channel().1;
         }
 
-        let (saved_cache, cache, cache_metrics) = if self.disable_state_cache {
-            (None, None, None)
-        } else {
-            let saved_cache = self.cache_for(env.parent_hash);
-            let cache = saved_cache.cache().clone();
-            let cache_metrics = saved_cache.metrics().clone();
-            (Some(saved_cache), Some(cache), Some(cache_metrics))
-        };
+        let saved_cache = self.disable_state_cache.not().then(|| self.cache_for(env.parent_hash));
 
         // configure prewarming
         let prewarm_ctx = PrewarmContext {
             env,
             evm_config: self.evm_config.clone(),
-            saved_cache,
+            saved_cache: saved_cache.clone(),
             provider: provider_builder,
             metrics: PrewarmMetrics::default(),
             terminate_execution: Arc::new(AtomicBool::new(false)),
@@ -465,7 +461,7 @@ where
             });
         }
 
-        CacheTaskHandle { cache, to_prewarm_task: Some(to_prewarm_task), cache_metrics }
+        CacheTaskHandle { saved_cache, to_prewarm_task: Some(to_prewarm_task) }
     }
 
     /// Returns the cache for the given parent hash.
@@ -651,12 +647,12 @@ impl<Tx, Err, R: Send + Sync + 'static> PayloadHandle<Tx, Err, R> {
 
     /// Returns a clone of the caches used by prewarming
     pub(super) fn caches(&self) -> Option<StateExecutionCache> {
-        self.prewarm_handle.cache.clone()
+        self.prewarm_handle.saved_cache.as_ref().map(|cache| cache.cache().clone())
     }
 
     /// Returns a clone of the cache metrics used by prewarming
     pub(super) fn cache_metrics(&self) -> Option<CachedStateMetrics> {
-        self.prewarm_handle.cache_metrics.clone()
+        self.prewarm_handle.saved_cache.as_ref().map(|cache| cache.metrics().clone())
     }
 
     /// Terminates the pre-warming transaction processing.
@@ -693,9 +689,7 @@ impl<Tx, Err, R: Send + Sync + 'static> PayloadHandle<Tx, Err, R> {
 #[derive(Debug)]
 pub(crate) struct CacheTaskHandle<R> {
     /// The shared cache the task operates with.
-    cache: Option<StateExecutionCache>,
-    /// Metrics for the caches
-    cache_metrics: Option<CachedStateMetrics>,
+    saved_cache: Option<SavedCache>,
     /// Channel to the spawned prewarm task if any
     to_prewarm_task: Option<std::sync::mpsc::Sender<PrewarmTaskEvent<R>>>,
 }

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

@@ -1,8 +1,6 @@
 //! Multiproof task related functionality.
 
-use crate::tree::{
-    cached_state::CachedStateProvider, payload_processor::bal::bal_to_hashed_post_state,
-};
+use crate::tree::payload_processor::bal::bal_to_hashed_post_state;
 use alloy_eip7928::BlockAccessList;
 use alloy_evm::block::StateChangeSource;
 use alloy_primitives::{keccak256, map::HashSet, B256};
@@ -879,7 +877,7 @@ impl MultiProofTask {
         msg: MultiProofMessage,
         ctx: &mut MultiproofBatchCtx,
         batch_metrics: &mut MultiproofBatchMetrics,
-        provider: &CachedStateProvider<P>,
+        provider: &P,
     ) -> bool
     where
         P: AccountReader,
@@ -1184,7 +1182,7 @@ impl MultiProofTask {
         target = "engine::tree::payload_processor::multiproof",
         skip_all
     )]
-    pub(crate) fn run<P>(mut self, provider: CachedStateProvider<P>)
+    pub(crate) fn run<P>(mut self, provider: P)
     where
         P: AccountReader,
     {
@@ -1501,7 +1499,7 @@ fn estimate_evm_state_targets(state: &EvmState) -> usize {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::tree::cached_state::ExecutionCacheBuilder;
+    use crate::tree::cached_state::{CachedStateProvider, ExecutionCacheBuilder};
     use alloy_eip7928::{AccountChanges, BalanceChange};
     use alloy_primitives::{map::B256Set, Address};
     use reth_provider::{