Axum Tower Logging Middleware

Instrument Axum with a tower middleware layer that measures request durations, stamps request ids, and logs completions.

RustResilience

Code

snippet
// logging middleware tower layer for axum
use axum::{
    body::Body,
    extract::Request,
    http::{HeaderValue, StatusCode},
    middleware::{self, Next},
    response::Response,
    routing::get,
    Router,
};
use std::time::Instant;
use tower_http::request_id::{MakeRequestUuid, SetRequestIdLayer, XRequestIdName};
use tower_http::trace::TraceLayer;
use tracing::{info, info_span};

async fn timing_middleware(request: Request, next: Next) -> Response {
    let started = Instant::now();
    let request_id = request
        .headers()
        .get("x-request-id")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("-")
        .to_string();

    let method = request.method().clone();
    let path = request.uri().path().to_string();

    let response = next.run(request).await;

    info!(
        request_id = %request_id,
        method = %method,
        path = %path,
        status = %response.status().as_u16(),
        elapsed_ms = started.elapsed().as_millis() as u64,
        "request completed"
    );

    if response.status().is_server_error() {
        tracing::warn!(path = %path, "server error served");
    }
    response
}

async fn handler() -> StatusCode {
    StatusCode::OK
}

async fn fail() -> StatusCode {
    StatusCode::BAD_GATEWAY
}

#[tokio::main]
async fn main() {
    let app: Router = Router::new()
        .route("/ok", get(handler))
        .route("/boom", get(fail))
        .layer(middleware::from_fn(timing_middleware))
        .layer(TraceLayer::new_for_http().make_span_with(|request: &Request<Body>| {
            let request_id = request
                .headers()
                .get(XRequestIdName::default())
                .cloned()
                .unwrap_or_else(|| HeaderValue::from_static("-"));
            info_span!("http", request_id = ?request_id)
        }))
        .layer(SetRequestIdLayer::new(
            XRequestIdName::default(),
            MakeRequestUuid,
        ));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}