1//! Allows a future to execute for a maximum amount of time.
2//!
3//! See [`Timeout`] documentation for more details.
4//!
5//! [`Timeout`]: struct@Timeout
6
7use crate::{
8 runtime::coop,
9 time::{error::Elapsed, sleep_until, Duration, Instant, Sleep},
10 util::trace,
11};
12
13use pin_project_lite::pin_project;
14use std::future::Future;
15use std::pin::Pin;
16use std::task::{self, Poll};
17
18/// Requires a `Future` to complete before the specified duration has elapsed.
19///
20/// If the future completes before the duration has elapsed, then the completed
21/// value is returned. Otherwise, an error is returned and the future is
22/// canceled.
23///
24/// Note that the timeout is checked before polling the future, so if the future
25/// does not yield during execution then it is possible for the future to complete
26/// and exceed the timeout _without_ returning an error.
27///
28/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the
29/// return type of the provided future.
30///
31/// If the provided future completes immediately, then the future returned from
32/// this function is guaranteed to complete immediately with an [`Ok`] variant
33/// no matter the provided duration.
34///
35/// [`Ok`]: std::result::Result::Ok
36/// [`Result`]: std::result::Result
37/// [`Elapsed`]: crate::time::error::Elapsed
38///
39/// # Cancellation
40///
41/// Cancelling a timeout is done by dropping the future. No additional cleanup
42/// or other work is required.
43///
44/// The original future may be obtained by calling [`Timeout::into_inner`]. This
45/// consumes the `Timeout`.
46///
47/// # Examples
48///
49/// Create a new `Timeout` set to expire in 10 milliseconds.
50///
51/// ```rust
52/// use tokio::time::timeout;
53/// use tokio::sync::oneshot;
54///
55/// use std::time::Duration;
56///
57/// # async fn dox() {
58/// let (tx, rx) = oneshot::channel();
59/// # tx.send(()).unwrap();
60///
61/// // Wrap the future with a `Timeout` set to expire in 10 milliseconds.
62/// if let Err(_) = timeout(Duration::from_millis(10), rx).await {
63/// println!("did not receive value within 10 ms");
64/// }
65/// # }
66/// ```
67///
68/// # Panics
69///
70/// This function panics if there is no current timer set.
71///
72/// It can be triggered when [`Builder::enable_time`] or
73/// [`Builder::enable_all`] are not included in the builder.
74///
75/// It can also panic whenever a timer is created outside of a
76/// Tokio runtime. That is why `rt.block_on(sleep(...))` will panic,
77/// since the function is executed outside of the runtime.
78/// Whereas `rt.block_on(async {sleep(...).await})` doesn't panic.
79/// And this is because wrapping the function on an async makes it lazy,
80/// and so gets executed inside the runtime successfully without
81/// panicking.
82///
83/// [`Builder::enable_time`]: crate::runtime::Builder::enable_time
84/// [`Builder::enable_all`]: crate::runtime::Builder::enable_all
85#[track_caller]
86pub fn timeout<F>(duration: Duration, future: F) -> Timeout<F>
87where
88 F: Future,
89{
90 let location = trace::caller_location();
91
92 let deadline = Instant::now().checked_add(duration);
93 let delay = match deadline {
94 Some(deadline) => Sleep::new_timeout(deadline, location),
95 None => Sleep::far_future(location),
96 };
97 Timeout::new_with_delay(future, delay)
98}
99
100/// Requires a `Future` to complete before the specified instant in time.
101///
102/// If the future completes before the instant is reached, then the completed
103/// value is returned. Otherwise, an error is returned.
104///
105/// This function returns a future whose return type is [`Result`]`<T,`[`Elapsed`]`>`, where `T` is the
106/// return type of the provided future.
107///
108/// If the provided future completes immediately, then the future returned from
109/// this function is guaranteed to complete immediately with an [`Ok`] variant
110/// no matter the provided deadline.
111///
112/// [`Ok`]: std::result::Result::Ok
113/// [`Result`]: std::result::Result
114/// [`Elapsed`]: crate::time::error::Elapsed
115///
116/// # Cancellation
117///
118/// Cancelling a timeout is done by dropping the future. No additional cleanup
119/// or other work is required.
120///
121/// The original future may be obtained by calling [`Timeout::into_inner`]. This
122/// consumes the `Timeout`.
123///
124/// # Examples
125///
126/// Create a new `Timeout` set to expire in 10 milliseconds.
127///
128/// ```rust
129/// use tokio::time::{Instant, timeout_at};
130/// use tokio::sync::oneshot;
131///
132/// use std::time::Duration;
133///
134/// # async fn dox() {
135/// let (tx, rx) = oneshot::channel();
136/// # tx.send(()).unwrap();
137///
138/// // Wrap the future with a `Timeout` set to expire 10 milliseconds into the
139/// // future.
140/// if let Err(_) = timeout_at(Instant::now() + Duration::from_millis(10), rx).await {
141/// println!("did not receive value within 10 ms");
142/// }
143/// # }
144/// ```
145pub fn timeout_at<F>(deadline: Instant, future: F) -> Timeout<F>
146where
147 F: Future,
148{
149 let delay = sleep_until(deadline);
150
151 Timeout {
152 value: future,
153 delay,
154 }
155}
156
157pin_project! {
158 /// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at).
159 #[must_use = "futures do nothing unless you `.await` or poll them"]
160 #[derive(Debug)]
161 pub struct Timeout<T> {
162 #[pin]
163 value: T,
164 #[pin]
165 delay: Sleep,
166 }
167}
168
169impl<T> Timeout<T> {
170 pub(crate) fn new_with_delay(value: T, delay: Sleep) -> Timeout<T> {
171 Timeout { value, delay }
172 }
173
174 /// Gets a reference to the underlying value in this timeout.
175 pub fn get_ref(&self) -> &T {
176 &self.value
177 }
178
179 /// Gets a mutable reference to the underlying value in this timeout.
180 pub fn get_mut(&mut self) -> &mut T {
181 &mut self.value
182 }
183
184 /// Consumes this timeout, returning the underlying value.
185 pub fn into_inner(self) -> T {
186 self.value
187 }
188}
189
190impl<T> Future for Timeout<T>
191where
192 T: Future,
193{
194 type Output = Result<T::Output, Elapsed>;
195
196 fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
197 let me = self.project();
198
199 let had_budget_before = coop::has_budget_remaining();
200
201 // First, try polling the future
202 if let Poll::Ready(v) = me.value.poll(cx) {
203 return Poll::Ready(Ok(v));
204 }
205
206 let has_budget_now = coop::has_budget_remaining();
207
208 let delay = me.delay;
209
210 let poll_delay = || -> Poll<Self::Output> {
211 match delay.poll(cx) {
212 Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())),
213 Poll::Pending => Poll::Pending,
214 }
215 };
216
217 if let (true, false) = (had_budget_before, has_budget_now) {
218 // if it is the underlying future that exhausted the budget, we poll
219 // the `delay` with an unconstrained one. This prevents pathological
220 // cases where the underlying future always exhausts the budget and
221 // we never get a chance to evaluate whether the timeout was hit or
222 // not.
223 coop::with_unconstrained(poll_delay)
224 } else {
225 poll_delay()
226 }
227 }
228}
229