feat(engine): support dual-frontier partial trie persistence

bin/reth-bench/README.md

@@ -39,7 +39,7 @@ Both `new-payload-fcu` and `new-payload-only` support `--rpc-block-fetch-retries
 to control how many times block fetches are retried after an RPC failure. The default is `10`.
 Use `--rpc-block-fetch-retries forever` to keep retrying indefinitely.
 
-When using `--wait-for-persistence`, the benchmark waits after every `(threshold + 1)` blocks, where the threshold defaults to the engine's persistence threshold (2). This can be customized with `--persistence-threshold <N>`.
+When using `--wait-for-persistence`, the benchmark waits after every `(threshold + 1)` blocks, where the threshold defaults to the engine's persistence threshold. This can be customized with `--persistence-threshold <N>`.
 
 By default, the WebSocket URL for persistence subscriptions is derived from `--engine-rpc-url` (converting to ws:// on port 8546). Use `--ws-rpc-url` to override this.
 

bin/reth-bench/src/bench/new_payload_fcu.rs

@@ -25,18 +25,17 @@ use alloy_provider::{
     network::{AnyNetwork, AnyRpcBlock},
     Provider, RootProvider,
 };
-use alloy_rpc_types_engine::{
-    ExecutionData, ExecutionPayloadEnvelopeV5, ForkchoiceState, PayloadAttributes,
-};
+use alloy_rpc_types_engine::{ExecutionData, ForkchoiceState, PayloadAttributes};
 use clap::Parser;
 use eyre::{bail, ensure, Context, OptionExt};
 use futures::{stream, Stream, StreamExt, TryStreamExt};
 use reth_cli_runner::CliContext;
 use reth_engine_primitives::config::DEFAULT_PERSISTENCE_THRESHOLD;
 use reth_node_api::{EngineApiMessageVersion, ExecutionPayload};
 use reth_node_core::args::{BenchmarkArgs, WaitForPersistence};
-use reth_rpc_api::{RethNewPayloadInput, TestingBuildBlockRequestV1};
+use reth_rpc_api::{RethNewPayloadInput, TestingBuildBlockRequestV1, TestingBuildBlockResponseV1};
 use std::{
+    path::Path,
     pin::Pin,
     time::{Duration, Instant},
 };
@@ -189,6 +188,7 @@ impl Command {
         if let Some(depth) = self.reorg {
             info!(target: "reth-bench", depth, "Using testing_buildBlockV1 reorg mode");
         }
+        let output_dir = self.benchmark.output.clone();
 
         let BenchContext {
             benchmark_mode,
@@ -213,13 +213,15 @@ impl Command {
         let buffer_size = self.rpc_block_buffer_size;
         let provider = block_provider.clone();
         let bench_mode = benchmark_mode.clone();
+        let artifact_output_dir = output_dir.clone();
         let mut blocks: Pin<Box<dyn Stream<Item = eyre::Result<Payload>> + Send>> = Box::pin(
             stream::iter((next_block..)
                 .take_while(move |next_block| {
                     bench_mode.contains(*next_block)
                 }))
                 .map(move |next_block| {
                     let block_provider = provider.clone();
+                    let artifact_output_dir = artifact_output_dir.clone();
                     async move {
                         let block_res = block_provider
                             .get_block_by_number(next_block.into())
@@ -238,10 +240,47 @@ impl Command {
                             };
 
 
-                        let bal = if !rlp_blocks &&
-                            (block.header.block_access_list_hash.is_some() || self.enable_bal)
-                        {
-                            Some(fetch_block_access_list(&block_provider, block.header.number).await?)
+                        let fetched_bal = if !rlp_blocks {
+                            match fetch_block_access_list(&block_provider, block.header.number).await {
+                                Ok(bal) => {
+                                    write_bal_artifact(
+                                        artifact_output_dir.as_deref(),
+                                        "real",
+                                        block.header.number,
+                                        block.header.hash,
+                                        Some(&bal),
+                                    )?;
+                                    Some(bal)
+                                }
+                                Err(err) => {
+                                    warn!(
+                                        target: "reth-bench",
+                                        block_number = block.header.number,
+                                        block_hash = %block.header.hash,
+                                        %err,
+                                        "Failed to fetch real block BAL artifact"
+                                    );
+                                    if is_unsupported_bal_rpc_error(&err) {
+                                        warn!(
+                                            target: "reth-bench",
+                                            "Remote RPC does not support BAL fetching; writing null real-block BAL artifact"
+                                        );
+                                    }
+                                    write_bal_artifact(
+                                        artifact_output_dir.as_deref(),
+                                        "real",
+                                        block.header.number,
+                                        block.header.hash,
+                                        None,
+                                    )?;
+                                    None
+                                }
+                            }
+                        } else {
+                            None
+                        };
+                        let bal = if block.header.block_access_list_hash.is_some() || self.enable_bal {
+                            fetched_bal
                         } else {
                             None
                         };
@@ -404,6 +443,7 @@ impl Command {
                                 .ok_or_eyre("missing deferred fork block for reorg branch start")?,
                             canonical_parent_hash,
                             no_wait_for_caches,
+                            output_dir.as_deref(),
                         )
                         .await?,
                     });
@@ -448,6 +488,7 @@ impl Command {
                         next_fork_block_number,
                         Some(prepared.block_hash),
                         no_wait_for_caches,
+                        output_dir.as_deref(),
                     )
                     .await?;
                 } else {
@@ -534,12 +575,13 @@ async fn prepare_built_block(
     block: &AnyRpcBlock,
     parent_block_hash: B256,
     no_wait_for_caches: bool,
+    output_dir: Option<&Path>,
 ) -> eyre::Result<PreparedBuiltBlock> {
     const MAX_BUILD_ATTEMPTS: usize = 10;
     const BUILD_RETRY_INTERVAL: Duration = Duration::from_millis(100);
 
     let request = build_block_request(block, parent_block_hash)?;
-    let built_payload: ExecutionPayloadEnvelopeV5 = {
+    let built_response: TestingBuildBlockResponseV1 = {
         let mut attempts_remaining = MAX_BUILD_ATTEMPTS;
 
         loop {
@@ -569,8 +611,16 @@ async fn prepare_built_block(
         }
     };
 
+    let built_payload = built_response.execution_payload_envelope;
     let payload = &built_payload.execution_payload.payload_inner.payload_inner;
     let block_hash = payload.block_hash;
+    write_bal_artifact(
+        output_dir,
+        "fork",
+        payload.block_number,
+        block_hash,
+        built_response.block_access_list.as_ref(),
+    )?;
     let (payload, sidecar) = built_payload
         .into_payload_and_sidecar(block.header.parent_beacon_block_root.unwrap_or_default());
     // Fork payloads are built immediately before the next `testing_buildBlockV1` call. Leaving
@@ -593,9 +643,10 @@ async fn queue_fork_block(
     block_number: u64,
     parent_block_hash: Option<B256>,
     no_wait_for_caches: bool,
+    output_dir: Option<&Path>,
 ) -> eyre::Result<Option<QueuedForkBlock>> {
     if !benchmark_mode.contains(block_number) {
-        return Ok(None)
+        return Ok(None);
     }
 
     let future_block = block_provider
@@ -613,17 +664,50 @@ async fn queue_fork_block(
             &future_block,
             parent_block_hash,
             no_wait_for_caches,
+            output_dir,
         )
         .await?,
     }))
 }
 
+fn write_bal_artifact(
+    output_dir: Option<&Path>,
+    kind: &str,
+    block_number: u64,
+    block_hash: B256,
+    block_access_list: Option<&BlockAccessList>,
+) -> eyre::Result<()> {
+    let Some(output_dir) = output_dir else { return Ok(()) };
+
+    let bal_dir = output_dir.join("block-access-lists");
+    std::fs::create_dir_all(&bal_dir)?;
+    let path = bal_dir.join(format!("bal-{kind}-{block_number}-{block_hash}.json"));
+    let value = serde_json::json!({
+        "kind": kind,
+        "blockNumber": block_number,
+        "blockHash": block_hash,
+        "blockAccessList": block_access_list,
+    });
+    let file = std::fs::File::create(&path)?;
+    serde_json::to_writer_pretty(file, &value)?;
+    debug!(target: "reth-bench", %kind, block_number, %block_hash, path = %path.display(), "Wrote BAL artifact");
+    Ok(())
+}
+
 fn is_retryable_build_block_error(err: &alloy_transport::TransportError) -> bool {
     let message = err.to_string();
     message.contains("block not found: hash") ||
         message.contains("block hash not found for block number")
 }
 
+fn is_unsupported_bal_rpc_error(err: &eyre::Report) -> bool {
+    let message = err.to_string();
+    message.contains("method ignored") ||
+        message.contains("Method not found") ||
+        message.contains("method not found") ||
+        message.contains("-32601")
+}
+
 fn build_block_request(
     block: &AnyRpcBlock,
     parent_block_hash: B256,

bin/reth-bench/src/bench/replay_payloads.rs

@@ -29,6 +29,7 @@ use reth_node_api::EngineApiMessageVersion;
 use reth_node_core::args::WaitForPersistence;
 use reth_rpc_api::RethNewPayloadInput;
 use std::{
+    collections::HashMap,
     path::PathBuf,
     time::{Duration, Instant},
 };
@@ -228,14 +229,15 @@ impl Command {
             );
         }
 
-        let mut parent_hash = initial_parent_hash;
+        let mut replayed_hashes = HashMap::from([(initial_parent_hash, initial_parent_hash)]);
 
         let mut results = Vec::new();
         let total_benchmark_duration = Instant::now();
 
         for (i, payload) in payloads.iter().enumerate() {
             let execution_data = &payload.execution_data;
             let mut block_hash = payload.block_hash;
+            let original_block_hash = block_hash;
             let v1 = execution_data.payload.as_v1();
 
             let gas_used = v1.gas_used;
@@ -274,11 +276,16 @@ impl Command {
                     .unwrap_or(WaitForPersistence::Never)
                     .rpc_value(block_number);
 
-                // Inject sidecar BAL into the inline V4 payload field when --bal is set.
-                // If the payload is not already V4 we upgrade it (V3→V4) so the BAL
-                // can be carried inline. This changes the block hash, so we recompute
-                // it and patch parent_hash to maintain the chain.
                 let mut execution_data = execution_data.clone();
+                let original_parent_hash = execution_data.payload.as_v1().parent_hash;
+                let mut payload_modified = false;
+                if let Some(remapped_parent_hash) = replayed_hashes.get(&original_parent_hash) {
+                    if *remapped_parent_hash != original_parent_hash {
+                        execution_data.payload.as_v1_mut().parent_hash = *remapped_parent_hash;
+                        payload_modified = true;
+                    }
+                }
+
                 if self.bal &&
                     let Some(bal) = &payload.block_access_list
                 {
@@ -292,12 +299,10 @@ impl Command {
                         execution_data.payload.as_v4_mut().unwrap().block_access_list = encoded_bal;
                     }
 
-                    // Patch parent_hash so this block chains off the (possibly
-                    // rehashed) previous block.
-                    execution_data.payload.as_v1_mut().parent_hash = parent_hash;
+                    payload_modified = true;
+                }
 
-                    // Recompute block hash after payload modification and update
-                    // the hash stored in the payload itself.
+                if payload_modified {
                     block_hash = compute_payload_block_hash(&execution_data)?;
                     execution_data.payload.as_v1_mut().block_hash = block_hash;
                 }
@@ -349,8 +354,8 @@ impl Command {
 
             let fcu_state = ForkchoiceState {
                 head_block_hash: block_hash,
-                safe_block_hash: parent_hash,
-                finalized_block_hash: parent_hash,
+                safe_block_hash: initial_parent_hash,
+                finalized_block_hash: initial_parent_hash,
             };
 
             let fcu_start = Instant::now();
@@ -390,7 +395,7 @@ impl Command {
                 TotalGasRow { block_number, transaction_count, gas_used, time: current_duration };
             results.push((gas_row, combined_result));
 
-            parent_hash = block_hash;
+            replayed_hashes.insert(original_block_hash, block_hash);
         }
 
         let (gas_output_results, combined_results): (Vec<TotalGasRow>, Vec<CombinedResult>) =

crates/chain-state/src/in_memory.rs

@@ -320,6 +320,19 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
     /// This will update the links between blocks and remove all blocks that are [..
     /// `persisted_height`].
     pub fn remove_persisted_blocks(&self, persisted_num_hash: BlockNumHash) {
+        self.remove_persisted_blocks_until(persisted_num_hash, persisted_num_hash.number);
+    }
+
+    /// Removes blocks from the in-memory state through `remove_until` while still reporting the
+    /// provided block as the persisted tip.
+    ///
+    /// This is used when block bodies/plain state have been persisted further than trie data, so a
+    /// suffix still needs to remain in memory for trie-backed operations.
+    pub fn remove_persisted_blocks_until(
+        &self,
+        persisted_num_hash: BlockNumHash,
+        remove_until: BlockNumber,
+    ) {
         self.set_persisted(persisted_num_hash);
         // if the persisted hash is not in the canonical in memory state, do nothing, because it
         // means canonical blocks were not actually persisted.
@@ -337,16 +350,15 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
             let mut numbers = self.inner.in_memory_state.numbers.write();
             let mut blocks = self.inner.in_memory_state.blocks.write();
 
-            let BlockNumHash { number: persisted_height, hash: _ } = persisted_num_hash;
+            let remove_until = remove_until.min(persisted_num_hash.number);
 
             // clear all numbers
             numbers.clear();
 
-            // drain all blocks and only keep the ones that are not persisted (below the persisted
-            // height)
+            // Drain all blocks and keep only the suffix that still has to stay in memory.
             let mut old_blocks = blocks
                 .drain()
-                .filter(|(_, b)| b.block_ref().recovered_block().number() > persisted_height)
+                .filter(|(_, b)| b.block_ref().recovered_block().number() > remove_until)
                 .map(|(_, b)| b.block.clone())
                 .collect::<Vec<_>>();