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
16pub fn metrics_layer<H>(hook: H) -> MetricsLayer<H>
21where
22 H: HttpMetricsHook,
23{
24 MetricsLayer::new(hook)
25}
26
27pub 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
38pub trait HttpMetricsHook: Clone + Send + Sync + 'static {
45 fn on_request(&self, method: &Method, route: Option<&str>);
47
48 fn on_response(
50 &self,
51 method: &Method,
52 route: Option<&str>,
53 status: StatusCode,
54 latency: Duration,
55 );
56
57 fn on_error(&self, _method: &Method, _route: Option<&str>, _latency: Duration) {}
59}
60
61#[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 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 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 pub fn with_max_series(mut self, max_series: usize) -> Self {
133 self.max_series = Some(max_series);
134 self
135 }
136
137 pub fn layer(&self) -> MetricsLayer<Self> {
143 MetricsLayer::new(self.clone())
144 }
145
146 pub fn routes(&self) -> Router {
148 self.routes_at("/metrics")
149 }
150
151 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 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 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 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#[derive(Clone, Debug, Default)]
338struct StatusSeries {
339 requests: u64,
340 errors: u64,
341 histogram: DurationHistogram,
342}
343
344impl PrometheusState {
345 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
366const 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#[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 pub fn new(hook: H) -> Self {
410 Self { hook, route: None }
411 }
412
413 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#[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 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
499enum 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}