at-jet 0.7.2

//! Request logging middleware

use {axum::http::Request,
     std::{future::Future,
           pin::Pin,
           task::{Context,
                  Poll},
           time::Instant},
     tower::{Layer,
             Service},
     tracing::{info,
               warn}};

/// Request logging layer
#[derive(Clone)]
pub struct RequestLoggingLayer;

impl<S> Layer<S> for RequestLoggingLayer {
  type Service = RequestLoggingMiddleware<S>;

  fn layer(&self, inner: S) -> Self::Service {
    RequestLoggingMiddleware { inner }
  }
}

/// Request logging middleware service
#[derive(Clone)]
pub struct RequestLoggingMiddleware<S> {
  inner: S,
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestLoggingMiddleware<S>
where
  S: Service<Request<ReqBody>, Response = axum::response::Response<ResBody>> + Clone + Send + 'static,
  S::Future: Send,
  ReqBody: Send + 'static,
  ResBody: Send + 'static,
{
  type Error = S::Error;
  type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
  type Response = S::Response;

  fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
    self.inner.poll_ready(cx)
  }

  fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
    let method = req.method().clone();
    let uri = req.uri().clone();
    let start = Instant::now();

    let mut inner = self.inner.clone();

    Box::pin(async move {
      let response = inner.call(req).await?;
      let duration = start.elapsed();
      let status = response.status();

      if status.is_success() {
        info!(
          method = %method,
          uri = %uri,
          status = %status.as_u16(),
          duration_ms = %duration.as_millis(),
          "Request completed"
        );
      } else {
        warn!(
          method = %method,
          uri = %uri,
          status = %status.as_u16(),
          duration_ms = %duration.as_millis(),
          "Request failed"
        );
      }

      Ok(response)
    })
  }
}