feat(perf): integrate OnStateHook in executor (#11345)

crates/ethereum/evm/src/execute.rs

@@ -14,7 +14,7 @@ use reth_evm::{
         BatchExecutor, BlockExecutionError, BlockExecutionInput, BlockExecutionOutput,
         BlockExecutorProvider, BlockValidationError, Executor, ProviderError,
     },
-    system_calls::SystemCaller,
+    system_calls::{NoopHook, OnStateHook, SystemCaller},
     ConfigureEvm,
 };
 use reth_execution_types::ExecutionOutcome;
@@ -126,20 +126,25 @@ where
     /// This applies the pre-execution and post-execution changes that require an [EVM](Evm), and
     /// executes the transactions.
     ///
+    /// The optional `state_hook` will be executed with the state changes if present.
+    ///
     /// # Note
     ///
     /// It does __not__ apply post-execution changes that do not require an [EVM](Evm), for that see
     /// [`EthBlockExecutor::post_execution`].
-    fn execute_state_transitions<Ext, DB>(
+    fn execute_state_transitions<Ext, DB, F>(
         &self,
         block: &BlockWithSenders,
         mut evm: Evm<'_, Ext, &mut State<DB>>,
+        state_hook: Option<F>,
     ) -> Result<EthExecuteOutput, BlockExecutionError>
     where
         DB: Database,
         DB::Error: Into<ProviderError> + Display,
+        F: OnStateHook,
     {
-        let mut system_caller = SystemCaller::new(&self.evm_config, &self.chain_spec);
+        let mut system_caller =
+            SystemCaller::new(&self.evm_config, &self.chain_spec).with_state_hook(state_hook);
 
         system_caller.apply_pre_execution_changes(block, &mut evm)?;
 
@@ -161,7 +166,7 @@ where
             self.evm_config.fill_tx_env(evm.tx_mut(), transaction, *sender);
 
             // Execute transaction.
-            let ResultAndState { result, state } = evm.transact().map_err(move |err| {
+            let result_and_state = evm.transact().map_err(move |err| {
                 let new_err = err.map_db_err(|e| e.into());
                 // Ensure hash is calculated for error log, if not already done
                 BlockValidationError::EVM {
@@ -169,6 +174,8 @@ where
                     error: Box::new(new_err),
                 }
             })?;
+            system_caller.on_state(&result_and_state);
+            let ResultAndState { result, state } = result_and_state;
             evm.db_mut().commit(state);
 
             // append gas used
@@ -260,17 +267,31 @@ where
         EnvWithHandlerCfg::new_with_cfg_env(cfg, block_env, Default::default())
     }
 
+    /// Convenience method to invoke `execute_without_verification_with_state_hook` setting the
+    /// state hook as `None`.
+    fn execute_without_verification(
+        &mut self,
+        block: &BlockWithSenders,
+        total_difficulty: U256,
+    ) -> Result<EthExecuteOutput, BlockExecutionError> {
+        self.execute_without_verification_with_state_hook(block, total_difficulty, None::<NoopHook>)
+    }
+
     /// Execute a single block and apply the state changes to the internal state.
     ///
     /// Returns the receipts of the transactions in the block, the total gas used and the list of
     /// EIP-7685 [requests](Request).
     ///
     /// Returns an error if execution fails.
-    fn execute_without_verification(
+    fn execute_without_verification_with_state_hook<F>(
         &mut self,
         block: &BlockWithSenders,
         total_difficulty: U256,
-    ) -> Result<EthExecuteOutput, BlockExecutionError> {
+        state_hook: Option<F>,
+    ) -> Result<EthExecuteOutput, BlockExecutionError>
+    where
+        F: OnStateHook,
+    {
         // 1. prepare state on new block
         self.on_new_block(&block.header);
 
@@ -278,7 +299,7 @@ where
         let env = self.evm_env_for_block(&block.header, total_difficulty);
         let output = {
             let evm = self.executor.evm_config.evm_with_env(&mut self.state, env);
-            self.executor.execute_state_transitions(block, evm)
+            self.executor.execute_state_transitions(block, evm, state_hook)
         }?;
 
         // 3. apply post execution changes
@@ -368,6 +389,27 @@ where
         witness(&self.state);
         Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
     }
+
+    fn execute_with_state_hook<F>(
+        mut self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        let BlockExecutionInput { block, total_difficulty } = input;
+        let EthExecuteOutput { receipts, requests, gas_used } = self
+            .execute_without_verification_with_state_hook(
+                block,
+                total_difficulty,
+                Some(state_hook),
+            )?;
+
+        // NOTE: we need to merge keep the reverts for the bundle retention
+        self.state.merge_transitions(BundleRetention::Reverts);
+        Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
+    }
 }
 /// An executor for a batch of blocks.
 ///

crates/evm/src/either.rs

@@ -2,7 +2,10 @@
 
 use core::fmt::Display;
 
-use crate::execute::{BatchExecutor, BlockExecutorProvider, Executor};
+use crate::{
+    execute::{BatchExecutor, BlockExecutorProvider, Executor},
+    system_calls::OnStateHook,
+};
 use alloy_primitives::BlockNumber;
 use reth_execution_errors::BlockExecutionError;
 use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, ExecutionOutcome};
@@ -87,6 +90,20 @@ where
             Self::Right(b) => b.execute_with_state_witness(input, witness),
         }
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        match self {
+            Self::Left(a) => a.execute_with_state_hook(input, state_hook),
+            Self::Right(b) => b.execute_with_state_hook(input, state_hook),
+        }
+    }
 }
 
 impl<A, B, DB> BatchExecutor<DB> for Either<A, B>

