/ src / tracing_extra.rs
tracing_extra.rs
 1  use std::fmt;
 2  use std::sync::atomic::{AtomicU64, Ordering};
 3  use std::sync::Arc;
 4  
 5  use axum::body::Body;
 6  use axum::extract::ConnectInfo;
 7  use axum::http::Request;
 8  use axum::middleware::Next;
 9  use axum::response::IntoResponse;
10  use axum::Extension;
11  use axum_listener::DualAddr;
12  use hyper::{Method, StatusCode, Uri, Version};
13  
14  pub use radicle_term::ansi::Paint;
15  
16  #[derive(Clone)]
17  pub struct RequestId(Arc<AtomicU64>);
18  
19  impl RequestId {
20      pub fn new() -> RequestId {
21          RequestId(Arc::new(0.into()))
22      }
23  
24      pub fn next(&mut self) -> u64 {
25          self.0.fetch_add(1, Ordering::SeqCst)
26      }
27  }
28  
29  #[derive(Clone)]
30  pub struct TracingInfo {
31      pub connect_info: ConnectInfo<DualAddr>,
32      pub method: Method,
33      pub version: Version,
34      pub uri: Uri,
35  }
36  
37  pub struct ColoredStatus(pub StatusCode);
38  
39  impl fmt::Display for ColoredStatus {
40      fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41          match self.0.as_u16() {
42              200..=299 => write!(f, "{}", Paint::green(self.0)),
43              300..=399 => write!(f, "{}", Paint::blue(self.0)),
44              400..=499 => write!(f, "{}", Paint::red(self.0)),
45              _ => write!(f, "{}", Paint::yellow(self.0)),
46          }
47      }
48  }
49  
50  pub async fn tracing_middleware(request: Request<Body>, next: Next) -> impl IntoResponse {
51      let connect_info = request
52          .extensions()
53          .get::<ConnectInfo<DualAddr>>()
54          .unwrap()
55          .clone();
56  
57      let method = request.method().clone();
58      let version = request.version();
59      let uri = request.uri().clone();
60  
61      let tracing_info = TracingInfo {
62          connect_info,
63          method,
64          version,
65          uri,
66      };
67  
68      let response = next.run(request).await;
69  
70      (Extension(tracing_info), response)
71  }