fix(bench): reth-bb multiple executors

bin/reth-bb/src/evm.rs

@@ -36,6 +36,7 @@ use tracing::{debug, trace};
 // ---------------------------------------------------------------------------
 
 /// Runtime state for segment boundary tracking.
+#[derive(Clone)]
 pub(crate) struct BbEvmPlan {
     /// The segment boundaries and environments.
     pub(crate) segments: Vec<BigBlockSegment>,
@@ -73,6 +74,10 @@ impl BbEvmPlan {
             .filter(|(n, _)| *n >= min && *n < block_number)
             .collect()
     }
+
+    pub(crate) fn segment_index_for_tx(&self, tx_index: usize) -> usize {
+        self.segments.partition_point(|segment| segment.start_tx <= tx_index).saturating_sub(1)
+    }
 }
 
 impl std::fmt::Debug for BbEvmPlan {
@@ -97,6 +102,9 @@ impl std::fmt::Debug for BbEvmPlan {
 /// segment boundaries without requiring additional trait bounds on `DB`.
 pub(crate) type BlockHashSeeder<DB> = fn(&mut DB, &[(u64, B256)]);
 
+/// Function pointer that reads the BAL index from the DB.
+pub(crate) type BalIndexReader<DB> = fn(&DB) -> u64;
+
 /// Block executor that wraps [`EthBlockExecutor`] and handles segment-boundary
 /// changes for big-block execution.
 ///
@@ -131,6 +139,10 @@ where
     /// Callback to reseed block hashes into the DB's cache at segment
     /// boundaries. See [`BlockHashSeeder`].
     block_hash_seeder: Option<BlockHashSeeder<DB>>,
+    /// Callback to read the BAL index from the DB.
+    bal_index_reader: Option<BalIndexReader<DB>>,
+    /// Whether the executor has selected its starting segment.
+    initialized: bool,
 }
 
 impl<'a, DB, I, P, Spec> BbBlockExecutor<'a, DB, I, P, Spec>
@@ -156,6 +168,7 @@ where
         receipt_builder: RethReceiptBuilder,
         plan: Option<BbEvmPlan>,
         block_hash_seeder: Option<BlockHashSeeder<DB>>,
+        bal_index_reader: Option<BalIndexReader<DB>>,
     ) -> Self {
         let inner = EthBlockExecutor::new(evm, ctx, spec, receipt_builder);
         Self {
@@ -166,7 +179,61 @@ where
             blob_gas_used_offset: 0,
             shared_hook: Arc::new(Mutex::new(None)),
             block_hash_seeder,
+            bal_index_reader,
+            initialized: false,
+        }
+    }
+
+    fn initialize(&mut self) -> Result<(), BlockExecutionError> {
+        if self.initialized {
+            return Ok(());
+        }
+
+        let plan = match &self.plan {
+            Some(plan) => plan,
+            None => return Ok(()),
+        };
+
+        self.initialized = true;
+
+        let bal_index =
+            self.bal_index_reader.map(|reader| reader(self.inner().evm().db())).unwrap_or(0);
+        let segment_idx =
+            if bal_index == 0 { 0 } else { plan.segment_index_for_tx((bal_index - 1) as usize) };
+        let segment = &plan.segments[segment_idx];
+
+        // Swap the EVM's block_env and executor ctx to the selected segment's
+        // values so that EIP-2935/EIP-4788 system calls use the correct block
+        // number and parent hash. Without this, the outer big block header's
+        // block_number (which is synthetic) would be used, writing to wrong
+        // EIP-2935 slots and corrupting state.
+        let block_env = segment.evm_env.block_env.clone();
+        let block_number = block_env.number.saturating_to::<u64>();
+        let mut cfg_env = segment.evm_env.cfg_env.clone();
+        cfg_env.disable_base_fee = true;
+        let ctx = EthBlockExecutionCtx {
+            parent_hash: segment.ctx.parent_hash,
+            parent_beacon_block_root: segment.ctx.parent_beacon_block_root,
+            ommers: segment.ctx.ommers,
+            withdrawals: segment.ctx.withdrawals.clone(),
+            extra_data: segment.ctx.extra_data.clone(),
+            tx_count_hint: segment.ctx.tx_count_hint,
+            slot_number: segment.ctx.slot_number,
+        };
+
+        let inner = self.inner_mut();
+        let evm_ctx = inner.evm.ctx_mut();
+        evm_ctx.block = block_env;
+        evm_ctx.cfg = cfg_env;
+        inner.ctx = ctx;
+
+        self.reseed_block_hashes_for(block_number);
+
+        if bal_index > 0 {
+            self.plan = None;
         }
+
+        Ok(())
     }
 
     /// Creates a forwarding `OnStateHook` that delegates to the shared hook.
