monoio/time/
timeout.rs

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 std::{
8    future::Future,
9    pin::Pin,
10    task::{self, Poll},
11};
12
13use pin_project_lite::pin_project;
14
15use crate::time::{error::Elapsed, sleep_until, Duration, Instant, Sleep};
16
17/// Require a `Future` to complete before the specified duration has elapsed.
18///
19/// If the future completes before the duration has elapsed, then the completed
20/// value is returned. Otherwise, an error is returned and the future is
21/// canceled.
22///
23/// # Cancelation
24///
25/// Cancelling a timeout is done by dropping the future. No additional cleanup
26/// or other work is required.
27///
28/// The original future may be obtained by calling [`Timeout::into_inner`]. This
29/// consumes the `Timeout`.
30pub fn timeout<T>(duration: Duration, future: T) -> Timeout<T>
31where
32    T: Future,
33{
34    let deadline = Instant::now().checked_add(duration);
35    let delay = match deadline {
36        Some(deadline) => Sleep::new_timeout(deadline),
37        None => Sleep::far_future(),
38    };
39    Timeout::new_with_delay(future, delay)
40}
41
42/// Require a `Future` to complete before the specified instant in time.
43///
44/// If the future completes before the instant is reached, then the completed
45/// value is returned. Otherwise, an error is returned.
46///
47/// # Cancelation
48///
49/// Cancelling a timeout is done by dropping the future. No additional cleanup
50/// or other work is required.
51///
52/// The original future may be obtained by calling [`Timeout::into_inner`]. This
53/// consumes the `Timeout`.
54pub fn timeout_at<T>(deadline: Instant, future: T) -> Timeout<T>
55where
56    T: Future,
57{
58    let delay = sleep_until(deadline);
59
60    Timeout {
61        value: future,
62        delay,
63    }
64}
65
66pin_project! {
67    /// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at).
68    #[must_use = "futures do nothing unless you `.await` or poll them"]
69    #[derive(Debug)]
70    pub struct Timeout<T> {
71        #[pin]
72        value: T,
73        #[pin]
74        delay: Sleep,
75    }
76}
77
78impl<T> Timeout<T> {
79    pub(crate) fn new_with_delay(value: T, delay: Sleep) -> Timeout<T> {
80        Timeout { value, delay }
81    }
82
83    /// Gets a reference to the underlying value in this timeout.
84    pub fn get_ref(&self) -> &T {
85        &self.value
86    }
87
88    /// Gets a mutable reference to the underlying value in this timeout.
89    pub fn get_mut(&mut self) -> &mut T {
90        &mut self.value
91    }
92
93    /// Consumes this timeout, returning the underlying value.
94    pub fn into_inner(self) -> T {
95        self.value
96    }
97}
98
99impl<T> Future for Timeout<T>
100where
101    T: Future,
102{
103    type Output = Result<T::Output, Elapsed>;
104
105    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
106        let me = self.project();
107
108        // First, try polling the future
109        if let Poll::Ready(v) = me.value.poll(cx) {
110            return Poll::Ready(Ok(v));
111        }
112
113        // Now check the timer
114        match me.delay.poll(cx) {
115            Poll::Ready(()) => Poll::Ready(Err(Elapsed::new())),
116            Poll::Pending => Poll::Pending,
117        }
118    }
119}