1 files changed, 64 insertions, 11 deletions

diff --git a/src/time/timeout.rs b/src/time/timeout.rs
index 6725caa..3bb98ea 100644
--- a/src/time/timeout.rs
+++ b/src/time/timeout.rs
@@ -5,6 +5,7 @@
 //! [`Timeout`]: struct@Timeout
 
 use crate::{
+    runtime::coop,
     time::{error::Elapsed, sleep_until, Duration, Instant, Sleep},
     util::trace,
 };
@@ -20,7 +21,17 @@ use std::task::{self, Poll};
 /// value is returned. Otherwise, an error is returned and the future is
 /// canceled.
 ///
-/// # Cancelation
+/// Note that the timeout is checked before polling the future, so if the future
+/// does not yield during execution then it is possible for the future to complete
+/// and exceed the timeout _without_ returning an error.
+///
+/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the
+/// return type of the provided future.
+///
+/// [`Result`]: std::result::Result
+/// [`Elapsed`]: crate::time::error::Elapsed
+///
+/// # Cancellation
 ///
 /// Cancelling a timeout is done by dropping the future. No additional cleanup
 /// or other work is required.
@@ -48,10 +59,28 @@ use std::task::{self, Poll};
 /// }
 /// # }
 /// ```
-#[cfg_attr(tokio_track_caller, track_caller)]
-pub fn timeout<T>(duration: Duration, future: T) -> Timeout<T>
+///
+/// # Panics
+///
+/// This function panics if there is no current timer set.
+///
+/// It can be triggered when [`Builder::enable_time`] or
+/// [`Builder::enable_all`] are not included in the builder.
+///
+/// It can also panic whenever a timer is created outside of a
+/// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic,
+/// since the function is executed outside of the runtime.
+/// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic.
+/// And this is because wrapping the function on an async makes it lazy,
+/// and so gets executed inside the runtime successfully without
+/// panicking.
+///
+/// [`Builder::enable_time`]: crate::runtime::Builder::enable_time
+/// [`Builder::enable_all`]: crate::runtime::Builder::enable_all
+#[track_caller]
+pub fn timeout<F>(duration: Duration, future: F) -> Timeout<F>
 where
-    T: Future,
+    F: Future,
 {
     let location = trace::caller_location();
 
@@ -68,7 +97,13 @@ where
 /// If the future completes before the instant is reached, then the completed
 /// value is returned. Otherwise, an error is returned.
 ///
-/// # Cancelation
+/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the
+/// return type of the provided future.
+///
+/// [`Result`]: std::result::Result
+/// [`Elapsed`]: crate::time::error::Elapsed
+///
+/// # Cancellation
 ///
 /// Cancelling a timeout is done by dropping the future. No additional cleanup
 /// or other work is required.
@@ -97,9 +132,9 @@ where
 /// }
 /// # }
 /// ```
-pub fn timeout_at<T>(deadline: Instant, future: T) -> Timeout<T>
+pub fn timeout_at<F>(deadline: Instant, future: F) -> Timeout<F>
 where
-    T: Future,
+    F: Future,
 {
     let delay = sleep_until(deadline);
 
@@ -151,15 +186,33 @@ where
     fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
         let me = self.project();
 
+        let had_budget_before = coop::has_budget_remaining();
+
         // First, try polling the future
         if let Poll::Ready(v) = me.value.poll(cx) {
             return Poll::Ready(Ok(v));
         }
 
-        // Now check the timer
-        match me.delay.poll(cx) {
-            Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())),
-            Poll::Pending => Poll::Pending,
+        let has_budget_now = coop::has_budget_remaining();
+
+        let delay = me.delay;
+
+        let poll_delay = || -> Poll<Self::Output> {
+            match delay.poll(cx) {
+                Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())),
+                Poll::Pending => Poll::Pending,
+            }
+        };
+
+        if let (true, false) = (had_budget_before, has_budget_now) {
+            // if it is the underlying future that exhausted the budget, we poll
+            // the `delay` with an unconstrained one. This prevents pathological
+            // cases where the underlying future always exhausts the budget and
+            // we never get a chance to evaluate whether the timeout was hit or
+            // not.
+            coop::with_unconstrained(poll_delay)
+        } else {
+            poll_delay()
         }
     }
 }