feat(rpc, reth-bench): reth_newPayload methods for reth-bench (#22133)

Co-authored-by: Amp <amp@ampcode.com>
Co-authored-by: Georgios Konstantopoulos <me@gakonst.com>

crates/engine/tree/src/launch.rs

@@ -10,7 +10,7 @@ use crate::{
     download::BasicBlockDownloader,
     engine::{EngineApiKind, EngineApiRequest, EngineApiRequestHandler, EngineHandler},
     persistence::PersistenceHandle,
-    tree::{EngineApiTreeHandler, EngineValidator, TreeConfig},
+    tree::{EngineApiTreeHandler, EngineValidator, TreeConfig, WaitForCaches},
 };
 use futures::Stream;
 use reth_consensus::FullConsensus;
@@ -76,7 +76,7 @@ where
     N: ProviderNodeTypes,
     Client: BlockClient<Block = <N::Primitives as NodePrimitives>::Block> + 'static,
     S: Stream<Item = BeaconEngineMessage<N::Payload>> + Send + Sync + Unpin + 'static,
-    V: EngineValidator<N::Payload>,
+    V: EngineValidator<N::Payload> + WaitForCaches,
     C: ConfigureEvm<Primitives = N::Primitives> + 'static,
 {
     let downloader = BasicBlockDownloader::new(client, consensus.clone());

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

@@ -19,7 +19,7 @@ use reth_chain_state::{
 use reth_consensus::{Consensus, FullConsensus};
 use reth_engine_primitives::{
     BeaconEngineMessage, BeaconOnNewPayloadError, ConsensusEngineEvent, ExecutionPayload,
-    ForkchoiceStateTracker, OnForkChoiceUpdated,
+    ForkchoiceStateTracker, NewPayloadTimings, OnForkChoiceUpdated,
 };
 use reth_errors::{ConsensusError, ProviderResult};
 use reth_evm::ConfigureEvm;
@@ -323,7 +323,7 @@ where
         + StorageSettingsCache,
     C: ConfigureEvm<Primitives = N> + 'static,
     T: PayloadTypes<BuiltPayload: BuiltPayload<Primitives = N>>,
-    V: EngineValidator<T>,
+    V: EngineValidator<T> + WaitForCaches,
 {
     /// Creates a new [`EngineApiTreeHandler`].
     #[expect(clippy::too_many_arguments)]
@@ -1555,6 +1555,81 @@ where
                                 // handle the event if any
                                 self.on_maybe_tree_event(maybe_event)?;
                             }
+                            BeaconEngineMessage::RethNewPayload { payload, tx } => {
+                                let pending_persistence = self.persistence_state.rx.take();
+                                let validator = &self.payload_validator;
+
+                                debug!(target: "engine::tree", "Waiting for persistence and caches in parallel before processing reth_newPayload");
+                                let (persistence_tx, persistence_rx) = std::sync::mpsc::channel();
+                                if let Some((rx, start_time, _action)) = pending_persistence {
+                                    tokio::task::spawn_blocking(move || {
+                                        let start = Instant::now();
+                                        let result = rx.recv().ok();
+                                        let _ = persistence_tx.send((
+                                            result,
+                                            start_time,
+                                            start.elapsed(),
+                                        ));
+                                    });
+                                }
+
+                                let cache_wait = validator.wait_for_caches();
+                                let persistence_result = persistence_rx.try_recv().ok();
+
+                                let persistence_wait =
+                                    if let Some((result, start_time, wait_duration)) =
+                                        persistence_result
+                                    {
+                                        let _ = self
+                                            .on_persistence_complete(result.flatten(), start_time);
+                                        Some(wait_duration)
+                                    } else {
+                                        None
+                                    };
+
+                                debug!(
+                                    target: "engine::tree",
+                                    persistence_wait = ?persistence_wait,
+                                    execution_cache_wait = ?cache_wait.execution_cache,
+                                    sparse_trie_wait = ?cache_wait.sparse_trie,
+                                    "Peresistence finished and caches updated for reth_newPayload"
+                                );
+
+                                let start = Instant::now();
+                                let gas_used = payload.gas_used();
+                                let num_hash = payload.num_hash();
+                                let mut output = self.on_new_payload(payload);
+                                let latency = start.elapsed();
+                                self.metrics.engine.new_payload.update_response_metrics(
+                                    start,
+                                    &mut self.metrics.engine.forkchoice_updated.latest_finish_at,
+                                    &output,
+                                    gas_used,
+                                );
+
+                                let maybe_event =
+                                    output.as_mut().ok().and_then(|out| out.event.take());
+
+                                let timings = NewPayloadTimings {
+                                    latency,
+                                    persistence_wait,
+                                    execution_cache_wait: cache_wait.execution_cache,
+                                    sparse_trie_wait: cache_wait.sparse_trie,
+                                };
+                                if let Err(err) =
+                                    tx.send(output.map(|o| (o.outcome, timings)).map_err(|e| {
+                                        BeaconOnNewPayloadError::Internal(Box::new(e))
+                                    }))
+                                {
+                                    error!(target: "engine::tree", payload=?num_hash, elapsed=?start.elapsed(), "Failed to send event: {err:?}");
+                                    self.metrics
+                                        .engine
+                                        .failed_new_payload_response_deliveries
+                                        .increment(1);
+                                }
+
+                                self.on_maybe_tree_event(maybe_event)?;
+                            }
                         }
                     }
                 }

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