@@ -349,42 +416,17 @@ where
     type Result = EthTxResult<HaltReason, alloy_consensus::TxType>;
 
     fn apply_pre_execution_changes(&mut self) -> Result<(), BlockExecutionError> {
-        // Swap the EVM's block_env and executor ctx to the first segment's
-        // values so that the initial EIP-2935/EIP-4788 system calls use the
-        // correct block number and parent hash. Without this, the outer big
-        // block header's block_number (which is synthetic) would be used,
-        // writing to wrong EIP-2935 slots and corrupting state.
-        if let Some(seg0) = self.plan.as_ref().map(|p| &p.segments[0]) {
-            let block_env = seg0.evm_env.block_env.clone();
-            let block_number = block_env.number.saturating_to::<u64>();
-            let mut cfg_env = seg0.evm_env.cfg_env.clone();
-            cfg_env.disable_base_fee = true;
-            let seg0_ctx = EthBlockExecutionCtx {
-                parent_hash: seg0.ctx.parent_hash,
-                parent_beacon_block_root: seg0.ctx.parent_beacon_block_root,
-                ommers: seg0.ctx.ommers,
-                withdrawals: seg0.ctx.withdrawals.clone(),
-                extra_data: seg0.ctx.extra_data.clone(),
-                tx_count_hint: seg0.ctx.tx_count_hint,
-                slot_number: seg0.ctx.slot_number,
-            };
-
-            let inner = self.inner_mut();
-            let evm_ctx = inner.evm.ctx_mut();
-            evm_ctx.block = block_env;
-            evm_ctx.cfg = cfg_env;
-            inner.ctx = seg0_ctx;
-
-            self.reseed_block_hashes_for(block_number);
-        }
-
+        // The outer big-block header uses a synthetic block number, so start
+        // system calls must run against the selected real segment env.
+        self.initialize()?;
         self.inner_mut().apply_pre_execution_changes()
     }
 
     fn execute_transaction_without_commit(
         &mut self,
         tx: impl ExecutableTx<Self>,
     ) -> Result<Self::Result, BlockExecutionError> {
+        self.initialize()?;
         self.maybe_apply_boundary()?;
         self.inner_mut().execute_transaction_without_commit(tx)
     }
@@ -393,6 +435,7 @@ where
         &mut self,
         output: Self::Result,
     ) -> Result<GasOutput, BlockExecutionError> {
+        self.maybe_apply_boundary()?;
         let gas_used = self.inner_mut().commit_transaction(output)?;
 
         // Fix up cumulative_gas_used on the just-committed receipt so that
@@ -499,7 +542,7 @@ pub struct BbBlockExecutorFactory<Spec> {
     receipt_builder: RethReceiptBuilder,
     spec: Spec,
     evm_factory: EthEvmFactory,
-    /// Staged plan consumed by the next [`BbBlockExecutor`].
+    /// Staged plan cloned into each [`BbBlockExecutor`].
     pub(crate) staged_plan: Arc<Mutex<Option<BbEvmPlan>>>,
 }
 
@@ -528,23 +571,36 @@ impl<Spec> BbBlockExecutorFactory<Spec> {
         *self.staged_plan.lock().unwrap() = Some(plan);
     }
 
-    fn take_plan(&self) -> Option<BbEvmPlan> {
-        self.staged_plan.lock().unwrap().take()
+    pub(crate) fn clear_staged_plan(&self) {
+        *self.staged_plan.lock().unwrap() = None;
+    }
+
+    fn peek_plan(&self) -> Option<BbEvmPlan> {
+        self.staged_plan.lock().unwrap().clone()
     }
 
     pub(crate) fn create_executor_with_seeder<'a, DB, I>(
         &'a self,
         evm: EthEvm<DB, I, PrecompilesMap>,
         ctx: EthBlockExecutionCtx<'a>,
         block_hash_seeder: Option<BlockHashSeeder<DB>>,
+        bal_index_reader: Option<BalIndexReader<DB>>,
     ) -> BbBlockExecutor<'a, DB, I, PrecompilesMap, &'a Spec>
     where
         Spec: alloy_evm::eth::spec::EthExecutorSpec,
         DB: StateDB + 'a,
         I: Inspector<EthEvmContext<DB>> + 'a,
     {
-        let plan = self.take_plan();
-        BbBlockExecutor::new(evm, ctx, &self.spec, self.receipt_builder, plan, block_hash_seeder)
+        let plan = self.peek_plan();
+        BbBlockExecutor::new(
+            evm,
+            ctx,
+            &self.spec,
+            self.receipt_builder,
+            plan,
+            block_hash_seeder,
+            bal_index_reader,
+        )
     }
 }
 
