feat(trie): node iter metrics (#15438)

Author: shekhirin

Diff for: crates/trie/db/src/storage.rs

@@ -9,7 +9,7 @@ use reth_trie::{
 };
 
 #[cfg(feature = "metrics")]
-use reth_trie::metrics::{TrieRootMetrics, TrieType};
+use reth_trie::metrics::TrieRootMetrics;
 
 /// Extends [`StorageRoot`] with operations specific for working with a database transaction.
 pub trait DatabaseStorageRoot<'a, TX> {
@@ -44,7 +44,7 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             address,
             Default::default(),
             #[cfg(feature = "metrics")]
-            TrieRootMetrics::new(TrieType::Storage),
+            TrieRootMetrics::new(reth_trie::TrieType::Storage),
         )
     }
 
@@ -55,7 +55,7 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             hashed_address,
             Default::default(),
             #[cfg(feature = "metrics")]
-            TrieRootMetrics::new(TrieType::Storage),
+            TrieRootMetrics::new(reth_trie::TrieType::Storage),
         )
     }
 
@@ -73,7 +73,7 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             address,
             prefix_set,
             #[cfg(feature = "metrics")]
-            TrieRootMetrics::new(TrieType::Storage),
+            TrieRootMetrics::new(reth_trie::TrieType::Storage),
         )
         .root()
     }

Diff for: crates/trie/parallel/src/metrics.rs

@@ -1,7 +1,7 @@
 use crate::stats::ParallelTrieStats;
 use metrics::Histogram;
 use reth_metrics::Metrics;
-use reth_trie::metrics::{TrieRootMetrics, TrieType};
+use reth_trie::{metrics::TrieRootMetrics, TrieType};
 
 /// Parallel state root metrics.
 #[derive(Debug)]

Diff for: crates/trie/parallel/src/proof.rs

@@ -26,7 +26,7 @@ use reth_trie::{
     updates::TrieUpdatesSorted,
     walker::TrieWalker,
     HashBuilder, HashedPostStateSorted, MultiProof, MultiProofTargets, Nibbles, StorageMultiProof,
-    TRIE_ACCOUNT_RLP_MAX_SIZE,
+    TrieType, TRIE_ACCOUNT_RLP_MAX_SIZE,
 };
 use reth_trie_common::proof::ProofRetainer;
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
@@ -229,6 +229,7 @@ where
         let mut account_node_iter = TrieNodeIter::new(
             walker,
             hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,
+            TrieType::State,
         );
         while let Some(account_node) =
             account_node_iter.try_next().map_err(ProviderError::Database)?

Diff for: crates/trie/parallel/src/root.rs

@@ -16,7 +16,7 @@ use reth_trie::{
     trie_cursor::{InMemoryTrieCursorFactory, TrieCursorFactory},
     updates::TrieUpdates,
     walker::TrieWalker,
-    HashBuilder, Nibbles, StorageRoot, TrieInput, TRIE_ACCOUNT_RLP_MAX_SIZE,
+    HashBuilder, Nibbles, StorageRoot, TrieInput, TrieType, TRIE_ACCOUNT_RLP_MAX_SIZE,
 };
 use reth_trie_db::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
 use std::{collections::HashMap, sync::Arc};
@@ -153,6 +153,7 @@ where
         let mut account_node_iter = TrieNodeIter::new(
             walker,
             hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,
+            TrieType::State,
         );
 
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);

Diff for: crates/trie/sparse/benches/root.rs

@@ -10,7 +10,7 @@ use reth_trie::{
     trie_cursor::{noop::NoopStorageTrieCursor, InMemoryStorageTrieCursor},
     updates::StorageTrieUpdates,
     walker::TrieWalker,
-    HashedStorage,
+    HashedStorage, TrieType,
 };
 use reth_trie_common::{HashBuilder, Nibbles};
 use reth_trie_sparse::SparseTrie;
@@ -143,6 +143,7 @@ fn calculate_root_from_leaves_repeated(c: &mut Criterion) {
                                         NoopHashedStorageCursor::default(),
                                         Some(&storage_sorted),
                                     ),
+                                    TrieType::Storage,
                                 );
 
                                 let mut hb = HashBuilder::default().with_updates(true);

Diff for: crates/trie/sparse/src/trie.rs

