diff --git a/tower-http/src/on_early_drop/early_drops_as_failures.rs b/tower-http/src/on_early_drop/early_drops_as_failures.rs index 64e8d77ff..708f6b70e 100644 --- a/tower-http/src/on_early_drop/early_drops_as_failures.rs +++ b/tower-http/src/on_early_drop/early_drops_as_failures.rs @@ -4,9 +4,8 @@ use crate::on_early_drop::failure::{BodyDropped, DroppedFailure, FutureDropped}; use crate::on_early_drop::traits::{OnBodyDrop, OnDropCallback, OnFutureDrop}; -use crate::trace::OnFailure; +use crate::trace::{Clock, DefaultClock, OnFailure}; use http::{response, Request, StatusCode}; -use std::time::Instant; use tracing::Span; /// Bridges early-drop events to [`trace::OnFailure`](crate::trace::OnFailure). @@ -25,30 +24,52 @@ use tracing::Span; /// /// [`OnEarlyDropLayer`]: super::OnEarlyDropLayer #[derive(Debug, Clone, Copy)] -pub struct EarlyDropsAsFailures { +pub struct EarlyDropsAsFailures +where + Clk: Clock, +{ on_failure: F, + clock: Clk, } impl EarlyDropsAsFailures { - /// Wrap an [`OnFailure`] implementation. + /// Wrap an [`OnFailure`] implementation with [`DefaultClock`]. pub fn new(on_failure: F) -> Self { - Self { on_failure } + Self { + on_failure, + clock: DefaultClock, + } + } +} + +impl EarlyDropsAsFailures +where + Clk: Clock, +{ + /// Wrap an [`OnFailure`] implementation with a custom [`Clock`]. + pub fn with_clock(on_failure: F, clock: Clk) -> Self { + Self { on_failure, clock } } } /// Future-drop callback produced by [`EarlyDropsAsFailures`]. -pub struct FutureDropFailureCallback { - start: Instant, +pub struct FutureDropFailureCallback +where + Clk: Clock, +{ + start: Clk::Instant, on_failure: F, span: Span, + clock: Clk, } -impl OnDropCallback for FutureDropFailureCallback +impl OnDropCallback for FutureDropFailureCallback where - F: OnFailure + Send + 'static, + Clk: Clock + Send + 'static, + F: OnFailure + Send + 'static, { fn on_drop(mut self) { - let latency = self.start.elapsed(); + let latency = self.clock.elapsed(self.start); let _entered = self.span.enter(); self.on_failure .on_failure(DroppedFailure::Future(FutureDropped), latency, &self.span); @@ -58,26 +79,35 @@ where /// Intermediate produced by [`OnBodyDrop::make_at_call`] for /// [`EarlyDropsAsFailures`], carrying state forward to /// [`OnBodyDrop::make_at_response`]. -pub struct PreResponseBodyDropCallback { - start: Instant, +pub struct PreResponseBodyDropCallback +where + Clk: Clock, +{ + start: Clk::Instant, on_failure: F, span: Span, + clock: Clk, } /// Body-drop callback produced by [`EarlyDropsAsFailures`]. -pub struct BodyDropFailureCallback { - start: Instant, +pub struct BodyDropFailureCallback +where + Clk: Clock, +{ + start: Clk::Instant, on_failure: F, span: Span, status: StatusCode, + clock: Clk, } -impl OnDropCallback for BodyDropFailureCallback +impl OnDropCallback for BodyDropFailureCallback where - F: OnFailure + Send + 'static, + Clk: Clock + Send + 'static, + F: OnFailure + Send + 'static, { fn on_drop(mut self) { - let latency = self.start.elapsed(); + let latency = self.clock.elapsed(self.start); let _entered = self.span.enter(); self.on_failure.on_failure( DroppedFailure::Body(BodyDropped { @@ -89,33 +119,37 @@ where } } -impl OnFutureDrop for EarlyDropsAsFailures +impl OnFutureDrop for EarlyDropsAsFailures where - F: OnFailure + Clone + Send + 'static, + Clk: Clock + Send + 'static, + F: OnFailure + Clone + Send + 'static, { - type Callback = FutureDropFailureCallback; + type Callback = FutureDropFailureCallback; fn make(&mut self, _request: &Request) -> Self::Callback { FutureDropFailureCallback { - start: Instant::now(), + start: self.clock.now(), on_failure: self.on_failure.clone(), span: Span::current(), + clock: self.clock, } } } -impl OnBodyDrop for EarlyDropsAsFailures +impl OnBodyDrop for EarlyDropsAsFailures where - F: OnFailure + Clone + Send + 'static, + Clk: Clock + Send + 'static, + F: OnFailure + Clone + Send + 'static, { - type Intermediate = PreResponseBodyDropCallback; - type Callback = BodyDropFailureCallback; + type Intermediate = PreResponseBodyDropCallback; + type Callback = BodyDropFailureCallback; fn make_at_call(&mut self, _request: &Request) -> Self::Intermediate { PreResponseBodyDropCallback { - start: Instant::now(), + start: self.clock.now(), on_failure: self.on_failure.clone(), span: Span::current(), + clock: self.clock, } } @@ -129,6 +163,7 @@ where on_failure: intermediate.on_failure, span: intermediate.span, status: response_parts.status, + clock: intermediate.clock, } } } diff --git a/tower-http/src/trace/body.rs b/tower-http/src/trace/body.rs index 861e5c46d..8dce3e9eb 100644 --- a/tower-http/src/trace/body.rs +++ b/tower-http/src/trace/body.rs @@ -1,4 +1,7 @@ -use super::{DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, OnBodyChunk, OnEos, OnFailure}; +use super::{ + clock::Clock, DefaultClock, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, OnBodyChunk, + OnEos, OnFailure, +}; use crate::classify::ClassifyEos; use http_body::{Body, Frame}; use pin_project_lite::pin_project; @@ -6,7 +9,6 @@ use std::{ fmt, pin::Pin, task::{ready, Context, Poll}, - time::Instant, }; use tracing::Span; @@ -14,27 +16,36 @@ pin_project! { /// Response body for [`Trace`]. /// /// [`Trace`]: super::Trace - pub struct ResponseBody { + pub struct ResponseBody< + B, + C, + OnBodyChunk = DefaultOnBodyChunk, + OnEos = DefaultOnEos, + OnFailure = DefaultOnFailure, + Clk: Clock = DefaultClock + > { #[pin] pub(crate) inner: B, pub(crate) classify_eos: Option, - pub(crate) on_eos: Option<(OnEos, Instant)>, + pub(crate) on_eos: Option<(OnEos, Clk::Instant)>, pub(crate) on_body_chunk: OnBodyChunk, pub(crate) on_failure: Option, - pub(crate) start: Instant, + pub(crate) start: Clk::Instant, pub(crate) span: Span, + pub(crate) clock: Clk, } } -impl Body - for ResponseBody +impl Body + for ResponseBody where B: Body, B::Error: fmt::Display + 'static, C: ClassifyEos, - OnEosT: OnEos, - OnBodyChunkT: OnBodyChunk, - OnFailureT: OnFailure, + OnEosT: OnEos, + OnBodyChunkT: OnBodyChunk, + OnFailureT: OnFailure, + Clk: Clock, { type Data = B::Data; type Error = B::Error; @@ -47,8 +58,8 @@ where let _guard = this.span.enter(); let result = ready!(this.inner.as_mut().poll_frame(cx)); - let latency = this.start.elapsed(); - *this.start = Instant::now(); + let latency = this.clock.elapsed(*this.start); + *this.start = this.clock.now(); match result { Some(Ok(frame)) => { @@ -70,7 +81,11 @@ where } } if let Some((on_eos, stream_start)) = this.on_eos.take() { - on_eos.on_eos(Some(&trailers), stream_start.elapsed(), this.span); + on_eos.on_eos( + Some(&trailers), + this.clock.elapsed(stream_start), + this.span, + ); } Frame::trailers(trailers) } @@ -89,7 +104,7 @@ where } } if let Some((on_eos, stream_start)) = this.on_eos.take() { - on_eos.on_eos(None, stream_start.elapsed(), this.span); + on_eos.on_eos(None, this.clock.elapsed(stream_start), this.span); } } @@ -114,7 +129,7 @@ where } } if let Some((on_eos, stream_start)) = this.on_eos.take() { - on_eos.on_eos(None, stream_start.elapsed(), this.span); + on_eos.on_eos(None, this.clock.elapsed(stream_start), this.span); } Poll::Ready(None) diff --git a/tower-http/src/trace/clock.rs b/tower-http/src/trace/clock.rs new file mode 100644 index 000000000..0c0f5d784 --- /dev/null +++ b/tower-http/src/trace/clock.rs @@ -0,0 +1,77 @@ +use std::time::{Duration, Instant}; + +/// Extension trait for duration types that need to be formatted as latencies. +/// +/// The [`Clock`] trait's `Duration` associated type must implement this trait, +/// which provides conversion methods used by [`Latency`](super::Latency). +pub trait DurationExt { + /// Returns the duration in seconds as an `f64`. + fn as_secs_f64(&self) -> f64; + /// Returns the duration in whole milliseconds. + fn as_millis(&self) -> u128; + /// Returns the duration in whole microseconds. + fn as_micros(&self) -> u128; + /// Returns the duration in whole nanoseconds. + fn as_nanos(&self) -> u128; +} + +impl DurationExt for Duration { + fn as_secs_f64(&self) -> f64 { + Duration::as_secs_f64(self) + } + fn as_millis(&self) -> u128 { + Duration::as_millis(self) + } + fn as_micros(&self) -> u128 { + Duration::as_micros(self) + } + fn as_nanos(&self) -> u128 { + Duration::as_nanos(self) + } +} + +/// An abstract clock for measuring time. +/// +/// Used by the trace middleware to record latencies and stream durations. +/// The default implementation ([`DefaultClock`]) uses [`std::time::Instant`] +/// and [`std::time::Duration`], but custom clocks can be provided — for +/// example, to mock time in tests or to use a WASM-compatible time source. +pub trait Clock: Copy { + /// An instant in time. + type Instant: Copy + Send; + + /// A duration of time. + type Duration: Copy + DurationExt + Send; + + /// Returns the current instant. + fn now(&self) -> Self::Instant; + + /// Returns the duration elapsed from `instant` to now. + fn elapsed(&self, instant: Self::Instant) -> Self::Duration; +} + +/// The default [`Clock`] implementation, backed by [`std::time::Instant`]. +#[derive(Debug, Clone, Copy, Default)] +pub struct DefaultClock; + +impl DefaultClock { + /// Create a new `DefaultClock`. + pub fn new() -> Self { + Self + } +} + +impl Clock for DefaultClock { + type Instant = Instant; + type Duration = Duration; + + #[inline] + fn now(&self) -> Instant { + Instant::now() + } + + #[inline] + fn elapsed(&self, instant: Instant) -> Duration { + instant.elapsed() + } +} diff --git a/tower-http/src/trace/future.rs b/tower-http/src/trace/future.rs index 6a73b887a..f1312f2a7 100644 --- a/tower-http/src/trace/future.rs +++ b/tower-http/src/trace/future.rs @@ -1,6 +1,6 @@ use super::{ - DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnResponse, OnBodyChunk, OnEos, - OnFailure, OnResponse, ResponseBody, + clock::Clock, DefaultClock, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, + DefaultOnResponse, OnBodyChunk, OnEos, OnFailure, OnResponse, ResponseBody, }; use crate::classify::{ClassifiedResponse, ClassifyResponse}; use http::Response; @@ -10,7 +10,6 @@ use std::{ future::Future, pin::Pin, task::{ready, Context, Poll}, - time::Instant, }; use tracing::Span; @@ -18,7 +17,15 @@ pin_project! { /// Response future for [`Trace`]. /// /// [`Trace`]: super::Trace - pub struct ResponseFuture { + pub struct ResponseFuture< + F, + C, + OnResponse = DefaultOnResponse, + OnBodyChunk = DefaultOnBodyChunk, + OnEos = DefaultOnEos, + OnFailure = DefaultOnFailure, + Clk: Clock = DefaultClock + > { #[pin] pub(crate) inner: F, pub(crate) span: Span, @@ -27,33 +34,35 @@ pin_project! { pub(crate) on_body_chunk: Option, pub(crate) on_eos: Option, pub(crate) on_failure: Option, - pub(crate) start: Instant, + pub(crate) start: Clk::Instant, + pub(crate) clock: Clk, } } -impl Future - for ResponseFuture +impl Future + for ResponseFuture where Fut: Future, E>>, ResBody: Body, ResBody::Error: std::fmt::Display + 'static, E: std::fmt::Display + 'static, C: ClassifyResponse, - OnResponseT: OnResponse, - OnFailureT: OnFailure, - OnBodyChunkT: OnBodyChunk, - OnEosT: OnEos, + OnResponseT: OnResponse, + OnFailureT: OnFailure, + OnBodyChunkT: OnBodyChunk, + OnEosT: OnEos, + Clk: Clock, { type Output = Result< - Response>, + Response>, E, >; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.project(); + let mut this = self.project(); let _guard = this.span.enter(); - let result = ready!(this.inner.poll(cx)); - let latency = this.start.elapsed(); + let result = ready!(this.inner.as_mut().poll(cx)); + let latency = this.clock.elapsed(*this.start); let classifier = this.classifier.take().unwrap(); let on_eos = this.on_eos.take(); @@ -64,6 +73,7 @@ where Ok(res) => { let classification = classifier.classify_response(&res); let start = *this.start; + let clock = *this.clock; this.on_response .take() @@ -80,11 +90,12 @@ where let res = res.map(|body| ResponseBody { inner: body, classify_eos: None, - on_eos: on_eos.zip(Some(Instant::now())), + on_eos: on_eos.zip(Some(this.clock.now())), on_body_chunk, on_failure: Some(on_failure), start, span, + clock, }); Poll::Ready(Ok(res)) @@ -94,11 +105,12 @@ where let res = res.map(|body| ResponseBody { inner: body, classify_eos: Some(classify_eos), - on_eos: on_eos.zip(Some(Instant::now())), + on_eos: on_eos.zip(Some(this.clock.now())), on_body_chunk, on_failure: Some(on_failure), start, span, + clock, }); Poll::Ready(Ok(res)) diff --git a/tower-http/src/trace/layer.rs b/tower-http/src/trace/layer.rs index 21ff321c1..00f28794d 100644 --- a/tower-http/src/trace/layer.rs +++ b/tower-http/src/trace/layer.rs @@ -1,6 +1,6 @@ use super::{ - DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnRequest, - DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, Trace, + DefaultClock, DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, + DefaultOnRequest, DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, Trace, }; use crate::classify::{ GrpcErrorsAsFailures, MakeClassifier, ServerErrorsAsFailures, SharedClassifier, @@ -23,6 +23,7 @@ pub struct TraceLayer< OnBodyChunk = DefaultOnBodyChunk, OnEos = DefaultOnEos, OnFailure = DefaultOnFailure, + Clk = DefaultClock, > { pub(crate) make_classifier: M, pub(crate) make_span: MakeSpan, @@ -31,6 +32,7 @@ pub struct TraceLayer< pub(crate) on_body_chunk: OnBodyChunk, pub(crate) on_eos: OnEos, pub(crate) on_failure: OnFailure, + pub(crate) clock: Clk, } impl TraceLayer { @@ -47,12 +49,13 @@ impl TraceLayer { on_eos: DefaultOnEos::default(), on_body_chunk: DefaultOnBodyChunk::default(), on_response: DefaultOnResponse::default(), + clock: DefaultClock, } } } -impl - TraceLayer +impl + TraceLayer { /// Customize what to do when a request is received. /// @@ -62,7 +65,7 @@ impl pub fn on_request( self, new_on_request: NewOnRequest, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { on_request: new_on_request, on_failure: self.on_failure, @@ -71,6 +74,7 @@ impl make_span: self.make_span, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -82,7 +86,7 @@ impl pub fn on_response( self, new_on_response: NewOnResponse, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { on_response: new_on_response, on_request: self.on_request, @@ -91,6 +95,7 @@ impl on_failure: self.on_failure, make_span: self.make_span, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -102,7 +107,7 @@ impl pub fn on_body_chunk( self, new_on_body_chunk: NewOnBodyChunk, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { on_body_chunk: new_on_body_chunk, on_eos: self.on_eos, @@ -111,6 +116,7 @@ impl make_span: self.make_span, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -122,7 +128,7 @@ impl pub fn on_eos( self, new_on_eos: NewOnEos, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { on_eos: new_on_eos, on_body_chunk: self.on_body_chunk, @@ -131,6 +137,7 @@ impl make_span: self.make_span, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -142,7 +149,7 @@ impl pub fn on_failure( self, new_on_failure: NewOnFailure, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { on_failure: new_on_failure, on_request: self.on_request, @@ -151,6 +158,7 @@ impl make_span: self.make_span, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -163,7 +171,7 @@ impl pub fn make_span_with( self, new_make_span: NewMakeSpan, - ) -> TraceLayer { + ) -> TraceLayer { TraceLayer { make_span: new_make_span, on_request: self.on_request, @@ -172,6 +180,32 @@ impl on_eos: self.on_eos, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, + } + } + + /// Customize which clock to use for timing. + /// + /// `NewClock` is expected to implement [`Clock`]. + /// + /// Defaults to [`DefaultClock`] which uses [`std::time::Instant`]. + /// + /// [`Clock`]: super::Clock + /// [`DefaultClock`]: super::DefaultClock + pub fn with_clock( + self, + new_clock: NewClock, + ) -> TraceLayer + { + TraceLayer { + clock: new_clock, + on_request: self.on_request, + on_failure: self.on_failure, + on_eos: self.on_eos, + on_body_chunk: self.on_body_chunk, + make_span: self.make_span, + on_response: self.on_response, + make_classifier: self.make_classifier, } } } @@ -188,6 +222,7 @@ impl TraceLayer { on_body_chunk: DefaultOnBodyChunk::default(), on_eos: DefaultOnEos::default(), on_failure: DefaultOnFailure::default(), + clock: DefaultClock, } } } @@ -204,12 +239,13 @@ impl TraceLayer { on_body_chunk: DefaultOnBodyChunk::default(), on_eos: DefaultOnEos::default(), on_failure: DefaultOnFailure::default(), + clock: DefaultClock, } } } -impl Layer - for TraceLayer +impl Layer + for TraceLayer where M: Clone, MakeSpan: Clone, @@ -218,8 +254,9 @@ where OnEos: Clone, OnBodyChunk: Clone, OnFailure: Clone, + Clk: Clone, { - type Service = Trace; + type Service = Trace; fn layer(&self, inner: S) -> Self::Service { Trace { @@ -231,6 +268,7 @@ where on_body_chunk: self.on_body_chunk.clone(), on_response: self.on_response.clone(), on_failure: self.on_failure.clone(), + clock: self.clock.clone(), } } } diff --git a/tower-http/src/trace/mod.rs b/tower-http/src/trace/mod.rs index 2c9aca703..f10d45bba 100644 --- a/tower-http/src/trace/mod.rs +++ b/tower-http/src/trace/mod.rs @@ -384,12 +384,13 @@ //! [`ServerErrorsAsFailures`]: crate::classify::ServerErrorsAsFailures //! [`Body::poll_frame`]: http_body::Body::poll_frame -use std::{fmt, time::Duration}; +use std::fmt; use tracing::Level; pub use self::{ body::ResponseBody, + clock::{Clock, DefaultClock, DurationExt}, future::ResponseFuture, layer::TraceLayer, make_span::{DefaultMakeSpan, MakeSpan}, @@ -491,6 +492,7 @@ macro_rules! event_dynamic_lvl { } mod body; +mod clock; mod future; mod layer; mod make_span; @@ -504,12 +506,12 @@ mod service; const DEFAULT_MESSAGE_LEVEL: Level = Level::DEBUG; const DEFAULT_ERROR_LEVEL: Level = Level::ERROR; -struct Latency { +struct Latency { unit: LatencyUnit, - duration: Duration, + duration: D, } -impl fmt::Display for Latency { +impl fmt::Display for Latency { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self.unit { LatencyUnit::Seconds => write!(f, "{} s", self.duration.as_secs_f64()), @@ -529,7 +531,7 @@ mod tests { use http::{HeaderMap, Request, Response}; use once_cell::sync::Lazy; use std::{ - sync::atomic::{AtomicU32, Ordering}, + sync::atomic::{AtomicU32, AtomicU64, Ordering}, time::Duration, }; use tower::{BoxError, Service, ServiceBuilder, ServiceExt}; @@ -852,6 +854,226 @@ mod tests { assert_eq!(1, ON_EOS.load(Ordering::SeqCst), "eos"); } + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + struct DummyDuration(u64); + + impl DurationExt for DummyDuration { + fn as_secs_f64(&self) -> f64 { + self.0 as f64 + } + fn as_millis(&self) -> u128 { + self.0 as u128 + } + fn as_micros(&self) -> u128 { + self.0 as u128 + } + fn as_nanos(&self) -> u128 { + self.0 as u128 + } + } + + #[derive(Debug, Clone, Copy)] + struct DummyClock(&'static AtomicU64); + + impl DummyClock { + fn new() -> Self { + Self(Box::leak(Box::new(AtomicU64::new(0)))) + } + } + + impl Clock for DummyClock { + type Instant = u64; + type Duration = DummyDuration; + + fn now(&self) -> Self::Instant { + self.0.fetch_add(1, Ordering::SeqCst) + } + + fn elapsed(&self, instant: Self::Instant) -> Self::Duration { + DummyDuration(self.0.load(Ordering::SeqCst) - instant) + } + } + + #[test] + fn dummy_clock_increments_and_elapsed() { + let clock = DummyClock::new(); + + let t0 = clock.now(); + assert_eq!(t0, 0, "first now() should return 0"); + let t1 = clock.now(); + assert_eq!(t1, 1, "second now() should return 1"); + let t2 = clock.now(); + assert_eq!(t2, 2, "third now() should return 2"); + + let d = clock.elapsed(t0); + assert_eq!( + d, + DummyDuration(3), + "elapsed from t0 at counter=3 should be 3" + ); + let d = clock.elapsed(t1); + assert_eq!( + d, + DummyDuration(2), + "elapsed from t1 at counter=3 should be 2" + ); + } + + #[test] + fn dummy_clock_latency_display() { + assert_eq!( + Latency { + unit: LatencyUnit::Seconds, + duration: DummyDuration(3), + } + .to_string(), + "3 s" + ); + assert_eq!( + Latency { + unit: LatencyUnit::Millis, + duration: DummyDuration(42), + } + .to_string(), + "42 ms" + ); + assert_eq!( + Latency { + unit: LatencyUnit::Micros, + duration: DummyDuration(7), + } + .to_string(), + "7 μs" + ); + assert_eq!( + Latency { + unit: LatencyUnit::Nanos, + duration: DummyDuration(100), + } + .to_string(), + "100 ns" + ); + } + + #[tokio::test] + async fn trace_with_dummy_clock_on_response() { + static ON_RESPONSE_LATENCY: Lazy = Lazy::new(|| AtomicU64::new(0)); + + let clock = DummyClock::new(); + + let mut svc = ServiceBuilder::new() + .layer( + TraceLayer::new_for_http() + .with_clock(clock) + .on_request(()) + .on_body_chunk(()) + .on_eos(()) + .on_failure(()) + .on_response( + |_res: &Response, latency: DummyDuration, _span: &Span| { + ON_RESPONSE_LATENCY.store(latency.0, Ordering::SeqCst); + }, + ), + ) + .service_fn(echo); + + let _res = svc + .ready() + .await + .unwrap() + .call(Request::new(Body::from("foobar"))) + .await + .unwrap(); + + assert!( + ON_RESPONSE_LATENCY.load(Ordering::SeqCst) > 0, + "on_response should receive a non-zero latency from DummyClock" + ); + } + + #[tokio::test] + async fn trace_with_dummy_clock_streaming() { + static ON_BODY_CHUNK_LATENCY: Lazy = Lazy::new(|| AtomicU64::new(0)); + static ON_EOS_DURATION: Lazy = Lazy::new(|| AtomicU64::new(0)); + + let clock = DummyClock::new(); + + let mut svc = ServiceBuilder::new() + .layer( + TraceLayer::new_for_http() + .with_clock(clock) + .on_request(()) + .on_response(()) + .on_failure(()) + .on_body_chunk(|_chunk: &Bytes, latency: DummyDuration, _span: &Span| { + ON_BODY_CHUNK_LATENCY.store(latency.0, Ordering::SeqCst); + }) + .on_eos( + |_trailers: Option<&HeaderMap>, duration: DummyDuration, _span: &Span| { + ON_EOS_DURATION.store(duration.0, Ordering::SeqCst); + }, + ), + ) + .service_fn(streaming_body); + + let res = svc + .ready() + .await + .unwrap() + .call(Request::new(Body::empty())) + .await + .unwrap(); + + crate::test_helpers::to_bytes(res.into_body()) + .await + .unwrap(); + + assert!( + ON_BODY_CHUNK_LATENCY.load(Ordering::SeqCst) > 0, + "on_body_chunk should receive a non-zero latency from DummyClock" + ); + assert!( + ON_EOS_DURATION.load(Ordering::SeqCst) > 0, + "on_eos should receive a non-zero stream duration from DummyClock" + ); + } + + #[tokio::test] + async fn trace_with_dummy_clock_on_failure() { + static ON_FAILURE_LATENCY: Lazy = Lazy::new(|| AtomicU64::new(0)); + + let clock = DummyClock::new(); + + let mut svc = ServiceBuilder::new() + .layer( + TraceLayer::new_for_http() + .with_clock(clock) + .on_request(()) + .on_response(()) + .on_body_chunk(()) + .on_eos(()) + .on_failure( + |_class: ServerErrorsFailureClass, latency: DummyDuration, _span: &Span| { + ON_FAILURE_LATENCY.store(latency.0, Ordering::SeqCst); + }, + ), + ) + .service_fn(service_error); + + let _err = svc + .ready() + .await + .unwrap() + .call(Request::new(Body::empty())) + .await; + + assert!(_err.is_err(), "service_error should return an error"); + assert!( + ON_FAILURE_LATENCY.load(Ordering::SeqCst) > 0, + "on_failure should receive a non-zero latency from DummyClock" + ); + } + async fn echo(req: Request) -> Result, BoxError> { Ok(Response::new(req.into_body())) } @@ -877,6 +1099,10 @@ mod tests { Ok(Response::new(body)) } + async fn service_error(_req: Request) -> Result, BoxError> { + Err(BoxError::from("test error")) + } + #[derive(Clone)] struct TestClassify { reject: bool, diff --git a/tower-http/src/trace/on_body_chunk.rs b/tower-http/src/trace/on_body_chunk.rs index 543f2a636..26a89ee4a 100644 --- a/tower-http/src/trace/on_body_chunk.rs +++ b/tower-http/src/trace/on_body_chunk.rs @@ -1,4 +1,4 @@ -use std::time::Duration; +use crate::trace::{Clock, DefaultClock}; use tracing::Span; /// Trait used to tell [`Trace`] what to do when a body chunk has been sent. @@ -7,7 +7,7 @@ use tracing::Span; /// `on_body_chunk` callback is called. /// /// [`Trace`]: super::Trace -pub trait OnBodyChunk { +pub trait OnBodyChunk { /// Do the thing. /// /// `latency` is the duration since the response was sent or since the last body chunk as sent. @@ -24,21 +24,21 @@ pub trait OnBodyChunk { /// [hyper]: https://hyper.rs /// [`Bytes`]: https://docs.rs/bytes/latest/bytes/struct.Bytes.html /// [`TraceLayer::make_span_with`]: crate::trace::TraceLayer::make_span_with - fn on_body_chunk(&mut self, chunk: &B, latency: Duration, span: &Span); + fn on_body_chunk(&mut self, chunk: &B, latency: Clk::Duration, span: &Span); } -impl OnBodyChunk for F +impl OnBodyChunk for F where - F: FnMut(&B, Duration, &Span), + F: FnMut(&B, Clk::Duration, &Span), { - fn on_body_chunk(&mut self, chunk: &B, latency: Duration, span: &Span) { + fn on_body_chunk(&mut self, chunk: &B, latency: Clk::Duration, span: &Span) { self(chunk, latency, span) } } -impl OnBodyChunk for () { +impl OnBodyChunk for () { #[inline] - fn on_body_chunk(&mut self, _: &B, _: Duration, _: &Span) {} + fn on_body_chunk(&mut self, _: &B, _: Clk::Duration, _: &Span) {} } /// The default [`OnBodyChunk`] implementation used by [`Trace`]. @@ -58,7 +58,7 @@ impl DefaultOnBodyChunk { } } -impl OnBodyChunk for DefaultOnBodyChunk { +impl OnBodyChunk for DefaultOnBodyChunk { #[inline] - fn on_body_chunk(&mut self, _: &B, _: Duration, _: &Span) {} + fn on_body_chunk(&mut self, _: &B, _: Clk::Duration, _: &Span) {} } diff --git a/tower-http/src/trace/on_eos.rs b/tower-http/src/trace/on_eos.rs index 92abb7ebd..5a7c58a42 100644 --- a/tower-http/src/trace/on_eos.rs +++ b/tower-http/src/trace/on_eos.rs @@ -1,7 +1,10 @@ use super::{Latency, DEFAULT_MESSAGE_LEVEL}; -use crate::{classify::grpc_errors_as_failures::ParsedGrpcStatus, LatencyUnit}; +use crate::{ + classify::grpc_errors_as_failures::ParsedGrpcStatus, + trace::{Clock, DefaultClock}, + LatencyUnit, +}; use http::header::HeaderMap; -use std::time::Duration; use tracing::{Level, Span}; /// Trait used to tell [`Trace`] what to do when a stream closes. @@ -10,7 +13,7 @@ use tracing::{Level, Span}; /// callback is called. /// /// [`Trace`]: super::Trace -pub trait OnEos { +pub trait OnEos { /// Do the thing. /// /// `stream_duration` is the duration since the response was sent. @@ -22,19 +25,19 @@ pub trait OnEos { /// [`Span`]: https://docs.rs/tracing/latest/tracing/span/index.html /// [record]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.record /// [`TraceLayer::make_span_with`]: crate::trace::TraceLayer::make_span_with - fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Duration, span: &Span); + fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Clk::Duration, span: &Span); } -impl OnEos for () { +impl OnEos for () { #[inline] - fn on_eos(self, _: Option<&HeaderMap>, _: Duration, _: &Span) {} + fn on_eos(self, _: Option<&HeaderMap>, _: Clk::Duration, _: &Span) {} } -impl OnEos for F +impl OnEos for F where - F: FnOnce(Option<&HeaderMap>, Duration, &Span), + F: FnOnce(Option<&HeaderMap>, Clk::Duration, &Span), { - fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Duration, span: &Span) { + fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Clk::Duration, span: &Span) { self(trailers, stream_duration, span) } } @@ -83,8 +86,8 @@ impl DefaultOnEos { } } -impl OnEos for DefaultOnEos { - fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Duration, span: &Span) { +impl OnEos for DefaultOnEos { + fn on_eos(self, trailers: Option<&HeaderMap>, stream_duration: Clk::Duration, span: &Span) { let stream_duration = Latency { unit: self.latency_unit, duration: stream_duration, diff --git a/tower-http/src/trace/on_failure.rs b/tower-http/src/trace/on_failure.rs index 4ed67f5bf..79a81b419 100644 --- a/tower-http/src/trace/on_failure.rs +++ b/tower-http/src/trace/on_failure.rs @@ -1,6 +1,9 @@ use super::{Latency, DEFAULT_ERROR_LEVEL}; -use crate::LatencyUnit; -use std::{fmt, time::Duration}; +use crate::{ + trace::{Clock, DefaultClock}, + LatencyUnit, +}; +use std::fmt; use tracing::{Level, Span}; /// Trait used to tell [`Trace`] what to do when a request fails. @@ -9,7 +12,7 @@ use tracing::{Level, Span}; /// `on_failure` callback is called. /// /// [`Trace`]: super::Trace -pub trait OnFailure { +pub trait OnFailure { /// Do the thing. /// /// `latency` is the duration since the request was received. @@ -21,19 +24,29 @@ pub trait OnFailure { /// [`Span`]: https://docs.rs/tracing/latest/tracing/span/index.html /// [record]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.record /// [`TraceLayer::make_span_with`]: crate::trace::TraceLayer::make_span_with - fn on_failure(&mut self, failure_classification: FailureClass, latency: Duration, span: &Span); + fn on_failure( + &mut self, + failure_classification: FailureClass, + latency: Clk::Duration, + span: &Span, + ); } -impl OnFailure for () { +impl OnFailure for () { #[inline] - fn on_failure(&mut self, _: FailureClass, _: Duration, _: &Span) {} + fn on_failure(&mut self, _: FailureClass, _: Clk::Duration, _: &Span) {} } -impl OnFailure for F +impl OnFailure for F where - F: FnMut(FailureClass, Duration, &Span), + F: FnMut(FailureClass, Clk::Duration, &Span), { - fn on_failure(&mut self, failure_classification: FailureClass, latency: Duration, span: &Span) { + fn on_failure( + &mut self, + failure_classification: FailureClass, + latency: Clk::Duration, + span: &Span, + ) { self(failure_classification, latency, span) } } @@ -81,11 +94,16 @@ impl DefaultOnFailure { } } -impl OnFailure for DefaultOnFailure +impl OnFailure for DefaultOnFailure where FailureClass: fmt::Display, { - fn on_failure(&mut self, failure_classification: FailureClass, latency: Duration, span: &Span) { + fn on_failure( + &mut self, + failure_classification: FailureClass, + latency: Clk::Duration, + span: &Span, + ) { let latency = Latency { unit: self.latency_unit, duration: latency, diff --git a/tower-http/src/trace/on_response.rs b/tower-http/src/trace/on_response.rs index a4d2249c5..768896d79 100644 --- a/tower-http/src/trace/on_response.rs +++ b/tower-http/src/trace/on_response.rs @@ -1,7 +1,9 @@ use super::{Latency, DEFAULT_MESSAGE_LEVEL}; -use crate::LatencyUnit; +use crate::{ + trace::{Clock, DefaultClock}, + LatencyUnit, +}; use http::Response; -use std::time::Duration; use tracing::Level; use tracing::Span; @@ -11,7 +13,7 @@ use tracing::Span; /// `on_response` callback is called. /// /// [`Trace`]: super::Trace -pub trait OnResponse { +pub trait OnResponse { /// Do the thing. /// /// `latency` is the duration since the request was received. @@ -23,19 +25,19 @@ pub trait OnResponse { /// [`Span`]: https://docs.rs/tracing/latest/tracing/span/index.html /// [record]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.record /// [`TraceLayer::make_span_with`]: crate::trace::TraceLayer::make_span_with - fn on_response(self, response: &Response, latency: Duration, span: &Span); + fn on_response(self, response: &Response, latency: Clk::Duration, span: &Span); } -impl OnResponse for () { +impl OnResponse for () { #[inline] - fn on_response(self, _: &Response, _: Duration, _: &Span) {} + fn on_response(self, _: &Response, _: Clk::Duration, _: &Span) {} } -impl OnResponse for F +impl OnResponse for F where - F: FnOnce(&Response, Duration, &Span), + F: FnOnce(&Response, Clk::Duration, &Span), { - fn on_response(self, response: &Response, latency: Duration, span: &Span) { + fn on_response(self, response: &Response, latency: Clk::Duration, span: &Span) { self(response, latency, span) } } @@ -101,8 +103,8 @@ impl DefaultOnResponse { } } -impl OnResponse for DefaultOnResponse { - fn on_response(self, response: &Response, latency: Duration, span: &Span) { +impl OnResponse for DefaultOnResponse { + fn on_response(self, response: &Response, latency: Clk::Duration, span: &Span) { let latency = Latency { unit: self.latency_unit, duration: latency, diff --git a/tower-http/src/trace/service.rs b/tower-http/src/trace/service.rs index 1ab4c1f00..fc2850121 100644 --- a/tower-http/src/trace/service.rs +++ b/tower-http/src/trace/service.rs @@ -1,7 +1,8 @@ use super::{ - DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnRequest, - DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, MakeSpan, OnBodyChunk, OnEos, - OnFailure, OnRequest, OnResponse, ResponseBody, ResponseFuture, TraceLayer, + clock::Clock, DefaultClock, DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, + DefaultOnFailure, DefaultOnRequest, DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, + MakeSpan, OnBodyChunk, OnEos, OnFailure, OnRequest, OnResponse, ResponseBody, ResponseFuture, + TraceLayer, }; use crate::classify::{ GrpcErrorsAsFailures, MakeClassifier, ServerErrorsAsFailures, SharedClassifier, @@ -11,7 +12,6 @@ use http_body::Body; use std::{ fmt, task::{Context, Poll}, - time::Instant, }; use tower_service::Service; @@ -31,6 +31,7 @@ pub struct Trace< OnBodyChunk = DefaultOnBodyChunk, OnEos = DefaultOnEos, OnFailure = DefaultOnFailure, + Clk = DefaultClock, > { pub(crate) inner: S, pub(crate) make_classifier: M, @@ -40,6 +41,7 @@ pub struct Trace< pub(crate) on_body_chunk: OnBodyChunk, pub(crate) on_eos: OnEos, pub(crate) on_failure: OnFailure, + pub(crate) clock: Clk, } impl Trace { @@ -57,6 +59,7 @@ impl Trace { on_body_chunk: DefaultOnBodyChunk::default(), on_eos: DefaultOnEos::default(), on_failure: DefaultOnFailure::default(), + clock: DefaultClock, } } @@ -71,8 +74,8 @@ impl Trace { } } -impl - Trace +impl + Trace { define_inner_service_accessors!(); @@ -84,7 +87,7 @@ impl pub fn on_request( self, new_on_request: NewOnRequest, - ) -> Trace { + ) -> Trace { Trace { on_request: new_on_request, inner: self.inner, @@ -94,6 +97,7 @@ impl make_span: self.make_span, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -105,7 +109,7 @@ impl pub fn on_response( self, new_on_response: NewOnResponse, - ) -> Trace { + ) -> Trace { Trace { on_response: new_on_response, inner: self.inner, @@ -115,6 +119,7 @@ impl on_eos: self.on_eos, make_span: self.make_span, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -126,7 +131,7 @@ impl pub fn on_body_chunk( self, new_on_body_chunk: NewOnBodyChunk, - ) -> Trace { + ) -> Trace { Trace { on_body_chunk: new_on_body_chunk, on_eos: self.on_eos, @@ -136,6 +141,7 @@ impl on_request: self.on_request, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -147,7 +153,7 @@ impl pub fn on_eos( self, new_on_eos: NewOnEos, - ) -> Trace { + ) -> Trace { Trace { on_eos: new_on_eos, make_span: self.make_span, @@ -157,6 +163,7 @@ impl on_body_chunk: self.on_body_chunk, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -168,7 +175,7 @@ impl pub fn on_failure( self, new_on_failure: NewOnFailure, - ) -> Trace { + ) -> Trace { Trace { on_failure: new_on_failure, inner: self.inner, @@ -178,6 +185,7 @@ impl on_eos: self.on_eos, on_response: self.on_response, make_classifier: self.make_classifier, + clock: self.clock, } } @@ -190,7 +198,7 @@ impl pub fn make_span_with( self, new_make_span: NewMakeSpan, - ) -> Trace { + ) -> Trace { Trace { make_span: new_make_span, inner: self.inner, @@ -200,6 +208,32 @@ impl on_response: self.on_response, on_eos: self.on_eos, make_classifier: self.make_classifier, + clock: self.clock, + } + } + + /// Customize which clock to use for timing. + /// + /// `NewClock` is expected to implement [`Clock`]. + /// + /// Defaults to [`DefaultClock`] which uses [`std::time::Instant`]. + /// + /// [`Clock`]: super::Clock + /// [`DefaultClock`]: super::DefaultClock + pub fn with_clock( + self, + new_clock: NewClock, + ) -> Trace { + Trace { + clock: new_clock, + on_request: self.on_request, + on_failure: self.on_failure, + on_eos: self.on_eos, + on_body_chunk: self.on_body_chunk, + make_span: self.make_span, + on_response: self.on_response, + make_classifier: self.make_classifier, + inner: self.inner, } } } @@ -228,6 +262,7 @@ impl on_body_chunk: DefaultOnBodyChunk::default(), on_eos: DefaultOnEos::default(), on_failure: DefaultOnFailure::default(), + clock: DefaultClock, } } } @@ -256,6 +291,7 @@ impl on_body_chunk: DefaultOnBodyChunk::default(), on_eos: DefaultOnEos::default(), on_failure: DefaultOnFailure::default(), + clock: DefaultClock, } } } @@ -271,8 +307,9 @@ impl< OnBodyChunkT, OnEosT, MakeSpanT, + Clk, > Service> - for Trace + for Trace where S: Service, Response = Response>, ReqBody: Body, @@ -283,23 +320,31 @@ where M::Classifier: Clone, MakeSpanT: MakeSpan, OnRequestT: OnRequest, - OnResponseT: OnResponse + Clone, - OnBodyChunkT: OnBodyChunk + Clone, - OnEosT: OnEos + Clone, - OnFailureT: OnFailure + Clone, + OnResponseT: OnResponse + Clone, + OnBodyChunkT: OnBodyChunk + Clone, + OnEosT: OnEos + Clone, + OnFailureT: OnFailure + Clone, + Clk: Clock, { type Response = - Response>; + Response>; type Error = S::Error; - type Future = - ResponseFuture; + type Future = ResponseFuture< + S::Future, + M::Classifier, + OnResponseT, + OnBodyChunkT, + OnEosT, + OnFailureT, + Clk, + >; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } fn call(&mut self, req: Request) -> Self::Future { - let start = Instant::now(); + let start = self.clock.now(); let span = self.make_span.make_span(&req); @@ -320,6 +365,7 @@ where on_eos: Some(self.on_eos.clone()), on_failure: Some(self.on_failure.clone()), start, + clock: self.clock, } } }