feat: add capacity metrics for tries (#19117)

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

@@ -172,4 +172,18 @@ impl SparseTrieInterface for ConfiguredSparseTrie {
             Self::Parallel(trie) => trie.updates_ref(),
         }
     }
+
+    fn node_capacity(&self) -> usize {
+        match self {
+            Self::Serial(trie) => trie.node_capacity(),
+            Self::Parallel(trie) => trie.node_capacity(),
+        }
+    }
+
+    fn value_capacity(&self) -> usize {
+        match self {
+            Self::Serial(trie) => trie.value_capacity(),
+            Self::Parallel(trie) => trie.value_capacity(),
+        }
+    }
 }

crates/trie/sparse-parallel/src/lower.rs

@@ -106,4 +106,20 @@ impl LowerSparseSubtrie {
             Self::Revealed(_) | Self::Blind(_) => None,
         }
     }
+
+    /// Returns the capacity of any maps containing trie nodes
+    pub(crate) fn node_capacity(&self) -> usize {
+        match self {
+            Self::Revealed(trie) | Self::Blind(Some(trie)) => trie.node_capacity(),
+            Self::Blind(None) => 0,
+        }
+    }
+
+    /// Returns the capacity of any maps containing trie values
+    pub(crate) fn value_capacity(&self) -> usize {
+        match self {
+            Self::Revealed(trie) | Self::Blind(Some(trie)) => trie.value_capacity(),
+            Self::Blind(None) => 0,
+        }
+    }
 }

crates/trie/sparse-parallel/src/trie.rs

@@ -873,6 +873,16 @@ impl SparseTrieInterface for ParallelSparseTrie {
             }
         }
     }
+
+    fn node_capacity(&self) -> usize {
+        self.upper_subtrie.node_capacity() +
+            self.lower_subtries.iter().map(|trie| trie.node_capacity()).sum::<usize>()
+    }
+
+    fn value_capacity(&self) -> usize {
+        self.upper_subtrie.value_capacity() +
+            self.lower_subtries.iter().map(|trie| trie.value_capacity()).sum::<usize>()
+    }
 }
 
 impl ParallelSparseTrie {
@@ -2091,6 +2101,16 @@ impl SparseSubtrie {
         self.nodes.clear();
         self.inner.clear();
     }
+
+    /// Returns the capacity of the map containing trie nodes.
+    pub(crate) fn node_capacity(&self) -> usize {
+        self.nodes.capacity()
+    }
+
+    /// Returns the capacity of the map containing trie values.
+    pub(crate) fn value_capacity(&self) -> usize {
+        self.inner.value_capacity()
+    }
 }
 
 /// Helper type for [`SparseSubtrie`] to mutably access only a subset of fields from the original
@@ -2424,6 +2444,11 @@ impl SparseSubtrieInner {
         self.values.clear();
         self.buffers.clear();
     }
+
+    /// Returns the capacity of the map storing leaf values
+    fn value_capacity(&self) -> usize {
+        self.values.capacity()
+    }
 }
 
 /// Represents the outcome of processing a node during leaf insertion

crates/trie/sparse/src/metrics.rs

@@ -1,5 +1,6 @@
 //! Metrics for the sparse state trie
 
+use metrics::Gauge;
 use reth_metrics::{metrics::Histogram, Metrics};
 
 /// Metrics for the sparse state trie
