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
use crate::{Chain, HandleFuture, Middleware, Request, INTERNAL_ERR}; use chrono::prelude::Utc; use futures::FutureExt; use hyper::header::{HeaderName, HeaderValue}; pub(crate) fn dur_to_string(i: i64) -> String { if i < 1000 { format!("{}us", i) } else if i < 1_000_000 { format!("{:.2}ms", (i as f32) / 1000.0) } else { format!("{:.2}s", (i as f32) / 1_000_000.0) } } #[derive(Debug, Clone)] pub struct Runtime { header: HeaderName, } impl Runtime { pub fn new(header: &str) -> Self { Self { header: HeaderName::from_lowercase(header.as_bytes()).unwrap(), } } pub fn default() -> Self { Self::new("x-runtime") } } impl Middleware for Runtime { fn handle<'m>(&'m self, req: &'m mut Request, chain: Chain<'m>) -> HandleFuture<'m> { async move { let start = Utc::now(); let mut response = chain.run(req).await?; let duration = Utc::now().signed_duration_since(start).num_microseconds(); if let Some(dur) = duration { response.headers_mut().insert( self.header.clone(), HeaderValue::from_str(&dur_to_string(dur)).expect(INTERNAL_ERR), ); } Ok(response) } .boxed() } } #[cfg(test)] mod test { use super::{dur_to_string, Runtime}; #[test] fn test_dur_to_string_micro_seconds() { assert_eq!(&dur_to_string(193), "193us"); } #[test] fn test_dur_to_string_milli_seconds() { assert_eq!(&dur_to_string(2940), "2.94ms"); } #[test] fn test_dur_to_string_seconds() { assert_eq!(&dur_to_string(3495773), "3.50s"); } #[test] #[should_panic] fn test_runtime_with_uppercase() { Runtime::new("X-Runtime"); } }