use std::time::Instant;
use axum::{
Router,
body::Body,
extract::{MatchedPath, State},
http::Request,
middleware::Next,
response::IntoResponse,
routing::get,
};
use crate::observability::{CorrelationContext, HttpMetricLabels, MetricsRegistry};
use tracing::Instrument;
pub async fn record_metrics_middleware(
State(registry): State<MetricsRegistry>,
mut request: Request<Body>,
next: Next,
) -> impl IntoResponse {
let method = request.method().to_string();
let route = request
.extensions()
.get::<MatchedPath>()
.map(|path| path.as_str().to_string())
.unwrap_or_else(|| "unknown".to_string());
let correlation = CorrelationContext::from_http_headers(
None,
method.clone(),
Some(&route),
request.headers(),
);
let request_id = correlation.request_id().unwrap_or("").to_string();
let incoming_traceparent = correlation.traceparent();
let started = Instant::now();
let span = tracing::info_span!(
"rs_zero.http.request",
http.method = %method,
http.route = %route,
service = "unknown",
transport = "http",
route = %route,
method = %method,
request_id = %request_id,
traceparent = incoming_traceparent.unwrap_or(""),
status = tracing::field::Empty,
trace_id = tracing::field::Empty,
span_id = tracing::field::Empty
);
registry.increment_http_in_flight();
if !request_id.is_empty() {
request
.extensions_mut()
.insert(crate::observability::CurrentRequestId(request_id.clone()));
}
#[cfg(feature = "otlp")]
crate::observability::set_span_parent_from_headers(&span, request.headers());
let response_future = next.run(request).instrument(span.clone());
let response = if request_id.is_empty() {
response_future.await
} else {
with_current_request_id(request_id.clone(), response_future).await
};
registry.decrement_http_in_flight();
let status = response.status().as_u16();
span.record("status", tracing::field::display(status));
let correlation = correlation.with_status(status.to_string());
if let Some(trace_id) = correlation.trace_id() {
span.record("trace_id", tracing::field::display(trace_id));
}
if let Some(span_id) = correlation.span_id() {
span.record("span_id", tracing::field::display(span_id));
}
registry.record_http_request(
HttpMetricLabels::new(method, route, status),
started.elapsed(),
);
response
}
async fn with_current_request_id<T>(
request_id: String,
future: impl std::future::Future<Output = T>,
) -> T {
crate::layer::context::scope_request_id(request_id, future).await
}
pub fn metrics_router(registry: MetricsRegistry) -> Router {
Router::new().route(
"/metrics",
get(move || {
let registry = registry.clone();
async move { registry.render_prometheus() }
}),
)
}