@@ -15,24 +16,24 @@ pub(crate) struct SparseStateTrieMetrics {
     pub(crate) multiproof_skipped_storage_nodes: u64,
     /// Number of total storage nodes, including those that were skipped.
     pub(crate) multiproof_total_storage_nodes: u64,
-    /// The actual metrics we will record into the histogram
-    pub(crate) histograms: SparseStateTrieHistograms,
+    /// The actual metrics we will record
+    pub(crate) inner_metrics: SparseStateTrieInnerMetrics,
 }
 
 impl SparseStateTrieMetrics {
     /// Record the metrics into the histograms
     pub(crate) fn record(&mut self) {
         use core::mem::take;
-        self.histograms
+        self.inner_metrics
             .multiproof_skipped_account_nodes
             .record(take(&mut self.multiproof_skipped_account_nodes) as f64);
-        self.histograms
+        self.inner_metrics
             .multiproof_total_account_nodes
             .record(take(&mut self.multiproof_total_account_nodes) as f64);
-        self.histograms
+        self.inner_metrics
             .multiproof_skipped_storage_nodes
             .record(take(&mut self.multiproof_skipped_storage_nodes) as f64);
-        self.histograms
+        self.inner_metrics
             .multiproof_total_storage_nodes
             .record(take(&mut self.multiproof_total_storage_nodes) as f64);
     }
@@ -56,12 +57,28 @@ impl SparseStateTrieMetrics {
     pub(crate) const fn increment_total_storage_nodes(&mut self, count: u64) {
         self.multiproof_total_storage_nodes += count;
     }
+
+    /// Set the value capacity for the sparse state trie
+    pub(crate) fn set_value_capacity(&self, capacity: usize) {
+        self.inner_metrics.value_capacity.set(capacity as f64);
+    }
+
+    /// Set the node capacity for the sparse state trie
+    pub(crate) fn set_node_capacity(&self, capacity: usize) {
+        self.inner_metrics.node_capacity.set(capacity as f64);
+    }
+
+    /// Set the number of cleared and active storage tries
+    pub(crate) fn set_storage_trie_metrics(&self, cleared: usize, active: usize) {
+        self.inner_metrics.cleared_storage_tries.set(cleared as f64);
+        self.inner_metrics.active_storage_tries.set(active as f64);
+    }
 }
 
 /// Metrics for the sparse state trie
 #[derive(Metrics)]
 #[metrics(scope = "sparse_state_trie")]
-pub(crate) struct SparseStateTrieHistograms {
+pub(crate) struct SparseStateTrieInnerMetrics {
     /// Histogram of account nodes that were skipped during a multiproof reveal due to being
     /// redundant (i.e. they were already revealed)
     pub(crate) multiproof_skipped_account_nodes: Histogram,
@@ -72,4 +89,12 @@ pub(crate) struct SparseStateTrieHistograms {
     pub(crate) multiproof_skipped_storage_nodes: Histogram,
     /// Histogram of total storage nodes, including those that were skipped.
     pub(crate) multiproof_total_storage_nodes: Histogram,
+    /// Gauge for the trie's node capacity
+    pub(crate) node_capacity: Gauge,
+    /// Gauge for the trie's value capacity
+    pub(crate) value_capacity: Gauge,
+    /// The current number of cleared storage tries.
+    pub(crate) cleared_storage_tries: Gauge,
+    /// The number of currently active storage tries, i.e., not cleared
+    pub(crate) active_storage_tries: Gauge,
 }

crates/trie/sparse/src/state.rs

@@ -585,9 +585,17 @@ where
         &mut self,
         provider_factory: impl TrieNodeProviderFactory,
     ) -> SparseStateTrieResult<B256> {
-        // record revealed node metrics
+        // record revealed node metrics and capacity metrics
         #[cfg(feature = "metrics")]
-        self.metrics.record();
+        {
+            self.metrics.record();
+            self.metrics.set_node_capacity(self.node_capacity());
+            self.metrics.set_value_capacity(self.value_capacity());
+            self.metrics.set_storage_trie_metrics(
+                self.storage.cleared_tries.len(),
+                self.storage.tries.len(),
+            );
+        }
 
         Ok(self.revealed_trie_mut(provider_factory)?.root())
     }
@@ -598,9 +606,17 @@ where
         &mut self,
         provider_factory: impl TrieNodeProviderFactory,
     ) -> SparseStateTrieResult<(B256, TrieUpdates)> {
-        // record revealed node metrics
+        // record revealed node metrics and capacity metrics
         #[cfg(feature = "metrics")]
-        self.metrics.record();
+        {
+            self.metrics.record();
+            self.metrics.set_node_capacity(self.node_capacity());
+            self.metrics.set_value_capacity(self.value_capacity());
+            self.metrics.set_storage_trie_metrics(
+                self.storage.cleared_tries.len(),
+                self.storage.tries.len(),
+            );
+        }
 
         let storage_tries = self.storage_trie_updates();
         let revealed = self.revealed_trie_mut(provider_factory)?;
@@ -805,6 +821,16 @@ where
         storage_trie.remove_leaf(slot, provider)?;
         Ok(())
     }
+
+    /// The sum of the account trie's node capacity and the storage tries' node capacity
+    pub fn node_capacity(&self) -> usize {
+        self.state.node_capacity() + self.storage.total_node_capacity()
+    }
+
+    /// The sum of the account trie's value capacity and the storage tries' value capacity
+    pub fn value_capacity(&self) -> usize {
+        self.state.value_capacity() + self.storage.total_value_capacity()
+    }
 }
 
 /// The fields of [`SparseStateTrie`] related to storage tries. This is kept separate from the rest