crates/evm/src/execute.rs

@@ -12,6 +12,8 @@ use reth_prune_types::PruneModes;
 use revm::State;
 use revm_primitives::db::Database;
 
+use crate::system_calls::OnStateHook;
+
 /// A general purpose executor trait that executes an input (e.g. block) and produces an output
 /// (e.g. state changes and receipts).
 ///
@@ -43,6 +45,16 @@ pub trait Executor<DB> {
     ) -> Result<Self::Output, Self::Error>
     where
         F: FnMut(&State<DB>);
+
+    /// Executes the EVM with the given input and accepts a state hook closure that is invoked with
+    /// the EVM state after execution.
+    fn execute_with_state_hook<F>(
+        self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook;
 }
 
 /// A general purpose executor that can execute multiple inputs in sequence, validate the outputs,
@@ -199,6 +211,17 @@ mod tests {
         {
             Err(BlockExecutionError::msg("execution unavailable for tests"))
         }
+
+        fn execute_with_state_hook<F>(
+            self,
+            _: Self::Input<'_>,
+            _: F,
+        ) -> Result<Self::Output, Self::Error>
+        where
+            F: OnStateHook,
+        {
+            Err(BlockExecutionError::msg("execution unavailable for tests"))
+        }
     }
 
     impl<DB> BatchExecutor<DB> for TestExecutor<DB> {

crates/evm/src/noop.rs

@@ -10,7 +10,10 @@ use reth_storage_errors::provider::ProviderError;
 use revm::State;
 use revm_primitives::db::Database;
 
-use crate::execute::{BatchExecutor, BlockExecutorProvider, Executor};
+use crate::{
+    execute::{BatchExecutor, BlockExecutorProvider, Executor},
+    system_calls::OnStateHook,
+};
 
 const UNAVAILABLE_FOR_NOOP: &str = "execution unavailable for noop";
 
@@ -58,6 +61,17 @@ impl<DB> Executor<DB> for NoopBlockExecutorProvider {
     {
         Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
+    }
 }
 
 impl<DB> BatchExecutor<DB> for NoopBlockExecutorProvider {

crates/evm/src/system_calls/mod.rs

@@ -49,22 +49,19 @@ pub struct SystemCaller<'a, EvmConfig, Chainspec, Hook = NoopHook> {
     hook: Option<Hook>,
 }
 
-impl<'a, EvmConfig, Chainspec> SystemCaller<'a, EvmConfig, Chainspec> {
+impl<'a, EvmConfig, Chainspec> SystemCaller<'a, EvmConfig, Chainspec, NoopHook> {
     /// Create a new system caller with the given EVM config, database, and chain spec, and creates
     /// the EVM with the given initialized config and block environment.
     pub const fn new(evm_config: &'a EvmConfig, chain_spec: Chainspec) -> Self {
         Self { evm_config, chain_spec, hook: None }
     }
-}
-
-impl<'a, EvmConfig, Chainspec, Hook> SystemCaller<'a, EvmConfig, Chainspec, Hook> {
     /// Installs a custom hook to be called after each state change.
     pub fn with_state_hook<H: OnStateHook>(
         self,
-        hook: H,
+        hook: Option<H>,
     ) -> SystemCaller<'a, EvmConfig, Chainspec, H> {
         let Self { evm_config, chain_spec, .. } = self;
-        SystemCaller { evm_config, chain_spec, hook: Some(hook) }
+        SystemCaller { evm_config, chain_spec, hook }
     }
     /// Convenience method to consume the type and drop borrowed fields
     pub fn finish(self) {}
@@ -321,4 +318,11 @@ where
 
         eip7251::post_commit(result_and_state.result)
     }
+
+    /// Delegate to stored `OnStateHook`, noop if hook is `None`.
+    pub fn on_state(&mut self, state: &ResultAndState) {
+        if let Some(ref mut hook) = &mut self.hook {
+            hook.on_state(state);
+        }
+    }
 }

crates/evm/src/test_utils.rs

@@ -1,7 +1,10 @@
 //! Helpers for testing.
 
