Skip to main content

nidus_http/middleware/
metrics.rs

1use std::{
2    borrow::Cow,
3    collections::{BTreeMap, BTreeSet},
4    fmt::Write as _,
5    future::Future,
6    pin::Pin,
7    sync::{Arc, LazyLock, Mutex},
8    task::{Context, Poll},
9    time::{Duration, Instant},
10};
11
12use axum::{Router, routing::get};
13use http::{Method, Request, Response, StatusCode};
14use tower::{Layer, Service};
15
16/// Creates a metrics hook layer without a stable route label.
17///
18/// The layer will use Axum's [`axum::extract::MatchedPath`] extension when it is
19/// available, otherwise metrics are recorded with route `"<unknown>"`.
20pub fn metrics_layer<H>(hook: H) -> MetricsLayer<H>
21where
22    H: HttpMetricsHook,
23{
24    MetricsLayer::new(hook)
25}
26
27/// Creates a metrics hook layer that records a stable route label.
28///
29/// Use this for route-specific layers when you want stable labels independent
30/// of Axum extension timing.
31pub fn route_metrics_layer<H>(route: impl Into<Cow<'static, str>>, hook: H) -> MetricsLayer<H>
32where
33    H: HttpMetricsHook,
34{
35    MetricsLayer::new(hook).route(route)
36}
37
38/// Backend-neutral hook for recording HTTP request metrics.
39///
40/// Implement this trait to bridge Nidus' middleware lifecycle into a concrete
41/// metrics backend. Hooks are called in-process: one `on_request` before the
42/// inner service, one `on_response` after a response, or `on_error` if the inner
43/// service returns an error before producing a response.
44pub trait HttpMetricsHook: Clone + Send + Sync + 'static {
45    /// Records that a request entered the service.
46    fn on_request(&self, method: &Method, route: Option<&str>);
47
48    /// Records that a response left the service.
49    fn on_response(
50        &self,
51        method: &Method,
52        route: Option<&str>,
53        status: StatusCode,
54        latency: Duration,
55    );
56
57    /// Records that the inner service returned an error before producing a response.
58    fn on_error(&self, _method: &Method, _route: Option<&str>, _latency: Duration) {}
59}
60
61/// In-memory Prometheus-format HTTP metrics collector.
62///
63/// This collector stores counters and bounded duration histograms in process
64/// memory and renders Prometheus text exposition. It is useful for small
65/// services, examples, and tests; it is not a durable metrics store and values
66/// reset on process restart. The default exclusions are `/health/live`,
67/// `/health/ready`, and `/metrics`.
68///
69/// # Label cardinality
70///
71/// By default the collector records every distinct route label it observes, so
72/// the caller is responsible for keeping cardinality bounded. Prefer route
73/// patterns (e.g. `"/users/{id}"`) over concrete paths. To harden against
74/// accidental high-cardinality labels (which would grow memory without bound in
75/// a long-running process), apply [`PrometheusMetrics::with_max_series`]: once
76/// the configured number of distinct route labels has been admitted, every
77/// further distinct label collapses into a single `"<overflow>"` route.
78///
79/// ```
80/// use axum::{Router, routing::get};
81/// use nidus_http::middleware::{PrometheusMetrics, route_metrics_layer};
82/// # async fn show_user() -> &'static str { "user" }
83///
84/// let metrics = PrometheusMetrics::new();
85/// let app = Router::new()
86///     .route("/users/{id}", get(show_user))
87///     .route_layer(route_metrics_layer("/users/{id}", metrics.clone()))
88///     .merge(metrics.routes());
89/// # let _: Router = app;
90/// ```
91#[derive(Clone, Debug)]
92pub struct PrometheusMetrics {
93    state: Arc<Mutex<PrometheusState>>,
94    excluded_routes: Arc<BTreeSet<String>>,
95    max_series: Option<usize>,
96}
97
98impl PrometheusMetrics {
99    /// Creates a Prometheus metrics collector with default internal route exclusions.
100    ///
101    /// The collector is unbounded by default (every distinct route label is
102    /// recorded); use [`Self::with_max_series`] to cap label cardinality.
103    pub fn new() -> Self {
104        Self {
105            state: Arc::new(Mutex::new(PrometheusState::default())),
106            excluded_routes: Arc::new(BTreeSet::from([
107                "/health/live".to_owned(),
108                "/health/ready".to_owned(),
109                "/metrics".to_owned(),
110            ])),
111            max_series: None,
112        }
113    }
114
115    /// Adds a route pattern to exclude from recording.
116    ///
117    /// Match the exact route label emitted by the metrics layer, such as a
118    /// static route supplied to [`route_metrics_layer`] or an Axum matched path.
119    pub fn exclude_route(mut self, route: impl Into<String>) -> Self {
120        Arc::make_mut(&mut self.excluded_routes).insert(route.into());
121        self
122    }
123
124    /// Bounds the number of distinct route labels retained in memory.
125    ///
126    /// The first `max_series` distinct route labels are recorded normally; any
127    /// further distinct label collapses into a single shared `"<overflow>"`
128    /// route. This prevents unbounded memory growth when a layer accidentally
129    /// emits high-cardinality labels (for example concrete request paths) while
130    /// still keeping the already-admitted routes intact. Without this cap the
131    /// collector records every distinct label it observes.
132    pub fn with_max_series(mut self, max_series: usize) -> Self {
133        self.max_series = Some(max_series);
134        self
135    }
136
137    /// Creates a metrics layer backed by this collector.
138    ///
139    /// The layer records request totals, errors, in-flight counts, and bounded
140    /// duration histograms. It does not expose a scrape endpoint; use
141    /// [`Self::routes`] or [`Self::routes_at`] for that.
142    pub fn layer(&self) -> MetricsLayer<Self> {
143        MetricsLayer::new(self.clone())
144    }
145
146    /// Creates a `/metrics` route for this collector.
147    pub fn routes(&self) -> Router {
148        self.routes_at("/metrics")
149    }
150
151    /// Creates a metrics route at a custom path.
152    pub fn routes_at(&self, path: &'static str) -> Router {
153        let metrics = self.clone();
154        Router::new().route(path, get(move || async move { metrics.render() }))
155    }
156
157    /// Renders metrics in Prometheus text exposition format.
158    ///
159    /// The output includes `nidus_http_requests_total`,
160    /// `nidus_http_request_duration_seconds_count`,
161    /// `nidus_http_request_duration_seconds_sum`,
162    /// `nidus_http_in_flight_requests`, and `nidus_http_errors_total`.
163    pub fn render(&self) -> String {
164        let state = self.snapshot();
165        render_prometheus(&state)
166    }
167
168    fn snapshot(&self) -> PrometheusState {
169        self.state
170            .lock()
171            .unwrap_or_else(|poisoned| poisoned.into_inner())
172            .clone()
173    }
174
175    fn should_record(&self, route: Option<&str>) -> bool {
176        route
177            .map(|route| !self.excluded_routes.contains(route))
178            .unwrap_or(true)
179    }
180}
181
182fn render_prometheus(state: &PrometheusState) -> String {
183    // Writes into the output String are infallible; ignore the fmt::Result.
184    let mut output = String::new();
185    output.push_str("# TYPE nidus_http_requests_total counter\n");
186    for ((method, route, status), series) in &state.series {
187        let _ = writeln!(
188            output,
189            "nidus_http_requests_total{{method=\"{}\",route=\"{}\",status=\"{}\"}} {}",
190            escape_label(method.as_str()),
191            escape_label(route),
192            status,
193            series.requests
194        );
195    }
196    output.push_str("# TYPE nidus_http_request_duration_seconds histogram\n");
197    for ((method, route, status), series) in &state.series {
198        let histogram = &series.histogram;
199        let method = escape_label(method.as_str());
200        let route = escape_label(route);
201        for (bucket, count) in HTTP_DURATION_BUCKET_LABELS
202            .iter()
203            .zip(histogram.bucket_counts.iter())
204        {
205            let _ = writeln!(
206                output,
207                "nidus_http_request_duration_seconds_bucket{{method=\"{method}\",route=\"{route}\",status=\"{status}\",le=\"{bucket}\"}} {count}",
208            );
209        }
210        let _ = writeln!(
211            output,
212            "nidus_http_request_duration_seconds_bucket{{method=\"{method}\",route=\"{route}\",status=\"{status}\",le=\"+Inf\"}} {}",
213            histogram.count
214        );
215        let _ = writeln!(
216            output,
217            "nidus_http_request_duration_seconds_count{{method=\"{method}\",route=\"{route}\",status=\"{status}\"}} {}",
218            histogram.count
219        );
220        let _ = writeln!(
221            output,
222            "nidus_http_request_duration_seconds_sum{{method=\"{method}\",route=\"{route}\",status=\"{status}\"}} {:.6}",
223            histogram.sum
224        );
225    }
226    output.push_str("# TYPE nidus_http_in_flight_requests gauge\n");
227    for ((method, route), count) in &state.in_flight {
228        let _ = writeln!(
229            output,
230            "nidus_http_in_flight_requests{{method=\"{}\",route=\"{}\"}} {}",
231            escape_label(method.as_str()),
232            escape_label(route),
233            count
234        );
235    }
236    output.push_str("# TYPE nidus_http_errors_total counter\n");
237    for ((method, route, status), series) in &state.series {
238        if series.errors == 0 {
239            continue;
240        }
241        let _ = writeln!(
242            output,
243            "nidus_http_errors_total{{method=\"{}\",route=\"{}\",status=\"{}\"}} {}",
244            escape_label(method.as_str()),
245            escape_label(route),
246            status,
247            series.errors
248        );
249    }
250    output
251}
252
253impl Default for PrometheusMetrics {
254    fn default() -> Self {
255        Self::new()
256    }
257}
258
259impl HttpMetricsHook for PrometheusMetrics {
260    fn on_request(&self, method: &Method, route: Option<&str>) {
261        if !self.should_record(route) {
262            return;
263        }
264        let mut state = self
265            .state
266            .lock()
267            .unwrap_or_else(|poisoned| poisoned.into_inner());
268        let route = state.intern_route(route.unwrap_or("<unknown>"), self.max_series);
269        *state.in_flight.entry((method.clone(), route)).or_default() += 1;
270    }
271
272    fn on_response(
273        &self,
274        method: &Method,
275        route: Option<&str>,
276        status: StatusCode,
277        latency: Duration,
278    ) {
279        let is_error = status.is_client_error() || status.is_server_error();
280        self.record_completion(method, route, status.as_u16(), latency, is_error);
281    }
282
283    fn on_error(&self, method: &Method, route: Option<&str>, latency: Duration) {
284        self.record_completion(
285            method,
286            route,
287            StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
288            latency,
289            true,
290        );
291    }
292}
293
294impl PrometheusMetrics {
295    fn record_completion(
296        &self,
297        method: &Method,
298        route: Option<&str>,
299        status: u16,
300        latency: Duration,
301        is_error: bool,
302    ) {
303        if !self.should_record(route) {
304            return;
305        }
306        let mut state = self
307            .state
308            .lock()
309            .unwrap_or_else(|poisoned| poisoned.into_inner());
310        let route = state.intern_route(route.unwrap_or("<unknown>"), self.max_series);
311        // Update in-flight first so the owned key can then be moved into the
312        // series key without cloning either label.
313        let key = (method.clone(), route);
314        if let Some(count) = state.in_flight.get_mut(&key) {
315            *count = count.saturating_sub(1);
316        }
317        let (method, route) = key;
318        let series = state.series.entry((method, route, status)).or_default();
319        series.requests += 1;
320        series.histogram.observe(latency);
321        if is_error {
322            series.errors += 1;
323        }
324    }
325}
326
327#[derive(Clone, Debug, Default)]
328struct PrometheusState {
329    series: BTreeMap<(Method, Arc<str>, u16), StatusSeries>,
330    in_flight: BTreeMap<(Method, Arc<str>), u64>,
331    known_routes: BTreeSet<Arc<str>>,
332}
333
334static OVERFLOW_ROUTE: LazyLock<Arc<str>> = LazyLock::new(|| Arc::from("<overflow>"));
335
336/// Counters and histogram for one `(method, route, status)` label set.
337#[derive(Clone, Debug, Default)]
338struct StatusSeries {
339    requests: u64,
340    errors: u64,
341    histogram: DurationHistogram,
342}
343
344impl PrometheusState {
345    /// Interns stable route labels so repeated request/response observations
346    /// clone an `Arc` instead of allocating fresh `String` keys. When a cap is
347    /// configured, labels beyond it collapse into the shared overflow label.
348    fn intern_route(&mut self, route: &str, max_series: Option<usize>) -> Arc<str> {
349        if let Some(route) = self.known_routes.get(route) {
350            return Arc::clone(route);
351        }
352        if max_series.is_some_and(|max| self.known_routes.len() >= max) {
353            return Arc::clone(&OVERFLOW_ROUTE);
354        }
355
356        let route: Arc<str> = Arc::from(route);
357        self.known_routes.insert(Arc::clone(&route));
358        route
359    }
360}
361
362const HTTP_DURATION_BUCKETS: [f64; 11] = [
363    0.005, 0.010, 0.025, 0.050, 0.100, 0.250, 0.500, 1.000, 2.500, 5.000, 10.000,
364];
365
366/// Rendered `le` labels for [`HTTP_DURATION_BUCKETS`], kept in sync by test.
367const HTTP_DURATION_BUCKET_LABELS: [&str; 11] = [
368    "0.005", "0.01", "0.025", "0.05", "0.1", "0.25", "0.5", "1", "2.5", "5", "10",
369];
370
371#[derive(Clone, Debug, Default)]
372struct DurationHistogram {
373    count: u64,
374    sum: f64,
375    bucket_counts: [u64; HTTP_DURATION_BUCKETS.len()],
376}
377
378impl DurationHistogram {
379    fn observe(&mut self, latency: Duration) {
380        let seconds = latency.as_secs_f64();
381        self.count += 1;
382        self.sum += seconds;
383        for (bucket, count) in HTTP_DURATION_BUCKETS
384            .iter()
385            .zip(self.bucket_counts.iter_mut())
386        {
387            if seconds <= *bucket {
388                *count += 1;
389            }
390        }
391    }
392}
393
394/// Tower layer that invokes [`HttpMetricsHook`] for request lifecycle metrics.
395///
396/// Route labels come from [`Self::route`] when set, then from Axum
397/// [`axum::extract::MatchedPath`], and finally `"<unknown>"`.
398#[derive(Clone, Debug)]
399pub struct MetricsLayer<H> {
400    hook: H,
401    route: Option<Cow<'static, str>>,
402}
403
404impl<H> MetricsLayer<H>
405where
406    H: HttpMetricsHook,
407{
408    /// Creates a metrics layer without a route label.
409    pub fn new(hook: H) -> Self {
410        Self { hook, route: None }
411    }
412
413    /// Adds a stable route label to emitted metrics.
414    ///
415    /// Prefer route patterns such as `"/users/:id"` over concrete paths to keep
416    /// label cardinality bounded.
417    pub fn route(mut self, route: impl Into<Cow<'static, str>>) -> Self {
418        self.route = Some(route.into());
419        self
420    }
421}
422
423impl<S, H> Layer<S> for MetricsLayer<H>
424where
425    H: HttpMetricsHook,
426{
427    type Service = MetricsService<S, H>;
428
429    fn layer(&self, inner: S) -> Self::Service {
430        MetricsService {
431            inner,
432            hook: self.hook.clone(),
433            route: self.route.clone(),
434        }
435    }
436}
437
438/// Service produced by [`MetricsLayer`].
439#[derive(Clone, Debug)]
440pub struct MetricsService<S, H> {
441    inner: S,
442    hook: H,
443    route: Option<Cow<'static, str>>,
444}
445
446impl<S, H, RequestBody, ResponseBody> Service<Request<RequestBody>> for MetricsService<S, H>
447where
448    S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
449    S::Future: Send + 'static,
450    S::Error: Send + 'static,
451    H: HttpMetricsHook,
452    RequestBody: Send + 'static,
453    ResponseBody: Send + 'static,
454{
455    type Response = Response<ResponseBody>;
456    type Error = S::Error;
457    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
458
459    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
460        self.inner.poll_ready(cx)
461    }
462
463    fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
464        let method = request.method().clone();
465        let hook = self.hook.clone();
466        let route = match self.route.clone() {
467            Some(route) => RouteLabel::Fixed(route),
468            // MatchedPath wraps an Arc<str>; cloning it avoids copying the
469            // route pattern into a fresh String on every request.
470            None => match request.extensions().get::<axum::extract::MatchedPath>() {
471                Some(path) => RouteLabel::Matched(path.clone()),
472                None => RouteLabel::Unknown,
473            },
474        };
475        hook.on_request(&method, route.as_deref());
476        let started_at = Instant::now();
477        let future = self.inner.call(request);
478
479        Box::pin(async move {
480            match future.await {
481                Ok(response) => {
482                    hook.on_response(
483                        &method,
484                        route.as_deref(),
485                        response.status(),
486                        started_at.elapsed(),
487                    );
488                    Ok(response)
489                }
490                Err(error) => {
491                    hook.on_error(&method, route.as_deref(), started_at.elapsed());
492                    Err(error)
493                }
494            }
495        })
496    }
497}
498
499/// Route label carried across the metrics service future without copying
500/// matched route patterns.
501enum RouteLabel {
502    Fixed(Cow<'static, str>),
503    Matched(axum::extract::MatchedPath),
504    Unknown,
505}
506
507impl RouteLabel {
508    fn as_deref(&self) -> Option<&str> {
509        match self {
510            Self::Fixed(route) => Some(route),
511            Self::Matched(path) => Some(path.as_str()),
512            Self::Unknown => None,
513        }
514    }
515}
516
517fn escape_label(value: &str) -> Cow<'_, str> {
518    if value.contains(['\\', '\n', '"']) {
519        Cow::Owned(
520            value
521                .replace('\\', r"\\")
522                .replace('\n', r"\n")
523                .replace('"', r#"\""#),
524        )
525    } else {
526        Cow::Borrowed(value)
527    }
528}
529
530#[cfg(test)]
531fn format_bucket(bucket: f64) -> String {
532    if bucket.fract() == 0.0 {
533        format!("{bucket:.0}")
534    } else {
535        let formatted = format!("{bucket:.3}");
536        formatted.trim_end_matches('0').to_owned()
537    }
538}
539
540#[cfg(test)]
541mod tests {
542    use std::sync::Arc;
543
544    use super::{
545        HTTP_DURATION_BUCKET_LABELS, HTTP_DURATION_BUCKETS, PrometheusState, escape_label,
546        format_bucket,
547    };
548
549    #[test]
550    fn prometheus_state_reuses_interned_and_overflow_route_labels() {
551        let mut state = PrometheusState::default();
552        let first = state.intern_route("/users/{id}", None);
553        let repeated = state.intern_route("/users/{id}", None);
554        assert!(Arc::ptr_eq(&first, &repeated));
555
556        let first_overflow = state.intern_route("/projects/{id}", Some(1));
557        let second_overflow = state.intern_route("/teams/{id}", Some(1));
558        assert_eq!(&*first_overflow, "<overflow>");
559        assert!(Arc::ptr_eq(&first_overflow, &second_overflow));
560    }
561
562    #[test]
563    fn duration_bucket_labels_match_formatted_buckets() {
564        for (bucket, label) in HTTP_DURATION_BUCKETS
565            .iter()
566            .zip(HTTP_DURATION_BUCKET_LABELS.iter())
567        {
568            assert_eq!(&format_bucket(*bucket), label);
569        }
570    }
571
572    #[test]
573    fn escape_label_borrows_clean_values_and_escapes_special_characters() {
574        assert!(matches!(
575            escape_label("/users/{id}"),
576            std::borrow::Cow::Borrowed("/users/{id}")
577        ));
578        assert_eq!(escape_label("a\\b\nc\"d"), r#"a\\b\nc\"d"#);
579    }
580}