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
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),
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 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#[derive(Clone, Debug, Default)]
348struct StatusSeries {
349 requests: u64,
350 errors: u64,
351 histogram: DurationHistogram,
352}
353
354impl PrometheusState {
355 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
375const 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#[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 pub fn new(hook: H) -> Self {
419 Self { hook, route: None }
420 }
421
422 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#[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 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
508enum 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}