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