From bdb35ae30b6486c5cf8f06ba4aebad248ea8f861 Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>
Date: Thu, 27 Feb 2025 20:44:31 +0000
Subject: [PATCH] fix(tree): state root task duration (#14763)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 crates/engine/tree/src/tree/mod.rs                 | 10 ++++++----
 .../engine/tree/src/tree/payload_processor/mod.rs  |  2 +-
 .../tree/src/tree/payload_processor/multiproof.rs  |  4 +++-
 .../tree/src/tree/payload_processor/sparse_trie.rs | 14 +++++---------
 4 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index 50dd515b83..0bc099870e 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -15,6 +15,7 @@ use alloy_rpc_types_engine::{
     ForkchoiceState, PayloadStatus, PayloadStatusEnum, PayloadValidationError,
 };
 use error::{InsertBlockError, InsertBlockErrorKind, InsertBlockFatalError};
+use payload_processor::sparse_trie::StateRootComputeOutcome;
 use persistence_state::CurrentPersistenceAction;
 use reth_chain_state::{
     CanonicalInMemoryState, ExecutedBlock, ExecutedBlockWithTrieUpdates,
@@ -2386,7 +2387,8 @@ where
             &block,
             Box::new(handle.state_hook()),
         )?;
-        let execution_time = execution_start.elapsed();
+        let execution_finish = Instant::now();
+        let execution_time = execution_finish.duration_since(execution_start);
         trace!(target: "engine::tree", elapsed = ?execution_time, number=?block_num_hash.number, "Executed block");
 
         // after executing the block we can stop executing transactions
@@ -2421,12 +2423,12 @@ where
             // background task or try to compute it in parallel
             if self.config.use_state_root_task() {
                 match handle.state_root() {
-                    Ok(res) => {
+                    Ok(StateRootComputeOutcome { state_root, trie_updates }) => {
                         // we double check the state root here for good measure
                         // TODO: clean this ups
-                        if res.state_root.0 == block.header().state_root() {
+                        if state_root == block.header().state_root() {
                             maybe_state_root =
-                                Some((res.state_root.0, res.state_root.1, res.total_time))
+                                Some((state_root, trie_updates, execution_finish.elapsed()))
                         }
                     }
                     Err(error) => {
diff --git a/crates/engine/tree/src/tree/payload_processor/mod.rs b/crates/engine/tree/src/tree/payload_processor/mod.rs
index b8f7fbd423..2c2bc9edb5 100644
--- a/crates/engine/tree/src/tree/payload_processor/mod.rs
+++ b/crates/engine/tree/src/tree/payload_processor/mod.rs
@@ -36,10 +36,10 @@ use std::{
 };
 
 pub(crate) mod executor;
+pub(crate) mod sparse_trie;
 
 mod multiproof;
 mod prewarm;
-mod sparse_trie;
 
 /// Entrypoint for executing the payload.
 pub(super) struct PayloadProcessor<N, Evm> {
diff --git a/crates/engine/tree/src/tree/payload_processor/multiproof.rs b/crates/engine/tree/src/tree/payload_processor/multiproof.rs
index 6a736d6109..3e60182ffd 100644
--- a/crates/engine/tree/src/tree/payload_processor/multiproof.rs
+++ b/crates/engine/tree/src/tree/payload_processor/multiproof.rs
@@ -31,7 +31,7 @@ use tracing::{debug, error, trace};
 /// A trie update that can be applied to sparse trie alongside the proofs for touched parts of the
 /// state.
 #[derive(Default, Debug)]
-pub(super) struct SparseTrieUpdate {
+pub(crate) struct SparseTrieUpdate {
     /// The state update that was used to calculate the proof
     pub(crate) state: HashedPostState,
     /// The calculated multiproof
@@ -401,6 +401,8 @@ pub(crate) struct MultiProofTaskMetrics {
     pub sparse_trie_update_duration_histogram: Histogram,
     /// Histogram of sparse trie final update durations.
     pub sparse_trie_final_update_duration_histogram: Histogram,
+    /// Histogram of sparse trie total durations.
+    pub sparse_trie_total_duration_histogram: Histogram,
 
     /// Histogram of state updates received.
     pub state_updates_received_histogram: Histogram,
diff --git a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
index d7d3a4e14c..070c075951 100644
--- a/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
+++ b/crates/engine/tree/src/tree/payload_processor/sparse_trie.rs
@@ -99,22 +99,18 @@ where
         })?;
 
         self.metrics.sparse_trie_final_update_duration_histogram.record(start.elapsed());
+        self.metrics.sparse_trie_total_duration_histogram.record(now.elapsed());
 
-        Ok(StateRootComputeOutcome {
-            state_root: (state_root, trie_updates),
-            total_time: now.elapsed(),
-        })
+        Ok(StateRootComputeOutcome { state_root, trie_updates })
     }
 }
 
 /// Outcome of the state root computation, including the state root itself with
-/// the trie updates and the total time spent.
+/// the trie updates.
 #[derive(Debug)]
 pub(crate) struct StateRootComputeOutcome {
-    /// The computed state root and trie updates
-    pub state_root: (B256, TrieUpdates),
-    /// The total time spent calculating the state root
-    pub total_time: Duration,
+    pub state_root: B256,
+    pub trie_updates: TrieUpdates,
 }
 
 /// Aliased for now to not introduce too many changes at once.