feat(trie): metrics (#6943)

crates/trie/Cargo.toml

@@ -30,7 +30,11 @@ thiserror.workspace = true
 derive_more.workspace = true
 auto_impl = "1"
 
-# test-utils
+# `metrics` feature
+reth-metrics = { workspace = true, optional = true }
+metrics = { workspace = true, optional = true }
+
+# `test-utils` feature
 triehash = { version = "0.8", optional = true }
 
 [dev-dependencies]
@@ -52,6 +56,8 @@ similar-asserts.workspace = true
 criterion.workspace = true
 
 [features]
+default = ["metrics"]
+metrics = ["reth-metrics", "dep:metrics"]
 test-utils = ["triehash"]
 
 [[bench]]

crates/trie/src/lib.rs

@@ -50,6 +50,13 @@ pub mod updates;
 mod progress;
 pub use progress::{IntermediateStateRootState, StateRootProgress};
 
+/// Trie calculation stats.
+pub mod stats;
+
+/// Trie calculation metrics.
+#[cfg(feature = "metrics")]
+pub mod metrics;
+
 /// Collection of trie-related test utilities.
 #[cfg(any(test, feature = "test-utils"))]
 pub mod test_utils;

crates/trie/src/metrics.rs

@@ -0,0 +1,65 @@
+use crate::stats::TrieStats;
+use metrics::Histogram;
+use reth_metrics::Metrics;
+
+/// Wrapper for state root metrics.
+#[derive(Debug)]
+pub struct StateRootMetrics {
+    /// State trie metrics.
+    pub state_trie: TrieRootMetrics,
+    /// Storage trie metrics.
+    pub storage_trie: TrieRootMetrics,
+}
+
+impl Default for StateRootMetrics {
+    fn default() -> Self {
+        Self {
+            state_trie: TrieRootMetrics::new(TrieType::State),
+            storage_trie: TrieRootMetrics::new(TrieType::Storage),
+        }
+    }
+}
+
+/// Metrics for trie root calculation.
+#[derive(Clone, Metrics)]
+#[metrics(scope = "trie")]
+pub struct TrieRootMetrics {
+    /// The number of seconds trie root calculation lasted.
+    duration_seconds: Histogram,
+    /// The number of branches added during trie root calculation.
+    branches_added: Histogram,
+    /// The number of leaves added during trie root calculation.
+    leaves_added: Histogram,
+}
+
+impl TrieRootMetrics {
+    /// Create new metrics for the given trie type.
+    pub fn new(ty: TrieType) -> Self {
+        Self::new_with_labels(&[("type", ty.as_str())])
+    }
+
+    /// Record trie stats as metrics.
+    pub fn record(&self, stats: TrieStats) {
+        self.duration_seconds.record(stats.duration().as_secs_f64());
+        self.branches_added.record(stats.branches_added() as f64);
+        self.leaves_added.record(stats.leaves_added() as f64);
+    }
+}
+
+/// 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",
+        }
+    }
+}

crates/trie/src/stats.rs

@@ -0,0 +1,61 @@
+use std::time::{Duration, Instant};
+
+/// Trie stats.
+#[derive(Clone, Copy, Debug)]
+pub struct TrieStats {
+    duration: Duration,
+    branches_added: u64,
+    leaves_added: u64,
+}
+
+impl TrieStats {
+    /// Duration for root calculation.
+    pub fn duration(&self) -> Duration {
+        self.duration
+    }
+
+    /// Number of leaves added to the hash builder during the calculation.
+    pub fn leaves_added(&self) -> u64 {
+        self.leaves_added
+    }
+
+    /// Number of branches added to the hash builder during the calculation.
+    pub fn branches_added(&self) -> u64 {
+        self.branches_added
+    }
+}
+
+/// Trie metrics tracker.
+#[derive(Debug)]
+pub struct TrieTracker {
+    started_at: Instant,
+    branches_added: u64,
+    leaves_added: u64,
+}
+
+impl Default for TrieTracker {
+    fn default() -> Self {
+        Self { started_at: Instant::now(), branches_added: 0, leaves_added: 0 }
+    }
+}
+
+impl TrieTracker {
+    /// Increment the number of branches added to the hash builder during the calculation.
+    pub fn inc_branch(&mut self) {
+        self.branches_added += 1;
+    }
+
+    /// Increment the number of leaves added to the hash builder during the calculation.
+    pub fn inc_leaf(&mut self) {
+        self.leaves_added += 1;
+    }
+
+    /// Called when root calculation is finished to return trie statistics.
+    pub fn finish(self) -> TrieStats {
+        TrieStats {
+            duration: self.started_at.elapsed(),
+            branches_added: self.branches_added,
+            leaves_added: self.leaves_added,
+        }
+    }
+}