sui_security_watchdog/
metrics.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use prometheus::{IntGauge, Registry, register_int_gauge_with_registry};
5use std::collections::HashMap;
6use tokio::sync::Mutex;
7
8/// Defines a structure to hold and manage metrics for a watchdog service.
9/// This structure is thread-safe, allowing concurrent access and modification of metrics.
10#[derive(Debug)]
11pub struct WatchdogMetrics {
12    // The Prometheus registry to which metrics are registered.
13    registry: Registry,
14    // A HashMap to store IntGauge metrics, keyed by their names.
15    // Wrapped in a Mutex to ensure thread-safe access.
16    metrics: Mutex<HashMap<String, IntGauge>>,
17}
18
19impl WatchdogMetrics {
20    /// Constructs a new WatchdogMetrics instance with the given Prometheus registry
21    pub fn new(registry: &Registry) -> Self {
22        Self {
23            registry: registry.clone(),
24            metrics: Mutex::new(HashMap::new()),
25        }
26    }
27
28    /// Retrieves or creates an metric for the specified metric name.
29    /// The metric name is suffixed with "_exact" to denote its type.
30    pub async fn get(&self, metric_name: &str) -> anyhow::Result<IntGauge> {
31        let mut metrics = self.metrics.lock().await;
32        // If the metric doesn't exist, register it and insert into the map.
33        let metric = metrics.entry(metric_name.to_string()).or_insert_with(|| {
34            register_int_gauge_with_registry!(metric_name, metric_name, &self.registry).unwrap()
35        });
36        Ok(metric.clone())
37    }
38
39    /// Retrieves or creates an "exact" metric for the specified metric name.
40    /// The metric name is suffixed with "_exact" to denote its type.
41    pub async fn get_exact(&self, metric_name: &str) -> anyhow::Result<IntGauge> {
42        let mut metrics = self.metrics.lock().await;
43        let metric_name = format!("{}_exact", metric_name);
44        // If the metric doesn't exist, register it and insert into the map.
45        let metric = metrics.entry(metric_name.clone()).or_insert_with(|| {
46            register_int_gauge_with_registry!(&metric_name, &metric_name, &self.registry).unwrap()
47        });
48        Ok(metric.clone())
49    }
50
51    /// Similar to get_exact, but for "lower" bound metrics.
52    pub async fn get_lower(&self, metric_name: &str) -> anyhow::Result<IntGauge> {
53        let mut metrics = self.metrics.lock().await;
54        let metric_name = format!("{}_lower", metric_name);
55        let metric = metrics.entry(metric_name.clone()).or_insert_with(|| {
56            register_int_gauge_with_registry!(&metric_name, &metric_name, &self.registry).unwrap()
57        });
58        Ok(metric.clone())
59    }
60
61    /// Similar to get_exact, but for "upper" bound metrics.
62    pub async fn get_upper(&self, metric_name: &str) -> anyhow::Result<IntGauge> {
63        let mut metrics = self.metrics.lock().await;
64        let metric_name = format!("{}_upper", metric_name);
65        let metric = metrics.entry(metric_name.clone()).or_insert_with(|| {
66            register_int_gauge_with_registry!(&metric_name, &metric_name, &self.registry).unwrap()
67        });
68        Ok(metric.clone())
69    }
70}
sui_security_watchdog/metrics.rs

sui_security_watchdog/
metrics.rs
