From a3cf4fb3aacda94f73fbd0948f9e99475462e53d Mon Sep 17 00:00:00 2001
From: Ivan Kalinin <i.kalinin@dexpa.io>
Date: Tue, 11 Jun 2024 02:08:54 +0200
Subject: [PATCH] feat(collator): add `handle_state` histograms

---
 collator/src/state_node.rs | 14 ++++++++++++++
 scripts/gen-dashboard.py   | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/collator/src/state_node.rs b/collator/src/state_node.rs
index 20f2b0077..658d47e88 100644
--- a/collator/src/state_node.rs
+++ b/collator/src/state_node.rs
@@ -1,5 +1,6 @@
 use std::collections::{BTreeMap, HashMap};
 use std::sync::Arc;
+use std::time::Instant;
 
 use anyhow::{Context, Result};
 use async_trait::async_trait;
@@ -8,6 +9,7 @@ use tokio::sync::{broadcast, Mutex};
 use tycho_block_util::block::{BlockStuff, BlockStuffAug};
 use tycho_block_util::state::ShardStateStuff;
 use tycho_storage::{BlockHandle, Storage};
+use tycho_util::metrics::HistogramGuard;
 
 use crate::tracing_targets;
 use crate::types::BlockStuffForSync;
@@ -147,6 +149,8 @@ impl StateNodeAdapter for StateNodeAdapterStdImpl {
     }
 
     async fn handle_state(&self, state: &ShardStateStuff) -> Result<()> {
+        let _histogram = HistogramGuard::begin("tycho_collator_adapter_handle_state_time");
+
         tracing::debug!(target: tracing_targets::STATE_NODE_ADAPTER, "Handle block: {}", state.block_id().as_short_id());
         let block_id = *state.block_id();
 
@@ -176,15 +180,25 @@ impl StateNodeAdapter for StateNodeAdapterStdImpl {
 
                 match block {
                     None => {
+                        let _histogram = HistogramGuard::begin(
+                            "tycho_collator_adapter_on_block_accepted_ext_time",
+                        );
+
                         tracing::info!(target: tracing_targets::STATE_NODE_ADAPTER, "Block handled external: {:?}", block_id);
                         self.listener.on_block_accepted_external(state).await?;
                     }
                     Some(block) => {
+                        let _histogram =
+                            HistogramGuard::begin("tycho_collator_adapter_on_block_accepted_time");
+
                         tracing::info!(target: tracing_targets::STATE_NODE_ADAPTER, "Block handled: {:?}", block_id);
                         self.listener.on_block_accepted(&block.block_id).await?;
                     }
                 }
             } else {
+                let _histogram =
+                    HistogramGuard::begin("tycho_collator_adapter_on_block_accepted_alt_ext_time");
+
                 tracing::info!(target: tracing_targets::STATE_NODE_ADAPTER, "Block handled external. Shard ID not found in blocks buffer: {:?}", block_id);
                 self.listener.on_block_accepted_external(state).await?;
             }
diff --git a/scripts/gen-dashboard.py b/scripts/gen-dashboard.py
index 69015c90f..f87908671 100644
--- a/scripts/gen-dashboard.py
+++ b/scripts/gen-dashboard.py
@@ -356,6 +356,20 @@ def collator_do_collate() -> RowPanel:
         create_heatmap_panel("tycho_do_collate_execute_time", "Execution time"),
         create_heatmap_panel("tycho_do_collate_build_block_time", "Build block time"),
         create_heatmap_panel("tycho_do_collate_update_state_time", "Update state time"),
+        create_heatmap_panel(
+            "tycho_collator_adapter_handle_state_time", "Handle state by collator"
+        ),
+        create_heatmap_panel(
+            "tycho_collator_adapter_on_block_accepted_ext_time",
+            "on_block_accepted_external with blocks guard",
+        ),
+        create_heatmap_panel(
+            "tycho_collator_adapter_on_block_accepted_alt_ext_time",
+            "on_block_accepted_external without blocks guard",
+        ),
+        create_heatmap_panel(
+            "tycho_collator_adapter_on_block_accepted_time", "on_block_accepted"
+        ),
     ]
     return create_row("Collator Do Collate", metrics)