rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
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;

/// Records request metrics for an axum route.
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
}

/// Creates a router exposing Prometheus metrics at `/metrics`.
pub fn metrics_router(registry: MetricsRegistry) -> Router {
    Router::new().route(
        "/metrics",
        get(move || {
            let registry = registry.clone();
            async move { registry.render_prometheus() }
        }),
    )
}