introduce Timed to track how long a future is alive for

Author: Joel Scarfone

tokio-util/src/time/mod.rs

@@ -10,7 +10,7 @@
 
 use std::future::Future;
 use std::time::Duration;
-use tokio::time::Timeout;
+use tokio::time::{Timed, Timeout};
 
 mod wheel;
 
@@ -43,6 +43,36 @@ pub trait FutureExt: Future {
     {
         tokio::time::timeout(timeout, self)
     }
+
+    /// A wrapper around [`tokio::time::Timed`], with the advantage that it is easier to write
+    /// fluent call chains.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    ///     /// use tokio::{sync::oneshot, time::Duration};
+    /// use tokio_util::time::FutureExt;
+    /// use tokio::sync::oneshot;
+    /// use tokio::time::Duration;
+    ///
+    /// # async fn dox() {
+    /// let (tx, rx) = oneshot::channel::<()>();
+    ///
+    /// tokio::task::spawn(async move {
+    ///     tokio::time::sleep(Duration::from_millis(5)).await;
+    ///     tx.send(()).unwrap();
+    /// });
+    ///
+    /// let (res, duration) = rx.timed().await;
+    /// println!("Received {res:?} after {duration:?}");
+    /// # }
+    /// ```
+    fn timed(self) -> Timed<Self>
+    where
+        Self: Sized,
+    {
+        tokio::time::timed(self)
+    }
 }
 
 impl<T: Future + ?Sized> FutureExt for T {}

tokio/src/time/mod.rs

@@ -101,6 +101,9 @@ pub use interval::{interval, interval_at, Interval, MissedTickBehavior};
 mod sleep;
 pub use sleep::{sleep, sleep_until, Sleep};
 
+mod timed;
+pub use timed::{timed, Timed};
+
 mod timeout;
 #[doc(inline)]
 pub use timeout::{timeout, timeout_at, Timeout};

tokio/src/time/timed.rs

@@ -0,0 +1,80 @@
+use std::{
+    future::{Future, IntoFuture},
+    pin::Pin,
+    task::Poll,
+};
+
+use pin_project_lite::pin_project;
+
+use crate::time::Instant;
+
+/// Constructs a `Timed` future that wraps an underlying `Future`. The `Timed` future will return
+/// the result of the underlying future as well as the duration the `Timed` was constructed for.
+///
+/// /// # Examples
+///
+/// Track how long we waited on a oneshot rx.
+///
+/// ```rust
+/// use tokio::time::timed;
+/// use tokio::sync::oneshot;
+///
+/// # async fn dox() {
+/// let (tx, rx) = oneshot::channel();
+/// # tx.send(()).unwrap();
+///
+/// // Wrap the future with a `Timed` to see how long we were at this step.
+/// let (_received, duration) = timed(rx).await;
+/// println!("Received a value after waiting {duration:?}");
+/// # }
+/// ```
+pub fn timed<F>(future: F) -> Timed<F::IntoFuture>
+where
+    F: IntoFuture,
+{
+    Timed {
+        inner: future.into_future(),
+        start: Instant::now(),
+    }
+}
+
+pin_project! {
+    /// A helper future that wraps an inner future to also return how long the `Timed` was constructed.
+    /// If constructed right before an await, it will give accurate timing information on the
+    /// underlying future.
+    ///
+    /// # Examples
+    ///
+    /// Time something in a generic async workflow.
+    ///
+    ///  ```
+    ///  use tokio::time::timed;
+    ///
+    ///  async fn i_wish_i_knew_how_long_something_took_in_here() {
+    ///    let (_, elapsed) = timed(async {
+    ///       // do async work
+    ///    }).await;
+    ///  }
+    ///  ```
+    pub struct Timed<T> {
+        #[pin]
+        inner: T,
+        start: Instant,
+    }
+}
+
+impl<T> Future for Timed<T>
+where
+    T: Future,
+{
+    type Output = (T::Output, std::time::Duration);
+
+    fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        if let Poll::Ready(inner) = this.inner.poll(cx) {
+            Poll::Ready((inner, this.start.elapsed()))
+        } else {
+            Poll::Pending
+        }
+    }
+}

tokio/tests/time_timed.rs

@@ -0,0 +1,20 @@
+use tokio::time::{self, timed, Duration};
+
+#[tokio::test(start_paused = true)]
+async fn accuracy() {
+    let durations = vec![
+        Duration::from_millis(50),
+        Duration::from_millis(100),
+        Duration::from_millis(200),
+    ];
+    for expected in durations {
+        let ((), elapsed) = timed(time::sleep(expected)).await;
+        assert_eq!(elapsed, expected);
+    }
+}
+
+#[tokio::test(start_paused = true)]
+async fn immediate_future() {
+    let ((), elapsed) = timed(std::future::ready(())).await;
+    assert_eq!(elapsed, Duration::from_millis(0));
+}