feat(trace): add abstract Clock trait supporting custom timing impl

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<F> {
+pub struct EarlyDropsAsFailures<F, Clk = DefaultClock>
+where
+    Clk: Clock,
+{
     on_failure: F,
+    clock: Clk,
 }
 
 impl<F> EarlyDropsAsFailures<F> {
-    /// 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<F, Clk> EarlyDropsAsFailures<F, Clk>
+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<F> {
-    start: Instant,
+pub struct FutureDropFailureCallback<F, Clk = DefaultClock>
+where
+    Clk: Clock,
+{
+    start: Clk::Instant,
     on_failure: F,
     span: Span,
+    clock: Clk,
 }
 
-impl<F> OnDropCallback for FutureDropFailureCallback<F>
+impl<F, Clk> OnDropCallback for FutureDropFailureCallback<F, Clk>
 where
-    F: OnFailure<DroppedFailure> + Send + 'static,
+    Clk: Clock + Send + 'static,
+    F: OnFailure<DroppedFailure, Clk> + 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<F> {
-    start: Instant,
+pub struct PreResponseBodyDropCallback<F, Clk = DefaultClock>
+where
+    Clk: Clock,
+{
+    start: Clk::Instant,
     on_failure: F,
     span: Span,
+    clock: Clk,
 }
 
 /// Body-drop callback produced by [`EarlyDropsAsFailures`].
-pub struct BodyDropFailureCallback<F> {
-    start: Instant,
+pub struct BodyDropFailureCallback<F, Clk = DefaultClock>
+where
+    Clk: Clock,
+{
+    start: Clk::Instant,
     on_failure: F,
     span: Span,
     status: StatusCode,
+    clock: Clk,
 }
 
-impl<F> OnDropCallback for BodyDropFailureCallback<F>
+impl<F, Clk> OnDropCallback for BodyDropFailureCallback<F, Clk>
 where
-    F: OnFailure<DroppedFailure> + Send + 'static,
+    Clk: Clock + Send + 'static,
+    F: OnFailure<DroppedFailure, Clk> + 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<F, ReqB> OnFutureDrop<ReqB> for EarlyDropsAsFailures<F>
+impl<F, Clk, ReqB> OnFutureDrop<ReqB> for EarlyDropsAsFailures<F, Clk>
 where
-    F: OnFailure<DroppedFailure> + Clone + Send + 'static,
+    Clk: Clock + Send + 'static,
+    F: OnFailure<DroppedFailure, Clk> + Clone + Send + 'static,
 {
-    type Callback = FutureDropFailureCallback<F>;
+    type Callback = FutureDropFailureCallback<F, Clk>;
 
     fn make(&mut self, _request: &Request<ReqB>) -> Self::Callback {
         FutureDropFailureCallback {
-            start: Instant::now(),
+            start: self.clock.now(),
             on_failure: self.on_failure.clone(),
             span: Span::current(),
+            clock: self.clock,
         }
     }
 }
 
-impl<F, ReqB> OnBodyDrop<ReqB> for EarlyDropsAsFailures<F>
+impl<F, Clk, ReqB> OnBodyDrop<ReqB> for EarlyDropsAsFailures<F, Clk>
 where
-    F: OnFailure<DroppedFailure> + Clone + Send + 'static,
+    Clk: Clock + Send + 'static,
+    F: OnFailure<DroppedFailure, Clk> + Clone + Send + 'static,
 {
-    type Intermediate = PreResponseBodyDropCallback<F>;
-    type Callback = BodyDropFailureCallback<F>;
+    type Intermediate = PreResponseBodyDropCallback<F, Clk>;
+    type Callback = BodyDropFailureCallback<F, Clk>;
 
     fn make_at_call(&mut self, _request: &Request<ReqB>) -> 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,
         }
     }
 }

tower-http/src/trace/body.rs

@@ -1,40 +1,51 @@
-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;
 use std::{
     fmt,
     pin::Pin,
     task::{ready, Context, Poll},
-    time::Instant,
 };
 use tracing::Span;
 
 pin_project! {
     /// Response body for [`Trace`].
     ///
     /// [`Trace`]: super::Trace
-    pub struct ResponseBody<B, C, OnBodyChunk = DefaultOnBodyChunk, OnEos = DefaultOnEos, OnFailure = DefaultOnFailure> {
+    pub struct ResponseBody<
+        B,
+        C,
+        OnBodyChunk = DefaultOnBodyChunk,
+        OnEos = DefaultOnEos,
+        OnFailure = DefaultOnFailure,
+        Clk: Clock = DefaultClock
+    > {
         #[pin]
         pub(crate) inner: B,
         pub(crate) classify_eos: Option<C>,
-        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<OnFailure>,
-        pub(crate) start: Instant,
+        pub(crate) start: Clk::Instant,
         pub(crate) span: Span,
+        pub(crate) clock: Clk,
     }
 }
 
-impl<B, C, OnBodyChunkT, OnEosT, OnFailureT> Body
-    for ResponseBody<B, C, OnBodyChunkT, OnEosT, OnFailureT>
+impl<B, C, OnBodyChunkT, OnEosT, OnFailureT, Clk> Body
+    for ResponseBody<B, C, OnBodyChunkT, OnEosT, OnFailureT, Clk>
 where
     B: Body,
     B::Error: fmt::Display + 'static,
     C: ClassifyEos,
-    OnEosT: OnEos,
-    OnBodyChunkT: OnBodyChunk<B::Data>,
-    OnFailureT: OnFailure<C::FailureClass>,
+    OnEosT: OnEos<Clk>,
+    OnBodyChunkT: OnBodyChunk<B::Data, Clk>,
+    OnFailureT: OnFailure<C::FailureClass, Clk>,
+    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)

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()
+    }
+}