@@ -880,6 +906,46 @@ impl<S: SparseTrieInterface + Clone> StorageTries<S> {
             .remove(account)
             .unwrap_or_else(|| self.cleared_revealed_paths.pop().unwrap_or_default())
     }
+
+    /// Sums the total node capacity in `cleared_tries`
+    fn total_cleared_tries_node_capacity(&self) -> usize {
+        self.cleared_tries.iter().map(|trie| trie.node_capacity()).sum()
+    }
+
+    /// Sums the total value capacity in `cleared_tries`
+    fn total_cleared_tries_value_capacity(&self) -> usize {
+        self.cleared_tries.iter().map(|trie| trie.value_capacity()).sum()
+    }
+
+    /// Calculates the sum of the active storage trie node capacity, ie the tries in `tries`
+    fn total_active_tries_node_capacity(&self) -> usize {
+        self.tries.values().map(|trie| trie.node_capacity()).sum()
+    }
+
+    /// Calculates the sum of the active storage trie value capacity, ie the tries in `tries`
+    fn total_active_tries_value_capacity(&self) -> usize {
+        self.tries.values().map(|trie| trie.value_capacity()).sum()
+    }
+
+    /// Calculates the sum of active and cleared storage trie node capacity, i.e. the sum of
+    /// * [`StorageTries::total_active_tries_node_capacity`], and
+    /// * [`StorageTries::total_cleared_tries_node_capacity`]
+    /// * the default trie's node capacity
+    fn total_node_capacity(&self) -> usize {
+        self.total_active_tries_node_capacity() +
+            self.total_cleared_tries_node_capacity() +
+            self.default_trie.node_capacity()
+    }
+
+    /// Calculates the sum of active and cleared storage trie value capacity, i.e. the sum of
+    /// * [`StorageTries::total_active_tries_value_capacity`], and
+    /// * [`StorageTries::total_cleared_tries_value_capacity`], and
+    /// * the default trie's value capacity
+    fn total_value_capacity(&self) -> usize {
+        self.total_active_tries_value_capacity() +
+            self.total_cleared_tries_value_capacity() +
+            self.default_trie.value_capacity()
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Default)]

crates/trie/sparse/src/traits.rs

@@ -222,6 +222,12 @@ pub trait SparseTrieInterface: Sized + Debug + Send + Sync {
     ///
     /// This is useful for reusing the trie without needing to reallocate memory.
     fn clear(&mut self);
+
+    /// This returns the capacity of any inner data structures which store nodes.
+    fn node_capacity(&self) -> usize;
+
+    /// This returns the capacity of any inner data structures which store leaf values.
+    fn value_capacity(&self) -> usize;
 }
 
 /// Struct for passing around branch node mask information.

crates/trie/sparse/src/trie.rs

@@ -259,6 +259,22 @@ impl<T: SparseTrieInterface> SparseTrie<T> {
         revealed.remove_leaf(path, provider)?;
         Ok(())
     }
+
+    /// Returns the allocated capacity for sparse trie nodes.
+    pub fn node_capacity(&self) -> usize {
+        match self {
+            Self::Blind(Some(trie)) | Self::Revealed(trie) => trie.node_capacity(),
+            _ => 0,
+        }
+    }
+
+    /// Returns the allocated capacity for sparse trie values.
+    pub fn value_capacity(&self) -> usize {
+        match self {
+            Self::Blind(Some(trie)) | Self::Revealed(trie) => trie.value_capacity(),
+            _ => 0,
+        }
+    }
 }
 
 /// The representation of revealed sparse trie.
@@ -1064,6 +1080,14 @@ impl SparseTrieInterface for SerialSparseTrie {
         // If we get here, there's no leaf at the target path
         Ok(LeafLookup::NonExistent)
     }
+
+    fn node_capacity(&self) -> usize {
+        self.nodes.capacity()
+    }
+
+    fn value_capacity(&self) -> usize {
+        self.values.capacity()
+    }
 }
 
 impl SerialSparseTrie {