futures_timeout/
lib.rs

1use core::ops::{Deref, DerefMut};
2use core::pin::Pin;
3use core::task::{Context, Poll};
4use core::time::Duration;
5
6use futures_core::future::FusedFuture;
7use futures_core::stream::FusedStream;
8use futures_core::{Future, Stream};
9use futures_timer::Delay;
10use pin_project::pin_project;
11
12pub trait TimeoutStreamExt: Stream + Sized {
13    /// Requires a [`Stream`] to complete before the specific duration has elapsed.
14    ///
15    /// **Note: If a [`Stream`] returns an item, the timer will reset until `Poll::Ready(None)` is returned**
16    fn timeout(self, duration: Duration) -> Timeout<Self> {
17        Timeout {
18            inner: self,
19            timer: Some(Delay::new(duration)),
20            duration,
21        }
22    }
23}
24
25pub trait TimeoutFutureExt: Future + Sized {
26    /// Requires a [`Future`] to complete before the specific duration has elapsed.
27    fn timeout(self, duration: Duration) -> Timeout<Self> {
28        Timeout {
29            inner: self,
30            timer: Some(Delay::new(duration)),
31            duration,
32        }
33    }
34}
35
36impl<T> TimeoutFutureExt for T where T: Future {}
37impl<T> TimeoutStreamExt for T where T: Stream {}
38
39#[derive(Debug)]
40pub struct TimeoutError;
41impl core::fmt::Display for TimeoutError {
42    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
43        write!(f, "Operation timed out")
44    }
45}
46
47impl core::error::Error for TimeoutError {}
48
49#[derive(Debug)]
50#[pin_project]
51pub struct Timeout<T> {
52    #[pin]
53    inner: T,
54    timer: Option<Delay>,
55    duration: Duration,
56}
57
58impl<T> Timeout<T> {
59    /// Create a timeout from a stream
60    pub fn from_stream(inner: T, duration: Duration) -> Self
61    where
62        T: Stream,
63    {
64        Timeout {
65            inner,
66            timer: Some(Delay::new(duration)),
67            duration,
68        }
69    }
70
71    /// Create a timeout from a future
72    pub fn from_future(inner: T, duration: Duration) -> Self
73    where
74        T: Future,
75    {
76        Timeout {
77            inner,
78            timer: Some(Delay::new(duration)),
79            duration,
80        }
81    }
82
83    /// Consumes Timeout and returns the inner value
84    pub fn into_inner(self) -> T {
85        self.inner
86    }
87}
88
89impl<T> Deref for Timeout<T> {
90    type Target = T;
91    fn deref(&self) -> &Self::Target {
92        &self.inner
93    }
94}
95
96impl<T> DerefMut for Timeout<T> {
97    fn deref_mut(&mut self) -> &mut Self::Target {
98        &mut self.inner
99    }
100}
101
102impl<T: Future> Future for Timeout<T> {
103    type Output = Result<T::Output, TimeoutError>;
104
105    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
106        let this = self.project();
107
108        let Some(timer) = this.timer.as_mut() else {
109            return Poll::Ready(Err(TimeoutError));
110        };
111
112        match this.inner.poll(cx) {
113            Poll::Ready(value) => {
114                this.timer.take();
115                return Poll::Ready(Ok(value));
116            }
117            Poll::Pending => {}
118        }
119
120        core::task::ready!(Pin::new(timer).poll(cx));
121        this.timer.take();
122        Poll::Ready(Err(TimeoutError))
123    }
124}
125
126impl<T: Future> FusedFuture for Timeout<T> {
127    fn is_terminated(&self) -> bool {
128        self.timer.is_none()
129    }
130}
131
132impl<T: Stream> Stream for Timeout<T> {
133    type Item = Result<T::Item, TimeoutError>;
134
135    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
136        let this = self.project();
137
138        let Some(timer) = this.timer.as_mut() else {
139            return Poll::Ready(None);
140        };
141
142        match this.inner.poll_next(cx) {
143            Poll::Ready(Some(value)) => {
144                timer.reset(*this.duration);
145                return Poll::Ready(Some(Ok(value)));
146            }
147            Poll::Ready(None) => {
148                this.timer.take();
149                return Poll::Ready(None);
150            }
151            Poll::Pending => {}
152        }
153
154        core::task::ready!(Pin::new(timer).poll(cx));
155        this.timer.take();
156        Poll::Ready(Some(Err(TimeoutError)))
157    }
158
159    fn size_hint(&self) -> (usize, Option<usize>) {
160        self.inner.size_hint()
161    }
162}
163
164impl<T: Stream> FusedStream for Timeout<T> {
165    fn is_terminated(&self) -> bool {
166        self.timer.is_none()
167    }
168}
169
170#[cfg(test)]
171mod test {
172    use core::time::Duration;
173
174    use futures::{StreamExt, TryStreamExt};
175
176    #[test]
177    fn fut_timeout() {
178        use crate::TimeoutFutureExt;
179        futures::executor::block_on(
180            futures_timer::Delay::new(Duration::from_secs(10)).timeout(Duration::from_secs(5)),
181        )
182        .expect_err("timeout after timer elapsed");
183    }
184
185    #[test]
186    fn stream_timeout() {
187        use crate::TimeoutStreamExt;
188        futures::executor::block_on(async move {
189            let mut st = futures::stream::once(async move {
190                futures_timer::Delay::new(Duration::from_secs(10)).await;
191                0
192            })
193            .timeout(Duration::from_secs(5))
194            .boxed();
195
196            st.try_next()
197                .await
198                .expect_err("timeout after timer elapsed");
199        });
200    }
201}
futures_timeout/lib.rs

futures_timeout/
lib.rs