@@ -1591,7 +1591,7 @@ mod tests {
         node_iter::{TrieElement, TrieNodeIter},
         trie_cursor::{noop::NoopAccountTrieCursor, TrieCursor, TrieCursorFactory},
         walker::TrieWalker,
-        BranchNode, ExtensionNode, HashedPostState, LeafNode,
+        BranchNode, ExtensionNode, HashedPostState, LeafNode, TrieType,
     };
     use reth_trie_common::{
         proof::{ProofNodes, ProofRetainer},
@@ -1648,6 +1648,7 @@ mod tests {
                 NoopHashedAccountCursor::default(),
                 hashed_post_state.accounts(),
             ),
+            TrieType::State,
         );
 
         while let Some(node) = node_iter.try_next().unwrap() {

Diff for: crates/trie/trie/src/lib.rs

@@ -37,7 +37,7 @@ pub mod witness;
 
 /// The implementation of the Merkle Patricia Trie.
 mod trie;
-pub use trie::{StateRoot, StorageRoot};
+pub use trie::{StateRoot, StorageRoot, TrieType};
 
 /// Utilities for state root checkpoint progress.
 mod progress;

Diff for: crates/trie/trie/src/metrics.rs

@@ -1,4 +1,4 @@
-use crate::stats::TrieStats;
+use crate::{stats::TrieStats, trie::TrieType};
 use metrics::{Counter, Histogram};
 use reth_metrics::Metrics;
 
@@ -46,25 +46,7 @@ impl TrieRootMetrics {
     }
 }
 
-/// Trie type for differentiating between various trie calculations.
-#[derive(Clone, Copy, Debug)]
-pub enum TrieType {
-    /// State trie type.
-    State,
-    /// Storage trie type.
-    Storage,
-}
-
-impl TrieType {
-    pub(crate) const fn as_str(&self) -> &'static str {
-        match self {
-            Self::State => "state",
-            Self::Storage => "storage",
-        }
-    }
-}
-
-/// Metrics for trie walker
+/// Metrics for [`crate::walker::TrieWalker`].
 #[derive(Clone, Metrics)]
 #[metrics(scope = "trie.walker")]
 pub struct WalkerMetrics {
@@ -83,3 +65,58 @@ impl WalkerMetrics {
         self.out_of_order_subnode.increment(amount);
     }
 }
+
+/// Metrics for [`crate::node_iter::TrieNodeIter`].
+#[derive(Clone, Metrics)]
+#[metrics(scope = "trie.node_iter")]
+pub struct TrieNodeIterMetrics {
+    /// The number of branch nodes returned by the iterator.
+    branch_nodes_returned_total: Counter,
+    /// The number of times same leaf node was seeked multiple times in a row by the iterator.
+    leaf_nodes_same_seeked_total: Counter,
+    /// The number of times the same leaf node as we just advanced to was seeked by the iterator.
+    leaf_nodes_same_seeked_as_advanced_total: Counter,
+    /// The number of leaf nodes seeked by the iterator.
+    leaf_nodes_seeked_total: Counter,
+    /// The number of leaf nodes advanced by the iterator.
+    leaf_nodes_advanced_total: Counter,
+    /// The number of leaf nodes returned by the iterator.
+    leaf_nodes_returned_total: Counter,
+}
+
+impl TrieNodeIterMetrics {
+    /// Create new metrics for the given trie type.
+    pub fn new(ty: TrieType) -> Self {
+        Self::new_with_labels(&[("type", ty.as_str())])
+    }
+
+    /// Increment `branch_nodes_returned_total`.
+    pub fn inc_branch_nodes_returned(&self) {
+        self.branch_nodes_returned_total.increment(1);
+    }
+
+    /// Increment `leaf_nodes_same_seeked_total`.
+    pub fn inc_leaf_nodes_same_seeked(&self) {
+        self.leaf_nodes_same_seeked_total.increment(1);
+    }
+
+    /// Increment `leaf_nodes_same_seeked_as_advanced_total`.
+    pub fn inc_leaf_nodes_same_seeked_as_advanced(&self) {
+        self.leaf_nodes_same_seeked_as_advanced_total.increment(1);
+    }
+
+    /// Increment `leaf_nodes_seeked_total`.
+    pub fn inc_leaf_nodes_seeked(&self) {
+        self.leaf_nodes_seeked_total.increment(1);
+    }
+
+    /// Increment `leaf_nodes_advanced_total`.
+    pub fn inc_leaf_nodes_advanced(&self) {
+        self.leaf_nodes_advanced_total.increment(1);
+    }
+
+    /// Increment `leaf_nodes_returned_total`.
+    pub fn inc_leaf_nodes_returned(&self) {
+        self.leaf_nodes_returned_total.increment(1);
+    }
+}