@@ -571,7 +627,7 @@ where
         DB: StateDB + 'a,
         I: Inspector<EthEvmContext<DB>> + 'a,
     {
-        let plan = self.take_plan();
-        BbBlockExecutor::new(evm, ctx, &self.spec, self.receipt_builder, plan, None)
+        let plan = self.peek_plan();
+        BbBlockExecutor::new(evm, ctx, &self.spec, self.receipt_builder, plan, None, None)
     }
 }

bin/reth-bb/src/evm_config.rs

@@ -8,7 +8,7 @@
 pub(crate) use reth_engine_primitives::BigBlockData;
 
 use crate::{
-    evm::{BbBlockExecutorFactory, BbEvmPlan},
+    evm::{BalIndexReader, BbBlockExecutorFactory, BbEvmPlan},
     BigBlockMap,
 };
 use alloy_consensus::Header;
@@ -55,7 +55,7 @@ pub(crate) struct BigBlockSegment {
 ///
 /// Wraps [`EthEvmConfig`] and a shared [`BigBlockMap`]. When a big-block
 /// payload is received, the plan is staged on the [`BbBlockExecutorFactory`]
-/// and consumed when the executor is created. Block hashes for inter-segment
+/// and cloned when executors are created. Block hashes for inter-segment
 /// BLOCKHASH resolution are reseeded into `State::block_hashes` at each
 /// segment boundary via a [`BlockHashSeeder`](crate::evm::BlockHashSeeder)
 /// callback injected in [`ConfigureEvm::create_executor`].
@@ -106,6 +106,10 @@ fn seed_state_block_hashes<DB>(state: &mut &mut revm::database::State<DB>, hashe
     }
 }
 
+fn read_bal_index<DB>(state: &&mut revm::database::State<DB>) -> u64 {
+    state.bal_state.bal_index()
+}
+
 // ---------------------------------------------------------------------------
 // ConfigureEvm
 // ---------------------------------------------------------------------------
@@ -144,6 +148,12 @@ where
         &self,
         block: &'a SealedBlock<reth_ethereum_primitives::Block>,
     ) -> Result<EthBlockExecutionCtx<'a>, Self::Error> {
+        if let Some(plan) = self.plan_for_payload_hash(&block.hash()) {
+            self.executor_factory.stage_plan(plan);
+        } else {
+            self.executor_factory.clear_staged_plan();
+        }
+
         self.inner.context_for_block(block)
     }
 
@@ -169,15 +179,16 @@ where
         DB: Database,
         I: reth_evm::InspectorFor<Self, &'a mut revm::database::State<DB>> + 'a,
     {
-        // Use create_executor_with_seeder to inject a concrete seeder that
-        // can reseed State::block_hashes at segment boundaries. The seeder
-        // is a function pointer that knows the concrete State<DB> type,
-        // allowing the generic BbBlockExecutor to reseed without additional
-        // trait bounds on DB.
+        let bal_index_reader: Option<BalIndexReader<&'a mut revm::database::State<DB>>> =
+            Some(read_bal_index::<DB>);
+
+        // Inject concrete function pointers that know the `State<DB>` type so
+        // the generic executor can reseed block hashes and read `bal_index`.
         self.executor_factory.create_executor_with_seeder(
             evm,
             ctx,
             Some(seed_state_block_hashes::<DB>),
+            bal_index_reader,
         )
     }
 }
@@ -214,6 +225,7 @@ where
 
             Ok(env)
         } else {
+            self.executor_factory.clear_staged_plan();
             self.inner.evm_env_for_payload(payload)
         }
     }
@@ -248,10 +260,12 @@ where
     /// In practice, this is called from `evm_env_for_payload` in the
     /// engine pipeline.
     pub fn stage_plan_for_payload(&self, payload_hash: &B256) {
-        let bb = match self.pending.lock().unwrap().remove(payload_hash) {
-            Some(bb) => bb,
-            None => return,
-        };
+        let Some(plan) = self.plan_for_payload_hash(payload_hash) else { return };
+        self.executor_factory.stage_plan(plan);
+    }
+
+    fn plan_for_payload_hash(&self, payload_hash: &B256) -> Option<BbEvmPlan> {
+        let bb = self.pending.lock().unwrap().remove(payload_hash)?;
 
         let segments: Vec<_> = bb
             .env_switches
@@ -287,6 +301,6 @@ where
 
         plan.block_hashes_to_seed.sort_unstable_by_key(|(n, _)| *n);
 
-        self.executor_factory.stage_plan(plan);
+        Some(plan)
     }
 }