axum_layer 0.1.1

axum_layer
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use std::fmt::Debug;

use axum::http::Request;
use futures::{future::Map, FutureExt};
use tower::Service;

#[derive(Clone)]
pub struct Log;

impl<S> tower::Layer<S> for Log {
  type Service = ResponseTimeService<S>;

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

#[derive(Clone)]
pub struct ResponseTimeService<S> {
  inner: S,
}

impl<S, B> Service<Request<B>> for ResponseTimeService<S>
where
  S::Response: Debug,
  S::Error: Debug,
  S: Service<Request<B>>,
  B: Send + Debug,
{
  type Response = S::Response;
  type Error = S::Error;
  type Future =
    Map<S::Future, impl FnOnce(Result<S::Response, S::Error>) -> Result<S::Response, S::Error>>;

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

  fn call(&mut self, req: Request<B>) -> Self::Future {
    // #[cfg(not(feature = "stackdriver"))]
    let start_time: u64 = coarsetime::Clock::now_since_epoch().as_millis();
    // #[cfg(not(feature = "stackdriver"))]
    let method = req.method().to_owned();
    // #[cfg(not(feature = "stackdriver"))]
    let url = req.uri().to_owned();
    // #[cfg(not(feature = "stackdriver"))]
    // tracing::info!("{method} {url}");

    // #[cfg(feature = "stackdriver")]
    // let user_agent;
    // #[cfg(feature = "stackdriver")]
    // {
    //   let headers = req.headers();
    //   user_agent = match headers.get("user-agent") {
    //     Some(t) => Some(t.to_str().unwrap().to_owned()),
    //     None => None,
    //   };
    // }

    self.inner.call(req).map(move |response| {
      // #[cfg(not(feature = "stackdriver"))]
      {
        // tracing::error!(&response);
        let latency =
          (coarsetime::Clock::now_since_epoch().as_millis() - start_time) as f32 / 1000.0;
        tracing::info!("{} {} {}s", method, url, latency)
      }
      response
    })
  }
}

// struct TimingMiddleware<S> {
//   inner: S,
// }
//
// impl<S> TimingMiddleware<S> {
//   fn new(service: S) -> Self {
//     Self { inner: service }
//   }
// }
//
// impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for TimingMiddleware<S>
// where
//   S: Service<Request<ReqBody>, Response = Response<ResBody>>,
//   S::Error: Into<crate::BoxError>,
//   S::Future: Send + 'static,
//   ReqBody: Send + 'static,
//   ResBody: Send + 'static,
// {
//   type Response = S::Response;
//   type Error = S::Error;
//   type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
//
//   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 start_time = Instant::now();
//     let mut inner = self.inner.clone();
//
//     Box::pin(async move {
//       let response = inner.call(req).await?;
//       // let elapsed = start_time.elapsed();
//       info!("Request duration: {:?}", elapsed);
//       Ok(response)
//     })
//   }
// }
//
// // Middleware layer
// struct TimingLayer;
//
// impl<S> tower::Layer<S> for TimingLayer {
//   type Service = TimingMiddleware<S>;
//
//   fn layer(&self, service: S) -> Self::Service {
//     TimingMiddleware::new(service)
//   }
// }