Skip to main content

aster_forge_observability/axum/
mod.rs

1//! Injectable Axum exposition; authentication and paths belong to the product.
2
3#[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}