use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;
static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(1);
fn generate_request_id() -> String {
let seq = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed);
let ts = std::time::SystemTime::now()
.duration_since(std::time::SystemTime::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("{ts:x}-{seq:04x}")
}
pub async fn tracing_middleware(request: Request, next: Next) -> Response {
let method = request.method().clone();
let path = request.uri().path().to_owned();
let request_id = generate_request_id();
let start = Instant::now();
let mut response = next.run(request).await;
let latency_ms = start.elapsed().as_secs_f64() * 1000.0;
let status = response.status().as_u16();
tracing::info!(
method = %method,
path = %path,
status = status,
latency_ms = latency_ms,
request_id = %request_id,
"request completed"
);
if let Ok(val) = HeaderValue::from_str(&request_id) {
response.headers_mut().insert("x-request-id", val);
}
response
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::Body;
use axum::http::{Request as HttpRequest, StatusCode};
use axum::middleware;
use axum::routing::get;
use axum::Router;
use tower::ServiceExt;
fn traced_app() -> Router {
Router::new()
.route("/test", get(|| async { "ok" }))
.layer(middleware::from_fn(tracing_middleware))
}
#[tokio::test]
async fn test_request_id_generated() {
let app = traced_app();
let req = HttpRequest::builder()
.uri("/test")
.body(Body::empty())
.expect("request builder should succeed");
let resp = app
.oneshot(req)
.await
.expect("router should handle request");
assert!(resp.headers().contains_key("x-request-id"));
let rid = resp.headers()["x-request-id"]
.to_str()
.expect("header should be valid string");
assert!(!rid.is_empty());
}
#[tokio::test]
async fn test_request_ids_are_unique() {
let id1 = generate_request_id();
let id2 = generate_request_id();
assert_ne!(id1, id2, "consecutive request IDs must differ");
}
#[tokio::test]
async fn test_response_status_preserved() {
let app = traced_app();
let req = HttpRequest::builder()
.uri("/test")
.body(Body::empty())
.expect("request builder should succeed");
let resp = app
.oneshot(req)
.await
.expect("router should handle request");
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn test_404_still_traced() {
let app = traced_app();
let req = HttpRequest::builder()
.uri("/nonexistent")
.body(Body::empty())
.expect("request builder should succeed");
let resp = app
.oneshot(req)
.await
.expect("router should handle request");
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
assert!(resp.headers().contains_key("x-request-id"));
}
#[tokio::test]
async fn test_latency_is_non_negative() {
let app = traced_app();
let req = HttpRequest::builder()
.uri("/test")
.body(Body::empty())
.expect("request builder should succeed");
let resp = app
.oneshot(req)
.await
.expect("router should handle request");
assert_eq!(resp.status(), StatusCode::OK);
}
}