api_tools/server/axum/layers/
prometheus.rs

1//! Prometheus' metrics layer
2//!
3//! This module provides:
4//!
5//! 1. [`PrometheusLayer`] — a tower [`Layer`] that records per-request HTTP
6//!    metrics (`http_requests_total`, `http_requests_duration_seconds`).
7//!    The middleware overhead is in the microsecond range.
8//!
9//! 2. [`spawn_system_metrics_collector`] — a helper that spawns a background
10//!    Tokio task to collect host-level metrics (CPU, memory, swap, disk
11//!    usage). System metrics are intentionally **not** collected from the
12//!    request path: they do not change at request granularity, and collecting
13//!    them inline would add hundreds of milliseconds to every response (see
14//!    `sysinfo::MINIMUM_CPU_UPDATE_INTERVAL`).
15//!
16//! # Example
17//!
18//! ```ignore
19//! use std::path::PathBuf;
20//! use std::time::Duration;
21//! use api_tools::server::axum::layers::prometheus::{
22//!     PrometheusLayer, spawn_system_metrics_collector,
23//! };
24//!
25//! let layer = PrometheusLayer { service_name: "myapp".into() };
26//!
27//! // Once, at application startup:
28//! let _collector = spawn_system_metrics_collector(
29//!     "myapp".into(),
30//!     vec![PathBuf::from("/")],
31//!     Duration::from_secs(10),
32//! );
33//! ```
34
35use axum::body::Body;
36use axum::extract::MatchedPath;
37use axum::http::{Method, Request};
38use axum::response::Response;
39use futures::future::BoxFuture;
40use metrics::{SharedString, counter, gauge, histogram};
41use std::borrow::Cow;
42use std::path::PathBuf;
43use std::sync::Arc;
44use std::task::{Context, Poll};
45use std::time::{Duration, Instant};
46use sysinfo::{CpuRefreshKind, Disks, MemoryRefreshKind, RefreshKind, System};
47use tokio::task::JoinHandle;
48use tower::{Layer, Service};
49
50/// Prometheus metrics layer for Axum.
51///
52/// Records `http_requests_total` (counter) and
53/// `http_requests_duration_seconds` (histogram) for every request, labeled
54/// by `method`, `path` (the matched route — bounded cardinality), `service`
55/// and `status`. Requests to `/metrics` are excluded.
56///
57/// System metrics (CPU, memory, swap, disks) are **not** collected here.
58/// Use [`spawn_system_metrics_collector`] at startup instead.
59#[derive(Clone)]
60pub struct PrometheusLayer {
61    /// Service name used as a label on every metric.
62    pub service_name: String,
63}
64
65impl<S> Layer<S> for PrometheusLayer {
66    type Service = PrometheusMiddleware<S>;
67
68    fn layer(&self, inner: S) -> Self::Service {
69        PrometheusMiddleware {
70            inner,
71            // One-time conversion: subsequent per-request clones bump the
72            // refcount only.
73            service_name: Arc::from(self.service_name.as_str()),
74        }
75    }
76}
77
78#[derive(Clone)]
79pub struct PrometheusMiddleware<S> {
80    inner: S,
81    service_name: Arc<str>,
82}
83
84/// Map standard HTTP methods to a `&'static str` to avoid an allocation on
85/// every request. Falls back to an owned string for non-standard methods.
86fn method_label(method: &Method) -> Cow<'static, str> {
87    match *method {
88        Method::GET => Cow::Borrowed("GET"),
89        Method::POST => Cow::Borrowed("POST"),
90        Method::PUT => Cow::Borrowed("PUT"),
91        Method::DELETE => Cow::Borrowed("DELETE"),
92        Method::PATCH => Cow::Borrowed("PATCH"),
93        Method::HEAD => Cow::Borrowed("HEAD"),
94        Method::OPTIONS => Cow::Borrowed("OPTIONS"),
95        Method::CONNECT => Cow::Borrowed("CONNECT"),
96        Method::TRACE => Cow::Borrowed("TRACE"),
97        _ => Cow::Owned(method.as_str().to_owned()),
98    }
99}
100
101impl<S> Service<Request<Body>> for PrometheusMiddleware<S>
102where
103    S: Service<Request<Body>, Response = Response> + Send + 'static,
104    S::Future: Send + 'static,
105{
106    type Response = S::Response;
107    type Error = S::Error;
108    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
109    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
110
111    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
112        self.inner.poll_ready(cx)
113    }
114
115    fn call(&mut self, request: Request<Body>) -> Self::Future {
116        let path = if let Some(matched_path) = request.extensions().get::<MatchedPath>() {
117            matched_path.as_str().to_owned()
118        } else {
119            request.uri().path().to_owned()
120        };
121        let method = method_label(request.method());
122        let service_name = Arc::clone(&self.service_name);
123
124        let start = Instant::now();
125        let future = self.inner.call(request);
126        Box::pin(async move {
127            let response = future.await?;
128
129            // Exclude metrics endpoint
130            if path != "/metrics" {
131                let latency = start.elapsed().as_secs_f64();
132                let status = status_label(response.status().as_u16());
133                let labels: [(&'static str, SharedString); 4] = [
134                    ("method", method.into()),
135                    ("path", path.into()),
136                    ("service", service_name.into()),
137                    ("status", status.into()),
138                ];
139
140                counter!("http_requests_total", &labels).increment(1);
141                histogram!("http_requests_duration_seconds", &labels).record(latency);
142            }
143
144            Ok(response)
145        })
146    }
147}
148
149/// Map common HTTP status codes to a `&'static str` to avoid formatting an
150/// integer on every response. Falls back to an owned string for uncommon
151/// codes.
152fn status_label(code: u16) -> Cow<'static, str> {
153    match code {
154        200 => Cow::Borrowed("200"),
155        201 => Cow::Borrowed("201"),
156        204 => Cow::Borrowed("204"),
157        301 => Cow::Borrowed("301"),
158        302 => Cow::Borrowed("302"),
159        304 => Cow::Borrowed("304"),
160        400 => Cow::Borrowed("400"),
161        401 => Cow::Borrowed("401"),
162        403 => Cow::Borrowed("403"),
163        404 => Cow::Borrowed("404"),
164        409 => Cow::Borrowed("409"),
165        422 => Cow::Borrowed("422"),
166        500 => Cow::Borrowed("500"),
167        502 => Cow::Borrowed("502"),
168        503 => Cow::Borrowed("503"),
169        504 => Cow::Borrowed("504"),
170        _ => Cow::Owned(code.to_string()),
171    }
172}
173
174/// Spawn a background task that periodically refreshes host metrics and
175/// publishes them as Prometheus gauges.
176///
177/// Emitted gauges (all labeled by `service`):
178///
179/// - `system_cpu_usage` — average global CPU usage in percent
180/// - `system_total_memory` / `system_used_memory` — bytes
181/// - `system_total_swap` / `system_used_swap` — bytes
182/// - `system_total_disks_space` / `system_used_disks_space` — bytes,
183///   summed over `disk_mount_points`
184///
185/// The first tick reports `system_cpu_usage = 0.0` because `sysinfo` needs
186/// two snapshots to compute a delta. Subsequent ticks report the real value.
187///
188/// The returned [`JoinHandle`] can be aborted at shutdown; if dropped, the
189/// task continues running for the lifetime of the Tokio runtime.
190pub fn spawn_system_metrics_collector(
191    service_name: String,
192    disk_mount_points: Vec<PathBuf>,
193    interval: Duration,
194) -> JoinHandle<()> {
195    tokio::spawn(async move {
196        let refresh_kind = RefreshKind::nothing()
197            .with_cpu(CpuRefreshKind::nothing().with_cpu_usage())
198            .with_memory(MemoryRefreshKind::everything());
199        let mut sys = System::new_with_specifics(refresh_kind);
200
201        let mut ticker = tokio::time::interval(interval);
202        loop {
203            ticker.tick().await;
204
205            sys.refresh_cpu_usage();
206            sys.refresh_memory();
207
208            let cpu_usage = sys.global_cpu_usage();
209            let total_memory = sys.total_memory();
210            let used_memory = sys.used_memory();
211            let total_swap = sys.total_swap();
212            let used_swap = sys.used_swap();
213
214            let disks = Disks::new_with_refreshed_list();
215            let mut total_disks_space: u64 = 0;
216            let mut used_disks_space: u64 = 0;
217            for disk in &disks {
218                if disk_mount_points.contains(&disk.mount_point().to_path_buf()) {
219                    total_disks_space += disk.total_space();
220                    used_disks_space += disk.total_space().saturating_sub(disk.available_space());
221                }
222            }
223
224            gauge!("system_cpu_usage", "service" => service_name.clone()).set(cpu_usage);
225            gauge!("system_total_memory", "service" => service_name.clone()).set(total_memory as f64);
226            gauge!("system_used_memory", "service" => service_name.clone()).set(used_memory as f64);
227            gauge!("system_total_swap", "service" => service_name.clone()).set(total_swap as f64);
228            gauge!("system_used_swap", "service" => service_name.clone()).set(used_swap as f64);
229            gauge!("system_total_disks_space", "service" => service_name.clone()).set(total_disks_space as f64);
230            gauge!("system_used_disks_space", "service" => service_name.clone()).set(used_disks_space as f64);
231        }
232    })
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use axum::body::Body;
239    use axum::http::{Request, Response, StatusCode};
240    use std::convert::Infallible;
241    use tower::{ServiceBuilder, ServiceExt};
242
243    /// Sentinel test: the middleware must not block on host-metrics collection.
244    /// In 0.7.x, this call took ~200 ms because of an in-path
245    /// `tokio::time::sleep(MINIMUM_CPU_UPDATE_INTERVAL)`. After the refactor,
246    /// it should complete in well under 10 ms.
247    #[tokio::test]
248    async fn middleware_does_not_block_on_system_metrics() {
249        let svc = ServiceBuilder::new()
250            .layer(PrometheusLayer {
251                service_name: "test".into(),
252            })
253            .service(tower::service_fn(|_req: Request<Body>| async {
254                Ok::<_, Infallible>(Response::builder().status(StatusCode::OK).body(Body::empty()).unwrap())
255            }));
256
257        let start = Instant::now();
258        let response = svc
259            .oneshot(Request::builder().uri("/health").body(Body::empty()).unwrap())
260            .await
261            .unwrap();
262        let elapsed = start.elapsed();
263
264        assert_eq!(response.status(), StatusCode::OK);
265        assert!(
266            elapsed < Duration::from_millis(50),
267            "middleware took {elapsed:?}, expected < 50 ms",
268        );
269    }
270
271    /// Smoke test: the collector must start, tick at least twice without
272    /// panicking, and remain alive until aborted.
273    #[tokio::test]
274    async fn collector_ticks_without_panicking() {
275        let handle = spawn_system_metrics_collector("test".into(), vec![PathBuf::from("/")], Duration::from_millis(50));
276
277        tokio::time::sleep(Duration::from_millis(150)).await;
278
279        assert!(!handle.is_finished(), "collector ended prematurely");
280        handle.abort();
281    }
282
283    #[test]
284    fn test_method_label_standard_methods_are_borrowed() {
285        for (method, expected) in [
286            (Method::GET, "GET"),
287            (Method::POST, "POST"),
288            (Method::PUT, "PUT"),
289            (Method::DELETE, "DELETE"),
290            (Method::PATCH, "PATCH"),
291            (Method::HEAD, "HEAD"),
292            (Method::OPTIONS, "OPTIONS"),
293            (Method::CONNECT, "CONNECT"),
294            (Method::TRACE, "TRACE"),
295        ] {
296            let label = method_label(&method);
297            assert_eq!(label, expected);
298            assert!(
299                matches!(label, Cow::Borrowed(_)),
300                "{method} should be Borrowed, got Owned (allocation in hot path)",
301            );
302        }
303    }
304
305    #[test]
306    fn test_method_label_custom_method_is_owned() {
307        let custom = Method::from_bytes(b"PURGE").unwrap();
308        let label = method_label(&custom);
309        assert_eq!(label, "PURGE");
310        assert!(matches!(label, Cow::Owned(_)));
311    }
312
313    #[test]
314    fn test_status_label_common_codes_are_borrowed() {
315        for code in [
316            200, 201, 204, 301, 302, 304, 400, 401, 403, 404, 409, 422, 500, 502, 503, 504,
317        ] {
318            let label = status_label(code);
319            assert_eq!(label, code.to_string());
320            assert!(
321                matches!(label, Cow::Borrowed(_)),
322                "{code} should be Borrowed, got Owned (allocation in hot path)",
323            );
324        }
325    }
326
327    #[test]
328    fn test_status_label_uncommon_code_is_owned() {
329        let label = status_label(418);
330        assert_eq!(label, "418");
331        assert!(matches!(label, Cow::Owned(_)));
332    }
333
334    /// The middleware must short-circuit on `/metrics` requests (avoiding
335    /// observation loops). We can't easily inspect the global recorder, but
336    /// we can at least verify the path is exercised without panicking.
337    #[tokio::test]
338    async fn middleware_handles_metrics_path() {
339        let svc = ServiceBuilder::new()
340            .layer(PrometheusLayer {
341                service_name: "test".into(),
342            })
343            .service(tower::service_fn(|_req: Request<Body>| async {
344                Ok::<_, Infallible>(Response::builder().status(StatusCode::OK).body(Body::empty()).unwrap())
345            }));
346
347        let response = svc
348            .oneshot(Request::builder().uri("/metrics").body(Body::empty()).unwrap())
349            .await
350            .unwrap();
351        assert_eq!(response.status(), StatusCode::OK);
352    }
353}
api_tools/server/axum/layers/prometheus.rs

api_tools/server/axum/layers/
prometheus.rs