Diff for: crates/trie/trie/src/node_iter.rs

@@ -1,4 +1,7 @@
-use crate::{hashed_cursor::HashedCursor, trie_cursor::TrieCursor, walker::TrieWalker, Nibbles};
+use crate::{
+    hashed_cursor::HashedCursor, trie::TrieType, trie_cursor::TrieCursor, walker::TrieWalker,
+    Nibbles,
+};
 use alloy_primitives::B256;
 use reth_storage_errors::db::DatabaseError;
 use tracing::trace;
@@ -45,17 +48,30 @@ pub struct TrieNodeIter<C, H: HashedCursor> {
     current_hashed_entry: Option<(B256, <H as HashedCursor>::Value)>,
     /// Flag indicating whether we should check the current walker key.
     should_check_walker_key: bool,
+
+    #[cfg(feature = "metrics")]
+    metrics: crate::metrics::TrieNodeIterMetrics,
+    #[cfg(feature = "metrics")]
+    previously_seeked_key: Option<B256>,
+    #[cfg(feature = "metrics")]
+    previously_advanced_to_key: Option<B256>,
 }
 
 impl<C, H: HashedCursor> TrieNodeIter<C, H> {
     /// Creates a new [`TrieNodeIter`].
-    pub const fn new(walker: TrieWalker<C>, hashed_cursor: H) -> Self {
+    pub fn new(walker: TrieWalker<C>, hashed_cursor: H, trie_type: TrieType) -> Self {
         Self {
             walker,
             hashed_cursor,
             previous_hashed_key: None,
             current_hashed_entry: None,
             should_check_walker_key: false,
+            #[cfg(feature = "metrics")]
+            metrics: crate::metrics::TrieNodeIterMetrics::new(trie_type),
+            #[cfg(feature = "metrics")]
+            previously_seeked_key: None,
+            #[cfg(feature = "metrics")]
+            previously_advanced_to_key: None,
         }
     }
 
@@ -65,6 +81,41 @@ impl<C, H: HashedCursor> TrieNodeIter<C, H> {
         self.previous_hashed_key = Some(previous_hashed_key);
         self
     }
+
+    /// Seeks the hashed cursor to the given key.
+    ///
+    /// If `metrics` feature is enabled, also updates the metrics.
+    fn hashed_cursor_seek(&mut self, key: B256) -> Result<Option<(B256, H::Value)>, DatabaseError> {
+        #[cfg(feature = "metrics")]
+        {
+            self.metrics.inc_leaf_nodes_seeked();
+
+            if Some(key) == self.previously_seeked_key {
+                self.metrics.inc_leaf_nodes_same_seeked();
+            }
+            self.previously_seeked_key = Some(key);
+
+            if Some(key) == self.previously_advanced_to_key {
+                self.metrics.inc_leaf_nodes_same_seeked_as_advanced();
+            }
+        }
+        self.hashed_cursor.seek(key)
+    }
+
+    /// Advances the hashed cursor to the next entry.
+    ///
+    /// If `metrics` feature is enabled, also updates the metrics.
+    fn hashed_cursor_next(&mut self) -> Result<Option<(B256, H::Value)>, DatabaseError> {
+        let result = self.hashed_cursor.next();
+        #[cfg(feature = "metrics")]
+        {
+            self.metrics.inc_leaf_nodes_advanced();
+
+            self.previously_advanced_to_key =
+                result.as_ref().ok().and_then(|result| result.as_ref().map(|(k, _)| *k));
+        }
+        result
+    }
 }
 
 impl<C, H> TrieNodeIter<C, H>
@@ -97,6 +148,8 @@ where
                     self.should_check_walker_key = true;
                     // If it's possible to skip the current node in the walker, return a branch node
                     if self.walker.can_skip_current_node {
+                        #[cfg(feature = "metrics")]
+                        self.metrics.inc_branch_nodes_returned();
                         return Ok(Some(TrieElement::Branch(TrieBranchNode::new(
                             key.clone(),
                             self.walker.hash().unwrap(),
@@ -116,7 +169,10 @@ where
 
                 // Set the next hashed entry as a leaf node and return
                 trace!(target: "trie::node_iter", ?hashed_key, "next hashed entry");
-                self.current_hashed_entry = self.hashed_cursor.next()?;
+                self.current_hashed_entry = self.hashed_cursor_next()?;
+
+                #[cfg(feature = "metrics")]
+                self.metrics.inc_leaf_nodes_returned();
                 return Ok(Some(TrieElement::Leaf(hashed_key, value)))
             }
 
@@ -125,8 +181,8 @@ where
                 Some(hashed_key) => {
                     trace!(target: "trie::node_iter", ?hashed_key, "seeking to the previous hashed entry");
                     // Seek to the previous hashed key and get the next hashed entry
-                    self.hashed_cursor.seek(hashed_key)?;
-                    self.current_hashed_entry = self.hashed_cursor.next()?;
+                    self.hashed_cursor_seek(hashed_key)?;
+                    self.current_hashed_entry = self.hashed_cursor_next()?;
                 }
                 None => {
                     // Get the seek key and set the current hashed entry based on walker's next
@@ -176,7 +232,7 @@ where
                         continue
                     }
 
-                    self.current_hashed_entry = self.hashed_cursor.seek(seek_key)?;
+                    self.current_hashed_entry = self.hashed_cursor_seek(seek_key)?;
                 }
             }
         }
@@ -209,6 +265,7 @@ mod tests {
             HashedPostStateAccountCursor,
         },
         mock::{KeyVisit, KeyVisitType},
+        trie::TrieType,
         trie_cursor::{
             mock::MockTrieCursorFactory, noop::NoopAccountTrieCursor, TrieCursorFactory,
         },
@@ -240,6 +297,7 @@ mod tests {
                 NoopHashedAccountCursor::default(),
                 hashed_post_state.accounts(),
             ),
+            TrieType::State,
         );
 
         while let Some(node) = node_iter.try_next().unwrap() {
@@ -369,8 +427,11 @@ mod tests {
             B256Map::default(),
         );
 
-        let mut iter =
-            TrieNodeIter::new(walker, hashed_cursor_factory.hashed_account_cursor().unwrap());
+        let mut iter = TrieNodeIter::new(
+            walker,
+            hashed_cursor_factory.hashed_account_cursor().unwrap(),
+            TrieType::State,
+        );
 
         // Walk the iterator until it's exhausted.
         while iter.try_next().unwrap().is_some() {}

Diff for: crates/trie/trie/src/proof/mod.rs

@@ -2,6 +2,7 @@ use crate::{
     hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
     node_iter::{TrieElement, TrieNodeIter},
     prefix_set::{PrefixSetMut, TriePrefixSetsMut},
+    trie::TrieType,
     trie_cursor::TrieCursorFactory,
     walker::TrieWalker,
     HashBuilder, Nibbles, TRIE_ACCOUNT_RLP_MAX_SIZE,
@@ -124,7 +125,8 @@ where
         let mut storages: B256Map<_> =
             targets.keys().map(|key| (*key, StorageMultiProof::empty())).collect();
         let mut account_rlp = Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE);
-        let mut account_node_iter = TrieNodeIter::new(walker, hashed_account_cursor);
+        let mut account_node_iter =
+            TrieNodeIter::new(walker, hashed_account_cursor, TrieType::State);
         while let Some(account_node) = account_node_iter.try_next()? {
             match account_node {
                 TrieElement::Branch(node) => {
@@ -288,7 +290,8 @@ where
         let mut hash_builder = HashBuilder::default()
             .with_proof_retainer(retainer)
             .with_updates(self.collect_branch_node_masks);
-        let mut storage_node_iter = TrieNodeIter::new(walker, hashed_storage_cursor);
+        let mut storage_node_iter =
+            TrieNodeIter::new(walker, hashed_storage_cursor, TrieType::Storage);
         while let Some(node) = storage_node_iter.try_next()? {
             match node {
                 TrieElement::Branch(node) => {