use colored::Colorize; use http::{Request, Response, method::Method}; use pin_project::pin_project; use std::fmt::Debug; use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::time::Instant; use tower::{Layer, Service}; #[derive(Clone)] pub struct LoggerLayer {} impl LoggerLayer { pub fn new() -> Self { LoggerLayer {} } } impl Layer for LoggerLayer { type Service = Logger; fn layer(&self, inner: S) -> Self::Service { Logger::new(inner) } } #[derive(Clone)] pub struct Logger { inner: S, } impl Logger { pub fn new(inner: S) -> Self { Logger { inner } } } impl Service> for Logger where S: Service, Response = Response>, ResBody: Default, >>::Error: Debug, { type Response = S::Response; type Error = S::Error; type Future = LoggerFuture; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } fn call(&mut self, req: Request) -> Self::Future { let method = req.method().clone(); let path = req.uri().path().to_string(); LoggerFuture { future: self.inner.call(req), method, path, start: Instant::now(), } } } #[pin_project] pub struct LoggerFuture { #[pin] future: F, method: Method, path: String, start: Instant, } impl Future for LoggerFuture where F: Future, E>>, B: Default, E: Debug, { type Output = Result, E>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); let res: F::Output = match this.future.poll(cx) { Poll::Ready(res) => res, Poll::Pending => return Poll::Pending, }; if this.path.starts_with("/__tuono/data") { return Poll::Ready(res); } let status_code = res.as_ref().unwrap().status(); println!( " {} {} {} in {}ms", this.method, this.path, status_code.as_str().green(), this.start.elapsed().as_millis() ); Poll::Ready(res) } }