@@ -50,6 +50,7 @@ use std::{
         mpsc::{self, channel},
         Arc,
     },
+    time::Duration,
 };
 use tracing::{debug, debug_span, instrument, warn, Span};
 
@@ -62,6 +63,26 @@ pub mod sparse_trie;
 
 use preserved_sparse_trie::{PreservedSparseTrie, SharedPreservedSparseTrie};
 
+/// Result of waiting for caches to become available.
+#[derive(Debug, Clone, Copy, Default)]
+pub struct CacheWaitDurations {
+    /// Time spent waiting for the execution cache lock.
+    pub execution_cache: Duration,
+    /// Time spent waiting for the sparse trie lock.
+    pub sparse_trie: Duration,
+}
+
+/// Trait for types that can wait for execution cache and sparse trie locks to become available.
+///
+/// This is used by `reth_newPayload` endpoint to ensure that payload processing
+/// waits for any ongoing operations to complete before starting.
+pub trait WaitForCaches {
+    /// Waits for persistence and cache updates to complete.
+    ///
+    /// Returns the time spent waiting for each cache separately.
+    fn wait_for_caches(&self) -> CacheWaitDurations;
+}
+
 /// Default parallelism thresholds to use with the [`ParallelSparseTrie`].
 ///
 /// These values were determined by performing benchmarks using gradually increasing values to judge
@@ -178,6 +199,46 @@ where
     }
 }
 
+impl<Evm> WaitForCaches for PayloadProcessor<Evm>
+where
+    Evm: ConfigureEvm,
+{
+    fn wait_for_caches(&self) -> CacheWaitDurations {
+        debug!(target: "engine::tree::payload_processor", "Waiting for execution cache and sparse trie locks");
+
+        // Wait for both caches in parallel using std threads
+        let execution_cache = self.execution_cache.clone();
+        let sparse_trie = self.sparse_state_trie.clone();
+
+        // Use channels and spawn_blocking instead of std::thread::spawn
+        let (execution_tx, execution_rx) = std::sync::mpsc::channel();
+        let (sparse_trie_tx, sparse_trie_rx) = std::sync::mpsc::channel();
+
+        self.executor.spawn_blocking(move || {
+            let _ = execution_tx.send(execution_cache.wait_for_availability());
+        });
+        self.executor.spawn_blocking(move || {
+            let _ = sparse_trie_tx.send(sparse_trie.wait_for_availability());
+        });
+
+        let execution_cache_duration =
+            execution_rx.recv().expect("execution cache wait task failed to send result");
+        let sparse_trie_duration =
+            sparse_trie_rx.recv().expect("sparse trie wait task failed to send result");
+
+        debug!(
+            target: "engine::tree::payload_processor",
+            ?execution_cache_duration,
+            ?sparse_trie_duration,
+            "Execution cache and sparse trie locks acquired"
+        );
+        CacheWaitDurations {
+            execution_cache: execution_cache_duration,
+            sparse_trie: sparse_trie_duration,
+        }
+    }
+}
+
 impl<N, Evm> PayloadProcessor<Evm>
 where
     N: NodePrimitives,
