pingora_web 0.1.1

use crate::core::router::Handler;
use crate::{
    core::{Request, Response},
    middleware::Middleware,
};
use async_trait::async_trait;
use std::sync::Arc;
use tracing::{Instrument, info};

/// Tracing middleware that creates a span for each request with request_id context
/// This ensures all tracing calls within the request have the request_id automatically included
#[derive(Clone)]
pub struct TracingMiddleware;

impl TracingMiddleware {
    pub fn new() -> Self {
        Self
    }
}

impl Default for TracingMiddleware {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Middleware for TracingMiddleware {
    async fn handle(&self, req: Request, next: Arc<dyn Handler>) -> Response {
        let request_id = req
            .headers()
            .get("x-request-id")
            .and_then(|v| v.to_str().ok())
            .unwrap_or("")
            .to_string();
        let method = req.method().as_str().to_string();
        let path = req.path().to_string();

        // Create a span for this request with structured fields
        let span = tracing::info_span!(
            "request",
            request_id = request_id.as_str(),
            method = method.as_str(),
            path = path,
            status = tracing::field::Empty,
            latency_ms = tracing::field::Empty,
        );

        // Clone span for use in both the closure and instrument
        let span_for_record = span.clone();

        // Execute the rest of the middleware chain within this span
        async move {
            // Log request start
            info!(
                method = method.as_str(),
                path = path,
                request_id = request_id.as_str(),
                "Request started"
            );

            let start_time = std::time::Instant::now();

            let res = next.handle(req).await;

            let elapsed_ms = start_time.elapsed().as_millis();

            // Record the response status and latency in the span
            span_for_record.record("status", res.status.as_u16());
            span_for_record.record("latency_ms", elapsed_ms);

            // Log the request completion
            info!(
                status = res.status.as_u16(),
                latency_ms = elapsed_ms,
                "Request completed"
            );

            res
        }
        .instrument(span)
        .await
    }
}