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, 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),
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);
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),
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),
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 route = route.unwrap_or("<unknown>").to_owned();
265        let mut state = self
266            .state
267            .lock()
268            .unwrap_or_else(|poisoned| poisoned.into_inner());
269        let route = match self.max_series {
270            Some(max) => state.admit_route(route, max),
271            None => route,
272        };
273        *state
274            .in_flight
275            .entry((method.as_str().to_owned(), route))
276            .or_default() += 1;
277    }
278
279    fn on_response(
280        &self,
281        method: &Method,
282        route: Option<&str>,
283        status: StatusCode,
284        latency: Duration,
285    ) {
286        let is_error = status.is_client_error() || status.is_server_error();
287        self.record_completion(method, route, status.as_u16(), latency, is_error);
288    }
289
290    fn on_error(&self, method: &Method, route: Option<&str>, latency: Duration) {
291        self.record_completion(
292            method,
293            route,
294            StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
295            latency,
296            true,
297        );
298    }
299}
300
301impl PrometheusMetrics {
302    fn record_completion(
303        &self,
304        method: &Method,
305        route: Option<&str>,
306        status: u16,
307        latency: Duration,
308        is_error: bool,
309    ) {
310        if !self.should_record(route) {
311            return;
312        }
313        let method = method.as_str().to_owned();
314        let route = route.unwrap_or("<unknown>").to_owned();
315        let mut state = self
316            .state
317            .lock()
318            .unwrap_or_else(|poisoned| poisoned.into_inner());
319        let route = match self.max_series {
320            Some(max) => state.admit_route(route, max),
321            None => route,
322        };
323        // Update in-flight first so the owned key can then be moved into the
324        // series key without cloning either label.
325        let key = (method, route);
326        if let Some(count) = state.in_flight.get_mut(&key) {
327            *count = count.saturating_sub(1);
328        }
329        let (method, route) = key;
330        let series = state.series.entry((method, route, status)).or_default();
331        series.requests += 1;
332        series.histogram.observe(latency);
333        if is_error {
334            series.errors += 1;
335        }
336    }
337}
338
339#[derive(Clone, Debug, Default)]
340struct PrometheusState {
341    series: BTreeMap<(String, String, u16), StatusSeries>,
342    in_flight: BTreeMap<(String, String), u64>,
343    known_routes: BTreeSet<String>,
344}
345
346/// Counters and histogram for one `(method, route, status)` label set.
347#[derive(Clone, Debug, Default)]
348struct StatusSeries {
349    requests: u64,
350    errors: u64,
351    histogram: DurationHistogram,
352}
353
354impl PrometheusState {
355    /// Returns the label to record for `route`, honoring a cap on the number of
356    /// distinct route labels. Already-admitted routes are returned unchanged;
357    /// once the cap is reached, new labels collapse to `"<overflow>"`. Callers
358    /// with no cap must skip this call entirely (the uncapped path pays nothing).
359    fn admit_route(&mut self, route: String, max_series: usize) -> String {
360        if self.known_routes.contains(&route) {
361            route
362        } else if self.known_routes.len() < max_series {
363            self.known_routes.insert(route.clone());
364            route
365        } else {
366            "<overflow>".to_owned()
367        }
368    }
369}
370
371const HTTP_DURATION_BUCKETS: [f64; 11] = [
372    0.005, 0.010, 0.025, 0.050, 0.100, 0.250, 0.500, 1.000, 2.500, 5.000, 10.000,
373];
374
375/// Rendered `le` labels for [`HTTP_DURATION_BUCKETS`], kept in sync by test.
376const HTTP_DURATION_BUCKET_LABELS: [&str; 11] = [
377    "0.005", "0.01", "0.025", "0.05", "0.1", "0.25", "0.5", "1", "2.5", "5", "10",
378];
379
380#[derive(Clone, Debug, Default)]
381struct DurationHistogram {
382    count: u64,
383    sum: f64,
384    bucket_counts: [u64; HTTP_DURATION_BUCKETS.len()],
385}
386
387impl DurationHistogram {
388    fn observe(&mut self, latency: Duration) {
389        let seconds = latency.as_secs_f64();
390        self.count += 1;
391        self.sum += seconds;
392        for (bucket, count) in HTTP_DURATION_BUCKETS
393            .iter()
394            .zip(self.bucket_counts.iter_mut())
395        {
396            if seconds <= *bucket {
397                *count += 1;
398            }
399        }
400    }
401}
402
403/// Tower layer that invokes [`HttpMetricsHook`] for request lifecycle metrics.
404///
405/// Route labels come from [`Self::route`] when set, then from Axum
406/// [`axum::extract::MatchedPath`], and finally `"<unknown>"`.
407#[derive(Clone, Debug)]
408pub struct MetricsLayer<H> {
409    hook: H,
410    route: Option<Cow<'static, str>>,
411}
412
413impl<H> MetricsLayer<H>
414where
415    H: HttpMetricsHook,
416{
417    /// Creates a metrics layer without a route label.
418    pub fn new(hook: H) -> Self {
419        Self { hook, route: None }
420    }
421
422    /// Adds a stable route label to emitted metrics.
423    ///
424    /// Prefer route patterns such as `"/users/:id"` over concrete paths to keep
425    /// label cardinality bounded.
426    pub fn route(mut self, route: impl Into<Cow<'static, str>>) -> Self {
427        self.route = Some(route.into());
428        self
429    }
430}
431
432impl<S, H> Layer<S> for MetricsLayer<H>
433where
434    H: HttpMetricsHook,
435{
436    type Service = MetricsService<S, H>;
437
438    fn layer(&self, inner: S) -> Self::Service {
439        MetricsService {
440            inner,
441            hook: self.hook.clone(),
442            route: self.route.clone(),
443        }
444    }
445}
446
447/// Service produced by [`MetricsLayer`].
448#[derive(Clone, Debug)]
449pub struct MetricsService<S, H> {
450    inner: S,
451    hook: H,
452    route: Option<Cow<'static, str>>,
453}
454
455impl<S, H, RequestBody, ResponseBody> Service<Request<RequestBody>> for MetricsService<S, H>
456where
457    S: Service<Request<RequestBody>, Response = Response<ResponseBody>> + Send + 'static,
458    S::Future: Send + 'static,
459    S::Error: Send + 'static,
460    H: HttpMetricsHook,
461    RequestBody: Send + 'static,
462    ResponseBody: Send + 'static,
463{
464    type Response = Response<ResponseBody>;
465    type Error = S::Error;
466    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
467
468    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
469        self.inner.poll_ready(cx)
470    }
471
472    fn call(&mut self, request: Request<RequestBody>) -> Self::Future {
473        let method = request.method().clone();
474        let hook = self.hook.clone();
475        let route = match self.route.clone() {
476            Some(route) => RouteLabel::Fixed(route),
477            // MatchedPath wraps an Arc<str>; cloning it avoids copying the
478            // route pattern into a fresh String on every request.
479            None => match request.extensions().get::<axum::extract::MatchedPath>() {
480                Some(path) => RouteLabel::Matched(path.clone()),
481                None => RouteLabel::Unknown,
482            },
483        };
484        hook.on_request(&method, route.as_deref());
485        let started_at = Instant::now();
486        let future = self.inner.call(request);
487
488        Box::pin(async move {
489            match future.await {
490                Ok(response) => {
491                    hook.on_response(
492                        &method,
493                        route.as_deref(),
494                        response.status(),
495                        started_at.elapsed(),
496                    );
497                    Ok(response)
498                }
499                Err(error) => {
500                    hook.on_error(&method, route.as_deref(), started_at.elapsed());
501                    Err(error)
502                }
503            }
504        })
505    }
506}
507
508/// Route label carried across the metrics service future without copying
509/// matched route patterns.
510enum RouteLabel {
511    Fixed(Cow<'static, str>),
512    Matched(axum::extract::MatchedPath),
513    Unknown,
514}
515
516impl RouteLabel {
517    fn as_deref(&self) -> Option<&str> {
518        match self {
519            Self::Fixed(route) => Some(route),
520            Self::Matched(path) => Some(path.as_str()),
521            Self::Unknown => None,
522        }
523    }
524}
525
526fn escape_label(value: &str) -> Cow<'_, str> {
527    if value.contains(['\\', '\n', '"']) {
528        Cow::Owned(
529            value
530                .replace('\\', r"\\")
531                .replace('\n', r"\n")
532                .replace('"', r#"\""#),
533        )
534    } else {
535        Cow::Borrowed(value)
536    }
537}
538
539#[cfg(test)]
540fn format_bucket(bucket: f64) -> String {
541    if bucket.fract() == 0.0 {
542        format!("{bucket:.0}")
543    } else {
544        let formatted = format!("{bucket:.3}");
545        formatted.trim_end_matches('0').to_owned()
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use super::{HTTP_DURATION_BUCKET_LABELS, HTTP_DURATION_BUCKETS, escape_label, format_bucket};
552
553    #[test]
554    fn duration_bucket_labels_match_formatted_buckets() {
555        for (bucket, label) in HTTP_DURATION_BUCKETS
556            .iter()
557            .zip(HTTP_DURATION_BUCKET_LABELS.iter())
558        {
559            assert_eq!(&format_bucket(*bucket), label);
560        }
561    }
562
563    #[test]
564    fn escape_label_borrows_clean_values_and_escapes_special_characters() {
565        assert!(matches!(
566            escape_label("/users/{id}"),
567            std::borrow::Cow::Borrowed("/users/{id}")
568        ));
569        assert_eq!(escape_label("a\\b\nc\"d"), r#"a\\b\nc\"d"#);
570    }
571}