Skip to main content

iota_metrics/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    future::Future,
7    net::SocketAddr,
8    path::Path,
9    pin::Pin,
10    sync::Arc,
11    task::{Context, Poll},
12    time::Instant,
13};
14
15use axum::{
16    Router,
17    extract::{Extension, Request},
18    http::StatusCode,
19    middleware::Next,
20    response::Response,
21    routing::get,
22};
23use dashmap::DashMap;
24use once_cell::sync::OnceCell;
25use parking_lot::Mutex;
26use prometheus_filtered::{
27    Filter, Histogram, IntCounterVec, IntGaugeVec, Registry, TextEncoder,
28    core::{AtomicI64, GenericGauge},
29    register_histogram_with_registry, register_int_counter_vec_with_registry,
30    register_int_gauge_vec_with_registry,
31};
32pub use scopeguard;
33use simple_server_timing_header::Timer;
34use tap::TapFallible;
35use tracing::{Span, warn};
36use uuid::Uuid;
37
38mod guards;
39pub mod hardware_metrics;
40pub mod histogram;
41pub mod metered_channel;
42pub mod metric_groups;
43pub mod metrics_network;
44pub mod monitored_mpsc;
45pub mod peak;
46pub mod quantile_gauge;
47// Relies on tokio's `RuntimeMetrics`, which the deterministic simulator's tokio
48// fork does not provide; the node only starts these monitors outside simtests.
49#[cfg(not(msim))]
50pub mod runtime_metrics;
51#[cfg(any(test, feature = "test-utils"))]
52pub mod test_utils;
53pub mod thread_stall_monitor;
54pub use guards::*;
55pub use metric_groups::{MetricGroups, MetricLevel};
56
57pub const TX_TYPE_SINGLE_WRITER_TX: &str = "single_writer";
58pub const TX_TYPE_SHARED_OBJ_TX: &str = "shared_object";
59
60pub const SUBSECOND_LATENCY_SEC_BUCKETS: &[f64] = &[
61    0.005, 0.01, 0.02, 0.03, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.7, 1.,
62];
63
64pub const COARSE_LATENCY_SEC_BUCKETS: &[f64] = &[
65    0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.2, 0.3, 0.5, 0.7, 1., 2., 3., 5., 10., 20., 30., 60.,
66];
67
68pub const LATENCY_SEC_BUCKETS: &[f64] = &[
69    0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.6,
70    0.7, 0.8, 0.9, 1., 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2., 2.5, 3., 3.5, 4., 4.5, 5.,
71    6., 7., 8., 9., 10., 15., 20., 25., 30., 60., 90.,
72];
73
74pub const COUNT_BUCKETS: &[f64] = &[
75    2., 5., 10., 20., 50., 100., 200., 500., 1000., 2000., 5000., 10000.,
76];
77
78pub const BYTES_BUCKETS: &[f64] = &[
79    1024., 4096., 16384., 65536., 262144., 524288., 1048576., 2097152., 4194304., 8388608.,
80    16777216., 33554432., 67108864.,
81];
82
83#[derive(Debug)]
84pub struct Metrics {
85    pub tasks: IntGaugeVec,
86    pub futures: IntGaugeVec,
87    pub channel_inflight: IntGaugeVec,
88    pub channel_sent: IntGaugeVec,
89    pub channel_received: IntGaugeVec,
90    pub future_active_duration_ns: IntGaugeVec,
91    pub scope_iterations: IntGaugeVec,
92    pub scope_duration_ns: IntGaugeVec,
93    pub scope_entrance: IntGaugeVec,
94    pub thread_stall_duration_sec: Histogram,
95    pub system_invariant_violations: IntCounterVec,
96}
97
98impl Metrics {
99    /// Creates a new instance of the monitoring metrics, registering various
100    /// gauges and histograms with the provided metrics `Registry`. The
101    /// gauges track metrics such as the number of running tasks, pending
102    /// futures, channel items, and scope activities, while the histogram
103    /// measures the duration of thread stalls. Each metric is registered
104    /// with descriptive labels to facilitate performance monitoring and
105    /// analysis.
106    fn new(registry: &Registry) -> Self {
107        Self {
108            tasks: register_int_gauge_vec_with_registry!(
109                "monitored_tasks",
110                "Number of running tasks per callsite.",
111                &["callsite"],
112                registry;
113                MetricLevel::Warn,
114            )
115            .unwrap(),
116            futures: register_int_gauge_vec_with_registry!(
117                "monitored_futures",
118                "Number of pending futures per callsite.",
119                &["callsite"],
120                registry;
121                MetricLevel::Warn,
122            )
123            .unwrap(),
124            channel_inflight: register_int_gauge_vec_with_registry!(
125                "monitored_channel_inflight",
126                "Inflight items in channels.",
127                &["name"],
128                registry;
129                MetricLevel::Warn,
130            )
131            .unwrap(),
132            channel_sent: register_int_gauge_vec_with_registry!(
133                "monitored_channel_sent",
134                "Sent items in channels.",
135                &["name"],
136                registry,
137            )
138            .unwrap(),
139            channel_received: register_int_gauge_vec_with_registry!(
140                "monitored_channel_received",
141                "Received items in channels.",
142                &["name"],
143                registry,
144            )
145            .unwrap(),
146            future_active_duration_ns: register_int_gauge_vec_with_registry!(
147                "monitored_future_active_duration_ns",
148                "Total duration in nanosecs where the monitored future is active (consuming CPU time)",
149                &["name"],
150                registry,
151            )
152            .unwrap(),
153            scope_entrance: register_int_gauge_vec_with_registry!(
154                "monitored_scope_entrance",
155                "Number of entrance in the scope.",
156                &["name"],
157                registry,
158            )
159            .unwrap(),
160            scope_iterations: register_int_gauge_vec_with_registry!(
161                "monitored_scope_iterations",
162                "Total number of times where the monitored scope runs",
163                &["name"],
164                registry,
165            )
166            .unwrap(),
167            scope_duration_ns: register_int_gauge_vec_with_registry!(
168                "monitored_scope_duration_ns",
169                "Total duration in nanosecs where the monitored scope is running",
170                &["name"],
171                registry,
172            )
173            .unwrap(),
174            thread_stall_duration_sec: register_histogram_with_registry!(
175                "thread_stall_duration_sec",
176                "Duration of thread stalls in seconds.",
177                registry,
178            )
179            .unwrap(),
180            system_invariant_violations: register_int_counter_vec_with_registry!(
181                "system_invariant_violations",
182                "Number of system invariant violations",
183                &["name"],
184                registry,
185            ).unwrap(),
186        }
187    }
188}
189
190static METRICS: OnceCell<Metrics> = OnceCell::new();
191
192/// Initializes the global `METRICS` instance by setting it to a new `Metrics`
193/// object registered with the provided `Registry`. If `METRICS` is already set,
194/// a warning is logged indicating that the metrics registry was overwritten.
195/// This function is intended to be called once during initialization to set up
196/// metrics collection.
197pub fn init_metrics(registry: &Registry) {
198    let _ = METRICS
199        .set(Metrics::new(registry))
200        // this happens many times during tests
201        .tap_err(|_| warn!("init_metrics registry overwritten"));
202}
203
204/// Retrieves the global `METRICS` instance if it has been initialized.
205pub fn get_metrics() -> Option<&'static Metrics> {
206    METRICS.get()
207}
208
209tokio::task_local! {
210    static SERVER_TIMING: Arc<Mutex<Timer>>;
211}
212
213/// Create a new task-local ServerTiming context and run the provided future
214/// within it. Should be used at the top-most level of a request handler. Can be
215/// added to an axum router as a layer by using
216/// iota_service::server_timing_middleware.
217pub async fn with_new_server_timing<T>(fut: impl Future<Output = T> + Send + 'static) -> T {
218    let timer = Arc::new(Mutex::new(Timer::new()));
219
220    let mut ret = None;
221    SERVER_TIMING
222        .scope(timer, async {
223            ret = Some(fut.await);
224        })
225        .await;
226
227    ret.unwrap()
228}
229
230/// The `Server-Timing` HTTP header key.
231pub fn server_timing_header_key() -> &'static str {
232    Timer::header_key()
233}
234
235/// Add a final `finish_request` entry and write the collected timings to the
236/// `Server-Timing` response header. No-op outside a server-timing context.
237pub fn finish_and_set_server_timing_header(headers: &mut http::HeaderMap) {
238    let Some(timer) = get_server_timing() else {
239        return;
240    };
241    let header_value = {
242        let mut timer = timer.lock();
243        timer.add("finish_request");
244        timer.header_value()
245    };
246    if let Ok(value) = http::HeaderValue::try_from(header_value) {
247        headers.insert(server_timing_header_key(), value);
248    }
249}
250
251pub async fn server_timing_middleware(request: Request, next: Next) -> Response {
252    with_new_server_timing(async move {
253        let mut response = next.run(request).await;
254        finish_and_set_server_timing_header(response.headers_mut());
255        response
256    })
257    .await
258}
259
260/// Create a new task-local ServerTiming context and run the provided future
261/// within it. Only intended for use by macros within this module.
262pub async fn with_server_timing<T>(
263    timer: Arc<Mutex<Timer>>,
264    fut: impl Future<Output = T> + Send + 'static,
265) -> T {
266    let mut ret = None;
267    SERVER_TIMING
268        .scope(timer, async {
269            ret = Some(fut.await);
270        })
271        .await;
272
273    ret.unwrap()
274}
275
276/// Get the currently active ServerTiming context. Only intended for use by
277/// macros within this module.
278pub fn get_server_timing() -> Option<Arc<Mutex<Timer>>> {
279    SERVER_TIMING.try_with(|timer| timer.clone()).ok()
280}
281
282/// Add a new entry to the ServerTiming header.
283/// If the caller is not currently in a ServerTiming context (created with
284/// `with_new_server_timing`), an error is logged.
285pub fn add_server_timing(name: &str) {
286    let res = SERVER_TIMING.try_with(|timer| {
287        timer.lock().add(name);
288    });
289
290    if res.is_err() {
291        tracing::error!("Server timing context not found");
292    }
293}
294
295#[macro_export]
296macro_rules! monitored_future {
297    ($fut: expr) => {{ monitored_future!(futures, $fut, "", INFO, false) }};
298
299    ($metric: ident, $fut: expr, $name: expr, $logging_level: ident, $logging_enabled: expr) => {{
300        let location: &str = if $name.is_empty() {
301            concat!(file!(), ':', line!())
302        } else {
303            concat!(file!(), ':', $name)
304        };
305
306        async move {
307            let metrics = $crate::get_metrics();
308
309            let _metrics_guard = if let Some(m) = metrics {
310                m.$metric.with_label_values(&[location]).inc();
311                Some($crate::scopeguard::guard(m, |_| {
312                    m.$metric.with_label_values(&[location]).dec();
313                }))
314            } else {
315                None
316            };
317            let _logging_guard = if $logging_enabled {
318                Some($crate::scopeguard::guard((), |_| {
319                    tracing::event!(
320                        tracing::Level::$logging_level,
321                        "Future {} completed",
322                        location
323                    );
324                }))
325            } else {
326                None
327            };
328
329            if $logging_enabled {
330                tracing::event!(
331                    tracing::Level::$logging_level,
332                    "Spawning future {}",
333                    location
334                );
335            }
336
337            $fut.await
338        }
339    }};
340}
341
342#[macro_export]
343macro_rules! forward_server_timing_and_spawn {
344    ($fut: expr) => {
345        if let Some(timing) = $crate::get_server_timing() {
346            tokio::task::spawn(async move { $crate::with_server_timing(timing, $fut).await })
347        } else {
348            tokio::task::spawn($fut)
349        }
350    };
351}
352
353#[macro_export]
354macro_rules! spawn_monitored_task {
355    ($fut: expr) => {
356        $crate::forward_server_timing_and_spawn!($crate::monitored_future!(
357            tasks, $fut, "", INFO, false
358        ))
359    };
360}
361
362#[macro_export]
363macro_rules! spawn_logged_monitored_task {
364    ($fut: expr) => {
365        $crate::forward_server_timing_and_spawn!($crate::monitored_future!(
366            tasks, $fut, "", INFO, true
367        ))
368    };
369
370    ($fut: expr, $name: expr) => {
371        $crate::forward_server_timing_and_spawn!($crate::monitored_future!(
372            tasks, $fut, $name, INFO, true
373        ))
374    };
375
376    ($fut: expr, $name: expr, $logging_level: ident) => {
377        $crate::forward_server_timing_and_spawn!($crate::monitored_future!(
378            tasks,
379            $fut,
380            $name,
381            $logging_level,
382            true
383        ))
384    };
385}
386
387pub struct MonitoredScopeGuard {
388    metrics: &'static Metrics,
389    name: &'static str,
390    timer: Instant,
391}
392
393impl Drop for MonitoredScopeGuard {
394    fn drop(&mut self) {
395        self.metrics
396            .scope_duration_ns
397            .with_label_values(&[self.name])
398            .add(self.timer.elapsed().as_nanos() as i64);
399        self.metrics
400            .scope_entrance
401            .with_label_values(&[self.name])
402            .dec();
403    }
404}
405
406/// This function creates a named scoped object, that keeps track of
407/// - the total iterations where the scope is called in the `monitored_scope_iterations` metric.
408/// - and the total duration of the scope in the `monitored_scope_duration_ns` metric.
409///
410/// The monitored scope should be single threaded, e.g. the scoped object
411/// encompass the lifetime of a select loop or guarded by mutex.
412/// Then the rate of `monitored_scope_duration_ns`, converted to the unit of sec
413/// / sec, would be how full the single threaded scope is running.
414pub fn monitored_scope(name: &'static str) -> Option<MonitoredScopeGuard> {
415    let metrics = get_metrics();
416    if let Some(m) = metrics {
417        m.scope_iterations.with_label_values(&[name]).inc();
418        m.scope_entrance.with_label_values(&[name]).inc();
419        Some(MonitoredScopeGuard {
420            metrics: m,
421            name,
422            timer: Instant::now(),
423        })
424    } else {
425        None
426    }
427}
428
429/// A trait extension for `Future` to allow monitoring the execution of the
430/// future within a specific scope. Provides the `in_monitored_scope` method to
431/// wrap the future in a `MonitoredScopeFuture`, which tracks the future's
432/// execution using a `MonitoredScopeGuard` for monitoring purposes.
433pub trait MonitoredFutureExt: Future + Sized {
434    /// Wraps the current future in a `MonitoredScopeFuture` that is associated
435    /// with a specific monitored scope name. The scope helps track the
436    /// execution of the future for performance analysis and metrics collection.
437    fn in_monitored_scope(self, name: &'static str) -> MonitoredScopeFuture<Self>;
438}
439
440impl<F: Future> MonitoredFutureExt for F {
441    fn in_monitored_scope(self, name: &'static str) -> MonitoredScopeFuture<Self> {
442        MonitoredScopeFuture {
443            f: Box::pin(self),
444            active_duration_metric: get_metrics()
445                .map(|m| m.future_active_duration_ns.with_label_values(&[name])),
446            _scope: monitored_scope(name),
447        }
448    }
449}
450
451/// A future that runs within a monitored scope. This struct wraps a pinned
452/// future and holds an optional `MonitoredScopeGuard` to measure and monitor
453/// the execution of the future. It forwards polling operations
454/// to the underlying future while maintaining the monitoring scope.
455pub struct MonitoredScopeFuture<F: Sized> {
456    f: Pin<Box<F>>,
457    active_duration_metric: Option<GenericGauge<AtomicI64>>,
458    _scope: Option<MonitoredScopeGuard>,
459}
460
461impl<F: Future> Future for MonitoredScopeFuture<F> {
462    type Output = F::Output;
463
464    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
465        let active_timer = Instant::now();
466        let ret = self.f.as_mut().poll(cx);
467        if let Some(m) = &self.active_duration_metric {
468            m.add(active_timer.elapsed().as_nanos() as i64);
469        }
470        ret
471    }
472}
473
474/// A future that runs within a monitored scope. This struct wraps a pinned
475/// future and holds an optional `MonitoredScopeGuard` to measure and monitor
476/// the execution of the future. It forwards polling operations
477/// to the underlying future while maintaining the monitoring scope.
478pub struct CancelMonitor<F: Sized> {
479    finished: bool,
480    inner: Pin<Box<F>>,
481}
482
483impl<F> CancelMonitor<F>
484where
485    F: Future,
486{
487    /// Creates a new `CancelMonitor` that wraps the given future (`inner`). The
488    /// monitor tracks whether the future has completed.
489    pub fn new(inner: F) -> Self {
490        Self {
491            finished: false,
492            inner: Box::pin(inner),
493        }
494    }
495
496    /// Returns `true` if the future has completed; otherwise, `false`.
497    pub fn is_finished(&self) -> bool {
498        self.finished
499    }
500}
501
502impl<F> Future for CancelMonitor<F>
503where
504    F: Future,
505{
506    type Output = F::Output;
507
508    /// Polls the inner future to determine if it is ready or still pending. For
509    /// `CancelMonitor`, if the future completes (`Poll::Ready`), `finished`
510    /// is set to `true`. If it is still pending, the status remains
511    /// unchanged. This allows monitoring of the future's completion status.
512    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
513        match self.inner.as_mut().poll(cx) {
514            Poll::Ready(output) => {
515                self.finished = true;
516                Poll::Ready(output)
517            }
518            Poll::Pending => Poll::Pending,
519        }
520    }
521}
522
523impl<F: Sized> Drop for CancelMonitor<F> {
524    /// When the `CancelMonitor` is dropped, it checks whether the future has
525    /// finished executing. If the future was not completed (`finished` is
526    /// `false`), it records that the future was cancelled by logging the
527    /// cancellation status using the current span.
528    fn drop(&mut self) {
529        if !self.finished {
530            Span::current().record("cancelled", true);
531        }
532    }
533}
534
535/// MonitorCancellation records a cancelled = true span attribute if the future
536/// it is decorating is dropped before completion. The cancelled attribute must
537/// be added at span creation, as you cannot add new attributes after the span
538/// is created.
539pub trait MonitorCancellation {
540    fn monitor_cancellation(self) -> CancelMonitor<Self>
541    where
542        Self: Sized + Future;
543}
544
545impl<T> MonitorCancellation for T
546where
547    T: Future,
548{
549    fn monitor_cancellation(self) -> CancelMonitor<Self> {
550        CancelMonitor::new(self)
551    }
552}
553
554pub type RegistryID = Uuid;
555
556/// A service to manage the prometheus registries. This service allow us to
557/// create a new Registry on demand and keep it accessible for
558/// processing/polling. The service can be freely cloned/shared across threads.
559#[derive(Clone)]
560pub struct RegistryService {
561    // Holds a Registry that is supposed to be used
562    default_registry: Registry,
563    registries_by_id: Arc<DashMap<Uuid, Registry>>,
564    filter: Arc<Filter>,
565}
566
567impl RegistryService {
568    // Creates a new registry service and also adds the main/default registry that
569    // is supposed to be preserved and never get removed
570    pub fn new(default_registry: Registry) -> Self {
571        Self {
572            filter: default_registry.filter(),
573            default_registry,
574            registries_by_id: Arc::new(DashMap::new()),
575        }
576    }
577
578    // Returns the default registry for the service that someone can use
579    // if they don't want to create a new one.
580    pub fn default_registry(&self) -> Registry {
581        self.default_registry.clone()
582    }
583
584    /// Returns the metrics filter shared by the service's registries.
585    pub fn filter(&self) -> Arc<Filter> {
586        self.filter.clone()
587    }
588
589    // Creates a new registry that shares the service's metric filter. Prefer
590    // this over `Registry::new()`/`Registry::new_custom()` for registries added
591    // via `add`, so that the configured filter applies to their metrics too.
592    pub fn new_registry_custom(
593        &self,
594        prefix: Option<String>,
595        labels: Option<std::collections::HashMap<String, String>>,
596    ) -> prometheus_filtered::Result<Registry> {
597        Registry::new_custom(prefix, labels, Some(self.filter.clone()))
598    }
599
600    // Adds a new registry to the service. The corresponding RegistryID is returned
601    // so can later be used for removing the Registry. Method panics if we try
602    // to insert a registry with the same id. As this can be quite serious for
603    // the operation of the node we don't want to accidentally swap an existing
604    // registry - we expected a removal to happen explicitly.
605    pub fn add(&self, registry: Registry) -> RegistryID {
606        let registry_id = Uuid::new_v4();
607        if self
608            .registries_by_id
609            .insert(registry_id, registry)
610            .is_some()
611        {
612            panic!("Other Registry already detected for the same id {registry_id}");
613        }
614
615        registry_id
616    }
617
618    // Removes the registry from the service. If Registry existed then this method
619    // returns true, otherwise false is returned instead.
620    pub fn remove(&self, registry_id: RegistryID) -> bool {
621        self.registries_by_id.remove(&registry_id).is_some()
622    }
623
624    // Returns all the registries of the service
625    pub fn get_all(&self) -> Vec<Registry> {
626        let mut registries: Vec<Registry> = self
627            .registries_by_id
628            .iter()
629            .map(|r| r.value().clone())
630            .collect();
631        registries.push(self.default_registry.clone());
632
633        registries
634    }
635
636    // Returns all the metric families from the registries that a service holds.
637    pub fn gather_all(&self) -> Vec<prometheus_filtered::proto::MetricFamily> {
638        self.get_all().iter().flat_map(|r| r.gather()).collect()
639    }
640
641    /// Sets the runtime override on the shared filter; every registry's next
642    /// gather exposes metrics per the new directives. `filter` may use
643    /// [`MetricGroups`] names, which are expanded for matching while the
644    /// string is echoed back as given. Rejects the whole update if any
645    /// directive is invalid.
646    pub fn set_runtime_filter(&self, filter: &str) -> std::result::Result<(), String> {
647        let expanded = MetricGroups::expand_directives(filter)?;
648        self.filter
649            .set_runtime_filter(prometheus_filtered::FilterSource::with_display(
650                &expanded, filter,
651            ))
652    }
653
654    /// Drops the runtime override on the shared filter, restoring every
655    /// registry to its startup exposure.
656    pub fn reset_runtime_filter(&self) {
657        self.filter.reset_runtime_filter();
658    }
659}
660
661/// Create a metric that measures the uptime from when this metric was
662/// constructed. The metric is labeled with:
663/// - 'process': the process type, differentiating between validator and fullnode
664/// - 'version': binary version, generally be of the format: 'semver-gitrevision'
665/// - 'chain_identifier': the identifier of the network which this process is part of
666pub fn uptime_metric(
667    process: &str,
668    version: &'static str,
669    chain_identifier: &str,
670) -> Box<dyn prometheus_filtered::core::Collector> {
671    let opts = prometheus_filtered::opts!("uptime", "uptime of the node service in seconds")
672        .variable_label("process")
673        .variable_label("version")
674        .variable_label("chain_identifier")
675        .variable_label("os_version")
676        .variable_label("is_docker");
677
678    let start_time = std::time::Instant::now();
679    let uptime = move || start_time.elapsed().as_secs();
680    let metric = prometheus_closure_metric::ClosureMetric::new(
681        opts,
682        prometheus_closure_metric::ValueType::Counter,
683        uptime,
684        &[
685            process,
686            version,
687            chain_identifier,
688            &sysinfo::System::long_os_version()
689                .unwrap_or_else(|| "os_version_unavailable".to_string()),
690            &is_running_in_docker().to_string(),
691        ],
692    )
693    .unwrap();
694
695    Box::new(metric)
696}
697
698pub fn is_running_in_docker() -> bool {
699    // Check for .dockerenv file instead. This file exists in the debian:__-slim
700    // image we use at runtime.
701    Path::new("/.dockerenv").exists()
702}
703
704pub const METRICS_ROUTE: &str = "/metrics";
705
706// Creates a new http server that has as a sole purpose to expose
707// and endpoint that prometheus agent can use to poll for the metrics.
708// A RegistryService is returned that can be used to get access in prometheus
709// Registries.
710pub fn start_prometheus_server(addr: SocketAddr) -> RegistryService {
711    start_prometheus_server_with_filter(addr, Filter::from_env())
712}
713
714pub fn start_prometheus_server_with_filter(addr: SocketAddr, filter: Filter) -> RegistryService {
715    // The default registry has no prefix or labels, so its construction is
716    // infallible.
717    let registry = Registry::new_custom(None, None, Some(Arc::new(filter)))
718        .expect("unprefixed registry is infallible");
719
720    let registry_service = RegistryService::new(registry);
721
722    if cfg!(msim) {
723        // prometheus uses difficult-to-support features such as
724        // TcpSocket::from_raw_fd(), so we can't yet run it in the simulator.
725        warn!("not starting prometheus server in simulator");
726        return registry_service;
727    }
728
729    let app = Router::new()
730        .route(METRICS_ROUTE, get(metrics))
731        .layer(Extension(registry_service.clone()));
732
733    tokio::spawn(async move {
734        let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
735        axum::serve(listener, app.into_make_service())
736            .await
737            .unwrap();
738    });
739
740    registry_service
741}
742
743/// Handles a request to retrieve metrics, using the provided `RegistryService`
744/// to gather all registered metric families. The metrics are then encoded to a
745/// text format for easy consumption by monitoring systems. If successful, it
746/// returns the metrics string with an `OK` status. If an error occurs during
747/// encoding, it returns an `INTERNAL_SERVER_ERROR` status along with an error
748/// message. Returns a tuple containing the status code and either the metrics
749/// data or an error description.
750pub async fn metrics(
751    Extension(registry_service): Extension<RegistryService>,
752) -> (StatusCode, String) {
753    let metrics_families = registry_service.gather_all();
754    match TextEncoder.encode_to_string(&metrics_families) {
755        Ok(metrics) => (StatusCode::OK, metrics),
756        Err(error) => (
757            StatusCode::INTERNAL_SERVER_ERROR,
758            format!("unable to encode metrics: {error}"),
759        ),
760    }
761}
762
763#[cfg(test)]
764mod tests {
765    use prometheus_filtered::{IntCounter, Registry};
766
767    use crate::RegistryService;
768
769    #[test]
770    fn registry_service() {
771        // GIVEN
772        let default_registry =
773            Registry::new_custom(Some("default".to_string()), None, None).unwrap();
774
775        let registry_service = RegistryService::new(default_registry.clone());
776        let default_counter = IntCounter::new("counter", "counter_desc").unwrap();
777        default_counter.inc();
778        default_registry
779            .register(Box::new(default_counter))
780            .unwrap();
781
782        // AND add a metric to the default registry
783
784        // AND a registry with one metric
785        let registry_1 = Registry::new_custom(Some("iota".to_string()), None, None).unwrap();
786        registry_1
787            .register(Box::new(
788                IntCounter::new("counter_1", "counter_1_desc").unwrap(),
789            ))
790            .unwrap();
791
792        // WHEN
793        let registry_1_id = registry_service.add(registry_1);
794
795        // THEN
796        let mut metrics = registry_service.gather_all();
797        metrics.sort_by(|m1, m2| Ord::cmp(m1.name(), m2.name()));
798
799        assert_eq!(metrics.len(), 2);
800
801        let metric_default = metrics.remove(0);
802        assert_eq!(metric_default.name(), "default_counter");
803        assert_eq!(metric_default.help(), "counter_desc");
804
805        let metric_1: prometheus_filtered::proto::MetricFamily = metrics.remove(0);
806        assert_eq!(metric_1.name(), "iota_counter_1");
807        assert_eq!(metric_1.help(), "counter_1_desc");
808
809        // AND add a second registry with a metric
810        let registry_2 = Registry::new_custom(Some("iota".to_string()), None, None).unwrap();
811        registry_2
812            .register(Box::new(
813                IntCounter::new("counter_2", "counter_2_desc").unwrap(),
814            ))
815            .unwrap();
816        let _registry_2_id = registry_service.add(registry_2);
817
818        // THEN all the metrics should be returned
819        let mut metrics = registry_service.gather_all();
820        metrics.sort_by(|m1, m2| Ord::cmp(m1.name(), m2.name()));
821
822        assert_eq!(metrics.len(), 3);
823
824        let metric_default = metrics.remove(0);
825        assert_eq!(metric_default.name(), "default_counter");
826        assert_eq!(metric_default.help(), "counter_desc");
827
828        let metric_1 = metrics.remove(0);
829        assert_eq!(metric_1.name(), "iota_counter_1");
830        assert_eq!(metric_1.help(), "counter_1_desc");
831
832        let metric_2 = metrics.remove(0);
833        assert_eq!(metric_2.name(), "iota_counter_2");
834        assert_eq!(metric_2.help(), "counter_2_desc");
835
836        // AND remove first registry
837        assert!(registry_service.remove(registry_1_id));
838
839        // THEN metrics should now not contain metric of registry_1
840        let mut metrics = registry_service.gather_all();
841        metrics.sort_by(|m1, m2| Ord::cmp(m1.name(), m2.name()));
842
843        assert_eq!(metrics.len(), 2);
844
845        let metric_default = metrics.remove(0);
846        assert_eq!(metric_default.name(), "default_counter");
847        assert_eq!(metric_default.help(), "counter_desc");
848
849        let metric_1 = metrics.remove(0);
850        assert_eq!(metric_1.name(), "iota_counter_2");
851        assert_eq!(metric_1.help(), "counter_2_desc");
852    }
853
854    #[test]
855    fn set_runtime_filter_applies_to_all_registries() {
856        use std::sync::Arc;
857
858        use prometheus_filtered::{Filter, MetricLevel};
859
860        fn gathered_names(service: &RegistryService) -> Vec<String> {
861            let mut names: Vec<_> = service
862                .gather_all()
863                .iter()
864                .map(|f| f.name().to_owned())
865                .collect();
866            names.sort();
867            names
868        }
869
870        // Both registries share the service's filter, which starts at `warn`
871        // for this module, hiding the default-level (`debug`) gauges.
872        let filter = Arc::new(Filter::parse("iota_metrics=warn"));
873        let default_registry =
874            Registry::new_custom(Some("default".to_string()), None, Some(filter)).unwrap();
875        let registry_service = RegistryService::new(default_registry.clone());
876        let second_registry = registry_service
877            .new_registry_custom(Some("second".to_string()), None)
878            .unwrap();
879        registry_service.add(second_registry.clone());
880
881        for registry in [&default_registry, &second_registry] {
882            prometheus_filtered::register_int_gauge_with_registry!(
883                "g_warn", "h", registry; MetricLevel::Warn
884            )
885            .unwrap();
886            prometheus_filtered::register_int_gauge_with_registry!("g_debug", "h", registry)
887                .unwrap();
888        }
889        assert_eq!(
890            gathered_names(&registry_service),
891            ["default_g_warn", "second_g_warn"]
892        );
893
894        // One runtime update changes the exposure of every registry in the
895        // service.
896        registry_service
897            .set_runtime_filter("runtime=debug")
898            .unwrap();
899        assert_eq!(
900            gathered_names(&registry_service),
901            [
902                "default_g_debug",
903                "default_g_warn",
904                "second_g_debug",
905                "second_g_warn"
906            ]
907        );
908
909        // Reset restores the startup exposure across all registries.
910        registry_service.reset_runtime_filter();
911        assert_eq!(
912            gathered_names(&registry_service),
913            ["default_g_warn", "second_g_warn"]
914        );
915    }
916}