feat(trie): add blinded node metrics in ProofTaskManager (#17685)

crates/trie/parallel/src/lib.rs

@@ -25,3 +25,7 @@ pub mod proof_task;
 /// Parallel state root metrics.
 #[cfg(feature = "metrics")]
 pub mod metrics;
+
+/// Proof task manager metrics.
+#[cfg(feature = "metrics")]
+pub mod proof_task_metrics;

crates/trie/parallel/src/proof_task.rs

@@ -39,6 +39,9 @@ use std::{
 use tokio::runtime::Handle;
 use tracing::debug;
 
+#[cfg(feature = "metrics")]
+use crate::proof_task_metrics::ProofTaskMetrics;
+
 type StorageProofResult = Result<DecodedStorageMultiProof, ParallelStateRootError>;
 type TrieNodeProviderResult = Result<Option<RevealedNode>, SparseTrieError>;
 
@@ -70,6 +73,9 @@ pub struct ProofTaskManager<Factory: DatabaseProviderFactory> {
     /// Incremented in [`ProofTaskManagerHandle::new`] and decremented in
     /// [`ProofTaskManagerHandle::drop`].
     active_handles: Arc<AtomicUsize>,
+    /// Metrics tracking blinded node fetches.
+    #[cfg(feature = "metrics")]
+    metrics: ProofTaskMetrics,
 }
 
 impl<Factory: DatabaseProviderFactory> ProofTaskManager<Factory> {
@@ -95,6 +101,8 @@ impl<Factory: DatabaseProviderFactory> ProofTaskManager<Factory> {
             proof_task_rx,
             tx_sender,
             active_handles: Arc::new(AtomicUsize::new(0)),
+            #[cfg(feature = "metrics")]
+            metrics: ProofTaskMetrics::default(),
         }
     }
 
@@ -167,6 +175,17 @@ where
             match self.proof_task_rx.recv() {
                 Ok(message) => match message {
                     ProofTaskMessage::QueueTask(task) => {
+                        // Track metrics for blinded node requests
+                        #[cfg(feature = "metrics")]
+                        match &task {
+                            ProofTaskKind::BlindedAccountNode(_, _) => {
+                                self.metrics.account_nodes += 1;
+                            }
+                            ProofTaskKind::BlindedStorageNode(_, _, _) => {
+                                self.metrics.storage_nodes += 1;
+                            }
+                            _ => {}
+                        }
                         // queue the task
                         self.queue_proof_task(task)
                     }
@@ -174,7 +193,12 @@ where
                         // return the transaction to the pool
                         self.proof_task_txs.push(tx);
                     }
-                    ProofTaskMessage::Terminate => return Ok(()),
+                    ProofTaskMessage::Terminate => {
+                        // Record metrics before terminating
+                        #[cfg(feature = "metrics")]
+                        self.metrics.record();
+                        return Ok(())
+                    }
                 },
                 // All senders are disconnected, so we can terminate
                 // However this should never happen, as this struct stores a sender

crates/trie/parallel/src/proof_task_metrics.rs

@@ -0,0 +1,42 @@
+use reth_metrics::{metrics::Histogram, Metrics};
+
+/// Metrics for blinded node fetching for the duration of the proof task manager.
+#[derive(Clone, Debug, Default)]
+pub struct ProofTaskMetrics {
+    /// The actual metrics for blinded nodes.
+    pub task_metrics: ProofTaskTrieMetrics,
+    /// Count of blinded account node requests.
+    pub account_nodes: usize,
+    /// Count of blinded storage node requests.
+    pub storage_nodes: usize,
+}
+
+impl ProofTaskMetrics {
+    /// Record the blinded node counts into the histograms.
+    pub fn record(&self) {
+        self.task_metrics.record_account_nodes(self.account_nodes);
+        self.task_metrics.record_storage_nodes(self.storage_nodes);
+    }
+}
+
+/// Metrics for the proof task.
+#[derive(Clone, Metrics)]
+#[metrics(scope = "trie.proof_task")]
+pub struct ProofTaskTrieMetrics {
+    /// A histogram for the number of blinded account nodes fetched.
+    blinded_account_nodes: Histogram,
+    /// A histogram for the number of blinded storage nodes fetched.
+    blinded_storage_nodes: Histogram,
+}
+
+impl ProofTaskTrieMetrics {
+    /// Record account nodes fetched.
+    pub fn record_account_nodes(&self, count: usize) {
+        self.blinded_account_nodes.record(count as f64);
+    }
+
+    /// Record storage nodes fetched.
+    pub fn record_storage_nodes(&self, count: usize) {
+        self.blinded_storage_nodes.record(count as f64);
+    }
+}