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