-use crate::execute::{
-    BatchExecutor, BlockExecutionInput, BlockExecutionOutput, BlockExecutorProvider, Executor,
+use crate::{
+    execute::{
+        BatchExecutor, BlockExecutionInput, BlockExecutionOutput, BlockExecutorProvider, Executor,
+    },
+    system_calls::OnStateHook,
 };
 use alloy_primitives::BlockNumber;
 use parking_lot::Mutex;
@@ -73,6 +76,17 @@ impl<DB> Executor<DB> for MockExecutorProvider {
     {
         unimplemented!()
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        unimplemented!()
+    }
 }
 
 impl<DB> BatchExecutor<DB> for MockExecutorProvider {

crates/optimism/evm/src/execute.rs

@@ -10,7 +10,7 @@ use reth_evm::{
         BatchExecutor, BlockExecutionError, BlockExecutionInput, BlockExecutionOutput,
         BlockExecutorProvider, BlockValidationError, Executor, ProviderError,
     },
-    system_calls::SystemCaller,
+    system_calls::{NoopHook, OnStateHook, SystemCaller},
     ConfigureEvm,
 };
 use reth_execution_types::ExecutionOutcome;
@@ -108,18 +108,23 @@ where
     ///
     /// This applies the pre-execution changes, and executes the transactions.
     ///
+    /// The optional `state_hook` will be executed with the state changes if present.
+    ///
     /// # Note
     ///
     /// It does __not__ apply post-execution changes.
-    fn execute_pre_and_transactions<Ext, DB>(
+    fn execute_pre_and_transactions<Ext, DB, F>(
         &self,
         block: &BlockWithSenders,
         mut evm: Evm<'_, Ext, &mut State<DB>>,
+        state_hook: Option<F>,
     ) -> Result<(Vec<Receipt>, u64), BlockExecutionError>
     where
         DB: Database<Error: Into<ProviderError> + Display>,
+        F: OnStateHook,
     {
-        let mut system_caller = SystemCaller::new(&self.evm_config, &self.chain_spec);
+        let mut system_caller =
+            SystemCaller::new(&self.evm_config, &self.chain_spec).with_state_hook(state_hook);
 
         // apply pre execution changes
         system_caller.apply_beacon_root_contract_call(
@@ -178,7 +183,7 @@ where
             self.evm_config.fill_tx_env(evm.tx_mut(), transaction, *sender);
 
             // Execute transaction.
-            let ResultAndState { result, state } = evm.transact().map_err(move |err| {
+            let result_and_state = evm.transact().map_err(move |err| {
                 let new_err = err.map_db_err(|e| e.into());
                 // Ensure hash is calculated for error log, if not already done
                 BlockValidationError::EVM {
@@ -192,7 +197,8 @@ where
                 ?transaction,
                 "Executed transaction"
             );
-
+            system_caller.on_state(&result_and_state);
+            let ResultAndState { result, state } = result_and_state;
             evm.db_mut().commit(state);
 
             // append gas used
@@ -278,16 +284,30 @@ where
         EnvWithHandlerCfg::new_with_cfg_env(cfg, block_env, Default::default())
     }
 
-    /// Execute a single block and apply the state changes to the internal state.
-    ///
-    /// Returns the receipts of the transactions in the block and the total gas used.
-    ///
-    /// Returns an error if execution fails.
+    /// Convenience method to invoke `execute_without_verification_with_state_hook` setting the
+    /// state hook as `None`.
     fn execute_without_verification(
         &mut self,
         block: &BlockWithSenders,
         total_difficulty: U256,
     ) -> Result<(Vec<Receipt>, u64), BlockExecutionError> {
+        self.execute_without_verification_with_state_hook(block, total_difficulty, None::<NoopHook>)
+    }
+
+    /// Execute a single block and apply the state changes to the internal state.
+    ///
+    /// Returns the receipts of the transactions in the block and the total gas used.
+    ///
+    /// Returns an error if execution fails.
+    fn execute_without_verification_with_state_hook<F>(
+        &mut self,
+        block: &BlockWithSenders,
+        total_difficulty: U256,
+        state_hook: Option<F>,
+    ) -> Result<(Vec<Receipt>, u64), BlockExecutionError>
+    where
+        F: OnStateHook,
+    {
         // 1. prepare state on new block
         self.on_new_block(&block.header);
 
@@ -296,7 +316,7 @@ where
 
         let (receipts, gas_used) = {
             let evm = self.executor.evm_config.evm_with_env(&mut self.state, env);
-            self.executor.execute_pre_and_transactions(block, evm)
+            self.executor.execute_pre_and_transactions(block, evm, state_hook)
         }?;
 
         // 3. apply post execution changes
@@ -383,6 +403,32 @@ where
             gas_used,
         })
     }
+
+    fn execute_with_state_hook<F>(
+        mut self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        let BlockExecutionInput { block, total_difficulty } = input;
+        let (receipts, gas_used) = self.execute_without_verification_with_state_hook(
+            block,
+            total_difficulty,
+            Some(state_hook),
+        )?;
+
+        // NOTE: we need to merge keep the reverts for the bundle retention
+        self.state.merge_transitions(BundleRetention::Reverts);
+
+        Ok(BlockExecutionOutput {
+            state: self.state.take_bundle(),
+            receipts,
+            requests: vec![],
+            gas_used,
+        })
+    }
 }
 
 /// An executor for a batch of blocks.