diff --git a/crates/engine/tree/benches/state_root_task.rs b/crates/engine/tree/benches/state_root_task.rs index 71cac68543..c049763e8d 100644 --- a/crates/engine/tree/benches/state_root_task.rs +++ b/crates/engine/tree/benches/state_root_task.rs @@ -13,7 +13,7 @@ use reth_chain_state::EthPrimitives; use reth_chainspec::ChainSpec; use reth_db_common::init::init_genesis; use reth_engine_tree::tree::{ - executor::WorkloadExecutor, precompile_cache::PrecompileCache, PayloadProcessor, + executor::WorkloadExecutor, precompile_cache::PrecompileCacheMap, PayloadProcessor, StateProviderBuilder, TreeConfig, }; use reth_evm::OnStateHook; @@ -221,7 +221,7 @@ fn bench_state_root(c: &mut Criterion) { WorkloadExecutor::default(), EthEvmConfig::new(factory.chain_spec()), &TreeConfig::default(), - PrecompileCache::default(), + PrecompileCacheMap::default(), ); let provider = BlockchainProvider::new(factory).unwrap(); diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index c13d1aeebb..e298640637 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/crates/engine/tree/src/tree/mod.rs @@ -18,7 +18,7 @@ use error::{InsertBlockError, InsertBlockErrorKind, InsertBlockFatalError}; use instrumented_state::InstrumentedStateProvider; use payload_processor::sparse_trie::StateRootComputeOutcome; use persistence_state::CurrentPersistenceAction; -use precompile_cache::{CachedPrecompile, PrecompileCache}; +use precompile_cache::{CachedPrecompile, PrecompileCacheMap}; use reth_chain_state::{ CanonicalInMemoryState, ExecutedBlock, ExecutedBlockWithTrieUpdates, MemoryOverlayStateProvider, NewCanonicalChain, @@ -268,8 +268,8 @@ where payload_processor: PayloadProcessor, /// The EVM configuration. evm_config: C, - /// Precompile cache. - precompile_cache: PrecompileCache, + /// Precompile cache map. + precompile_cache_map: PrecompileCacheMap, } impl std::fmt::Debug @@ -334,13 +334,13 @@ where ) -> Self { let (incoming_tx, incoming) = std::sync::mpsc::channel(); - let precompile_cache = PrecompileCache::default(); + let precompile_cache_map = PrecompileCacheMap::default(); let payload_processor = PayloadProcessor::new( WorkloadExecutor::default(), evm_config.clone(), &config, - precompile_cache.clone(), + precompile_cache_map.clone(), ); Self { @@ -362,7 +362,7 @@ where engine_kind, payload_processor, evm_config, - precompile_cache, + precompile_cache_map, } } @@ -2275,7 +2275,7 @@ where /// Executes a block with the given state provider fn execute_block( - &self, + &mut self, state_provider: S, block: &RecoveredBlock, handle: &PayloadHandle, @@ -2289,8 +2289,11 @@ where let mut executor = self.evm_config.executor_for_block(&mut db, block); if self.config.precompile_cache_enabled() { - executor.evm_mut().precompiles_mut().map_precompiles(|_, precompile| { - CachedPrecompile::wrap(precompile, self.precompile_cache.clone()) + executor.evm_mut().precompiles_mut().map_precompiles(|address, precompile| { + CachedPrecompile::wrap( + precompile, + self.precompile_cache_map.entry(*address).or_default().clone(), + ) }); } diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs index 578c77b5bd..f17858a190 100644 --- a/crates/engine/tree/src/tree/payload_processor/mod.rs +++ b/crates/engine/tree/src/tree/payload_processor/mod.rs @@ -6,7 +6,6 @@ use crate::tree::{ prewarm::{PrewarmCacheTask, PrewarmContext, PrewarmTaskEvent}, sparse_trie::StateRootComputeOutcome, }, - precompile_cache::PrecompileCache, sparse_trie::SparseTrieTask, StateProviderBuilder, TreeConfig, }; @@ -39,6 +38,8 @@ use std::{ }, }; +use super::precompile_cache::PrecompileCacheMap; + pub mod executor; pub mod multiproof; pub mod prewarm; @@ -61,8 +62,8 @@ pub struct PayloadProcessor { evm_config: Evm, /// whether precompile cache should be enabled. precompile_cache_enabled: bool, - /// Precompile cache. - precompile_cache: PrecompileCache, + /// Precompile cache map. + precompile_cache_map: PrecompileCacheMap, _marker: std::marker::PhantomData, } @@ -72,7 +73,7 @@ impl PayloadProcessor { executor: WorkloadExecutor, evm_config: Evm, config: &TreeConfig, - precompile_cache: PrecompileCache, + precompile_cache_map: PrecompileCacheMap, ) -> Self { Self { executor, @@ -82,7 +83,7 @@ impl PayloadProcessor { disable_transaction_prewarming: config.disable_caching_and_prewarming(), evm_config, precompile_cache_enabled: config.precompile_cache_enabled(), - precompile_cache, + precompile_cache_map, _marker: Default::default(), } } @@ -266,7 +267,7 @@ where metrics: PrewarmMetrics::default(), terminate_execution: Arc::new(AtomicBool::new(false)), precompile_cache_enabled: self.precompile_cache_enabled, - precompile_cache: self.precompile_cache.clone(), + precompile_cache_map: self.precompile_cache_map.clone(), }; let prewarm_task = PrewarmCacheTask::new( @@ -434,13 +435,11 @@ impl ExecutionCache { #[cfg(test)] mod tests { - use std::sync::Arc; - use crate::tree::{ payload_processor::{ evm_state_to_hashed_post_state, executor::WorkloadExecutor, PayloadProcessor, }, - precompile_cache::PrecompileCache, + precompile_cache::PrecompileCacheMap, StateProviderBuilder, TreeConfig, }; use alloy_evm::block::StateChangeSource; @@ -460,6 +459,7 @@ mod tests { use reth_trie::{test_utils::state_root, HashedPostState, TrieInput}; use revm_primitives::{Address, HashMap, B256, KECCAK_EMPTY, U256}; use revm_state::{AccountInfo, AccountStatus, EvmState, EvmStorageSlot}; + use std::sync::Arc; fn create_mock_state_updates(num_accounts: usize, updates_per_account: usize) -> Vec { let mut rng = generators::rng(); @@ -563,7 +563,7 @@ mod tests { WorkloadExecutor::default(), EthEvmConfig::new(factory.chain_spec()), &TreeConfig::default(), - PrecompileCache::default(), + PrecompileCacheMap::default(), ); let provider = BlockchainProvider::new(factory).unwrap(); let mut handle = payload_processor.spawn( diff --git a/crates/engine/tree/src/tree/payload_processor/prewarm.rs b/crates/engine/tree/src/tree/payload_processor/prewarm.rs index a45ce70d8f..84466aa85b 100644 --- a/crates/engine/tree/src/tree/payload_processor/prewarm.rs +++ b/crates/engine/tree/src/tree/payload_processor/prewarm.rs @@ -5,7 +5,7 @@ use crate::tree::{ payload_processor::{ executor::WorkloadExecutor, multiproof::MultiProofMessage, ExecutionCache, }, - precompile_cache::{CachedPrecompile, PrecompileCache}, + precompile_cache::{CachedPrecompile, PrecompileCacheMap}, StateProviderBuilder, }; use alloy_consensus::transaction::Recovered; @@ -207,7 +207,7 @@ pub(super) struct PrewarmContext { /// An atomic bool that tells prewarm tasks to not start any more execution. pub(super) terminate_execution: Arc, pub(super) precompile_cache_enabled: bool, - pub(super) precompile_cache: PrecompileCache, + pub(super) precompile_cache_map: PrecompileCacheMap, } impl PrewarmContext @@ -230,7 +230,7 @@ where metrics, terminate_execution, precompile_cache_enabled, - precompile_cache, + mut precompile_cache_map, } = self; let state_provider = match provider.build() { @@ -261,8 +261,11 @@ where let mut evm = evm_config.evm_with_env(state_provider, evm_env); if precompile_cache_enabled { - evm.precompiles_mut().map_precompiles(|_, precompile| { - CachedPrecompile::wrap(precompile, precompile_cache.clone()) + evm.precompiles_mut().map_precompiles(|address, precompile| { + CachedPrecompile::wrap( + precompile, + precompile_cache_map.entry(*address).or_default().clone(), + ) }); } diff --git a/crates/engine/tree/src/tree/precompile_cache.rs b/crates/engine/tree/src/tree/precompile_cache.rs index fe3e797de2..ae23c0929e 100644 --- a/crates/engine/tree/src/tree/precompile_cache.rs +++ b/crates/engine/tree/src/tree/precompile_cache.rs @@ -1,10 +1,21 @@ //! Contains a precompile cache that is backed by a moka cache. +use alloy_primitives::map::Entry; use reth_evm::precompiles::{DynPrecompile, Precompile}; use revm::precompile::{PrecompileOutput, PrecompileResult}; -use revm_primitives::Bytes; +use revm_primitives::{Address, Bytes, HashMap}; use std::sync::Arc; +/// Stores caches for each precompile. +#[derive(Debug, Clone, Default)] +pub struct PrecompileCacheMap(HashMap); + +impl PrecompileCacheMap { + pub(crate) fn entry(&mut self, address: Address) -> Entry<'_, Address, PrecompileCache> { + self.0.entry(address) + } +} + /// Cache for precompiles, for each input stores the result. #[derive(Debug, Clone)] pub struct PrecompileCache(