@@ -966,6 +1027,27 @@ impl PayloadExecutionCache {
         self.inner.write().take();
     }
 
+    /// Waits until the execution cache becomes available for use.
+    ///
+    /// This acquires a write lock to ensure exclusive access, then immediately releases it.
+    /// This is useful for synchronization before starting payload processing.
+    ///
+    /// Returns the time spent waiting for the lock.
+    pub fn wait_for_availability(&self) -> Duration {
+        let start = Instant::now();
+        // Acquire write lock to wait for any current holders to finish
+        let _guard = self.inner.write();
+        let elapsed = start.elapsed();
+        if elapsed.as_millis() > 5 {
+            debug!(
+                target: "engine::tree::payload_processor",
+                blocked_for=?elapsed,
+                "Waited for execution cache to become available"
+            );
+        }
+        elapsed
+    }
+
     /// Updates the cache with a closure that has exclusive access to the guard.
     /// This ensures that all cache operations happen atomically.
     ///

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

@@ -3,7 +3,7 @@
 use alloy_primitives::B256;
 use parking_lot::Mutex;
 use reth_trie_sparse::SparseStateTrie;
-use std::sync::Arc;
+use std::{sync::Arc, time::Instant};
 use tracing::debug;
 
 /// Type alias for the sparse trie type used in preservation.
@@ -28,6 +28,27 @@ impl SharedPreservedSparseTrie {
     pub(super) fn lock(&self) -> PreservedTrieGuard<'_> {
         PreservedTrieGuard(self.0.lock())
     }
+
+    /// Waits until the sparse trie lock becomes available.
+    ///
+    /// This acquires and immediately releases the lock, ensuring that any
+    /// ongoing operations complete before returning. Useful for synchronization
+    /// before starting payload processing.
+    ///
+    /// Returns the time spent waiting for the lock.
+    pub(super) fn wait_for_availability(&self) -> std::time::Duration {
+        let start = Instant::now();
+        let _guard = self.0.lock();
+        let elapsed = start.elapsed();
+        if elapsed.as_millis() > 5 {
+            debug!(
+                target: "engine::tree::payload_processor",
+                blocked_for=?elapsed,
+                "Waited for preserved sparse trie to become available"
+            );
+        }
+        elapsed
+    }
 }
 
 /// Guard that holds the lock on the preserved trie.

crates/engine/tree/src/tree/payload_validator.rs

@@ -4,11 +4,11 @@ use crate::tree::{
     cached_state::CachedStateProvider,
     error::{InsertBlockError, InsertBlockErrorKind, InsertPayloadError},
     instrumented_state::InstrumentedStateProvider,
-    payload_processor::PayloadProcessor,
+    payload_processor::{CacheWaitDurations, PayloadProcessor},
     precompile_cache::{CachedPrecompile, CachedPrecompileMetrics, PrecompileCacheMap},
     sparse_trie::StateRootComputeOutcome,
     EngineApiMetrics, EngineApiTreeState, ExecutionEnv, PayloadHandle, StateProviderBuilder,
-    StateProviderDatabase, TreeConfig,
+    StateProviderDatabase, TreeConfig, WaitForCaches,
 };
 use alloy_consensus::transaction::{Either, TxHashRef};
 use alloy_eip7928::BlockAccessList;
@@ -1585,6 +1585,15 @@ where
     }
 }
 
+impl<P, Evm, V> WaitForCaches for BasicEngineValidator<P, Evm, V>
+where
+    Evm: ConfigureEvm,
+{
+    fn wait_for_caches(&self) -> CacheWaitDurations {
+        self.payload_processor.wait_for_caches()
+    }
+}
+
 /// Enum representing either block or payload being validated.
 #[derive(Debug)]
 pub enum BlockOrPayload<T: PayloadTypes> {