aster_forge_observability/axum/
mod.rs1#[cfg(feature = "prometheus")]
4pub fn prometheus_router(provider: aster_forge_metrics::MetricsProvider) -> axum::Router {
5 axum::Router::new()
6 .route("/metrics", axum::routing::get(prometheus_metrics))
7 .layer(axum::Extension(provider))
8}
9
10#[cfg(feature = "prometheus")]
11async fn prometheus_metrics(
12 axum::Extension(provider): axum::Extension<aster_forge_metrics::MetricsProvider>,
13) -> axum::response::Response {
14 use axum::{
15 http::{StatusCode, header},
16 response::IntoResponse,
17 };
18 match provider
19 .export_prometheus(std::time::Duration::from_secs(5))
20 .await
21 {
22 Ok(body) => (
23 StatusCode::OK,
24 [
25 (
26 header::CONTENT_TYPE,
27 "text/plain; version=0.0.4; charset=utf-8",
28 ),
29 (header::CACHE_CONTROL, "no-store"),
30 ],
31 body,
32 )
33 .into_response(),
34 Err(_) => (
35 StatusCode::SERVICE_UNAVAILABLE,
36 [(header::CACHE_CONTROL, "no-store")],
37 "metrics export is unavailable",
38 )
39 .into_response(),
40 }
41}
42
43#[cfg(all(test, feature = "prometheus"))]
44mod tests {
45 use super::*;
46 use aster_forge_metrics::{MetricDescriptor, MetricsProvider, ProductCounter};
47 use axum::{
48 body::Body,
49 http::{Request, StatusCode, header},
50 };
51 use tower::ServiceExt;
52
53 #[tokio::test]
54 async fn injected_axum_endpoint_preserves_format_status_and_provider_isolation() {
55 let provider = MetricsProvider::builder()
56 .with_prometheus()
57 .unwrap()
58 .build()
59 .unwrap();
60 let handle = provider
61 .register_batch(&[MetricDescriptor::counter(
62 "test",
63 "requests_total",
64 "Requests.",
65 &[],
66 )])
67 .unwrap()
68 .remove(0);
69 ProductCounter::from_handle(handle).unwrap().inc(&[], 3);
70 let router = prometheus_router(provider.clone());
71 let response = router
72 .clone()
73 .oneshot(
74 Request::builder()
75 .uri("/metrics")
76 .body(Body::empty())
77 .unwrap(),
78 )
79 .await
80 .unwrap();
81 assert_eq!(response.status(), StatusCode::OK);
82 assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
83 assert_eq!(
84 response.headers()[header::CONTENT_TYPE],
85 "text/plain; version=0.0.4; charset=utf-8"
86 );
87 let body = axum::body::to_bytes(response.into_body(), 1024 * 1024)
88 .await
89 .unwrap();
90 assert!(
91 String::from_utf8(body.to_vec())
92 .unwrap()
93 .contains("test_requests_total 3")
94 );
95 provider
96 .shutdown(std::time::Duration::from_secs(1))
97 .await
98 .unwrap();
99 let response = router
100 .oneshot(
101 Request::builder()
102 .uri("/metrics")
103 .body(Body::empty())
104 .unwrap(),
105 )
106 .await
107 .unwrap();
108 assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
109 assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
110 }
111}