iota_rest_api/
metrics.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{borrow::Cow, sync::Arc, time::Instant};
6
7use axum::http;
8use iota_network_stack::callback::{MakeCallbackHandler, ResponseHandler};
9use prometheus::{
10    HistogramVec, IntCounterVec, IntGaugeVec, Registry, register_histogram_vec_with_registry,
11    register_int_counter_vec_with_registry, register_int_gauge_vec_with_registry,
12};
13
14#[derive(Clone)]
15pub struct RestMetrics {
16    inflight_requests: IntGaugeVec,
17    num_requests: IntCounterVec,
18    request_latency: HistogramVec,
19}
20
21const LATENCY_SEC_BUCKETS: &[f64] = &[
22    0.001, 0.005, 0.01, 0.05, 0.1, 0.25, 0.5, 1., 2.5, 5., 10., 20., 30., 60., 90.,
23];
24
25impl RestMetrics {
26    pub fn new(registry: &Registry) -> Self {
27        Self {
28            inflight_requests: register_int_gauge_vec_with_registry!(
29                "rest_inflight_requests",
30                "Total in-flight REST requests per route",
31                &["path"],
32                registry,
33            )
34            .unwrap(),
35            num_requests: register_int_counter_vec_with_registry!(
36                "rest_requests",
37                "Total REST requests per route and their http status",
38                &["path", "status"],
39                registry,
40            )
41            .unwrap(),
42            request_latency: register_histogram_vec_with_registry!(
43                "rest_request_latency",
44                "Latency of REST requests per route",
45                &["path"],
46                LATENCY_SEC_BUCKETS.to_vec(),
47                registry,
48            )
49            .unwrap(),
50        }
51    }
52}
53
54#[derive(Clone)]
55pub struct RestMetricsMakeCallbackHandler {
56    metrics: Arc<RestMetrics>,
57}
58
59impl RestMetricsMakeCallbackHandler {
60    pub fn new(metrics: Arc<RestMetrics>) -> Self {
61        Self { metrics }
62    }
63}
64
65impl MakeCallbackHandler for RestMetricsMakeCallbackHandler {
66    type Handler = RestMetricsCallbackHandler;
67
68    fn make_handler(&self, request: &http::request::Parts) -> Self::Handler {
69        let start = Instant::now();
70        let metrics = self.metrics.clone();
71
72        let path = if let Some(path) = request.extensions.get::<axum::extract::MatchedPath>() {
73            Cow::Owned(path.as_str().to_owned())
74        } else {
75            Cow::Borrowed("unknown")
76        };
77
78        metrics
79            .inflight_requests
80            .with_label_values(&[path.as_ref()])
81            .inc();
82
83        RestMetricsCallbackHandler {
84            metrics,
85            path,
86            start,
87            counted_response: false,
88        }
89    }
90}
91
92pub struct RestMetricsCallbackHandler {
93    metrics: Arc<RestMetrics>,
94    path: Cow<'static, str>,
95    start: Instant,
96    // Indicates if we successfully counted the response. In some cases when a request is
97    // prematurely canceled this will remain false
98    counted_response: bool,
99}
100
101impl ResponseHandler for RestMetricsCallbackHandler {
102    fn on_response(mut self, response: &http::response::Parts) {
103        self.metrics
104            .num_requests
105            .with_label_values(&[self.path.as_ref(), response.status.as_str()])
106            .inc();
107
108        self.counted_response = true;
109    }
110
111    fn on_error<E>(self, _error: &E) {
112        // Do nothing if the whole service errored
113        //
114        // in Axum this isn't possible since all services are required to have
115        // an error type of Infallible
116    }
117}
118
119impl Drop for RestMetricsCallbackHandler {
120    fn drop(&mut self) {
121        self.metrics
122            .inflight_requests
123            .with_label_values(&[self.path.as_ref()])
124            .dec();
125
126        let latency = self.start.elapsed().as_secs_f64();
127        self.metrics
128            .request_latency
129            .with_label_values(&[self.path.as_ref()])
130            .observe(latency);
131
132        if !self.counted_response {
133            self.metrics
134                .num_requests
135                .with_label_values(&[self.path.as_ref(), "canceled"])
136                .inc();
137        }
138    }
139}
iota_rest_api/metrics.rs

iota_rest_api/
metrics.rs
