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
60    /// Create a timeout for Stream or Future.
61    pub fn new(inner: T, duration: Duration) -> Self {
62        Timeout {
63            inner,
64            timer: Some(Delay::new(duration)),
65            duration,
66        }
67    }
68
69    /// Create a timeout from a stream
70    pub fn from_stream(inner: T, duration: Duration) -> Self
71    where
72        T: Stream,
73    {
74        Timeout {
75            inner,
76            timer: Some(Delay::new(duration)),
77            duration,
78        }
79    }
80
81    /// Create a timeout from a future
82    pub fn from_future(inner: T, duration: Duration) -> Self
83    where
84        T: Future,
85    {
86        Timeout {
87            inner,
88            timer: Some(Delay::new(duration)),
89            duration,
90        }
91    }
92
93    /// Consumes Timeout and returns the inner value
94    pub fn into_inner(self) -> T {
95        self.inner
96    }
97}
98
99impl<T> Deref for Timeout<T> {
100    type Target = T;
101    fn deref(&self) -> &Self::Target {
102        &self.inner
103    }
104}
105
106impl<T> DerefMut for Timeout<T> {
107    fn deref_mut(&mut self) -> &mut Self::Target {
108        &mut self.inner
109    }
110}
111
112impl<T: Future> Future for Timeout<T> {
113    type Output = Result<T::Output, TimeoutError>;
114
115    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
116        let this = self.project();
117
118        let Some(timer) = this.timer.as_mut() else {
119            return Poll::Ready(Err(TimeoutError));
120        };
121
122        match this.inner.poll(cx) {
123            Poll::Ready(value) => {
124                this.timer.take();
125                return Poll::Ready(Ok(value));
126            }
127            Poll::Pending => {}
128        }
129
130        core::task::ready!(Pin::new(timer).poll(cx));
131        this.timer.take();
132        Poll::Ready(Err(TimeoutError))
133    }
134}
135
136impl<T: Future> FusedFuture for Timeout<T> {
137    fn is_terminated(&self) -> bool {
138        self.timer.is_none()
139    }
140}
141
142impl<T: Stream> Stream for Timeout<T> {
143    type Item = Result<T::Item, TimeoutError>;
144
145    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
146        let this = self.project();
147
148        let Some(timer) = this.timer.as_mut() else {
149            return Poll::Ready(None);
150        };
151
152        match this.inner.poll_next(cx) {
153            Poll::Ready(Some(value)) => {
154                timer.reset(*this.duration);
155                return Poll::Ready(Some(Ok(value)));
156            }
157            Poll::Ready(None) => {
158                this.timer.take();
159                return Poll::Ready(None);
160            }
161            Poll::Pending => {}
162        }
163
164        core::task::ready!(Pin::new(timer).poll(cx));
165        this.timer.take();
166        Poll::Ready(Some(Err(TimeoutError)))
167    }
168
169    fn size_hint(&self) -> (usize, Option<usize>) {
170        self.inner.size_hint()
171    }
172}
173
174impl<T: Stream> FusedStream for Timeout<T> {
175    fn is_terminated(&self) -> bool {
176        self.timer.is_none()
177    }
178}
179
180#[cfg(test)]
181mod test {
182    use core::time::Duration;
183
184    use futures::{StreamExt, TryStreamExt};
185
186    #[test]
187    fn fut_timeout() {
188        use crate::TimeoutFutureExt;
189        futures::executor::block_on(
190            futures_timer::Delay::new(Duration::from_secs(10)).timeout(Duration::from_secs(5)),
191        )
192        .expect_err("timeout after timer elapsed");
193    }
194
195    #[test]
196    fn stream_timeout() {
197        use crate::TimeoutStreamExt;
198        futures::executor::block_on(async move {
199            let mut st = futures::stream::once(async move {
200                futures_timer::Delay::new(Duration::from_secs(10)).await;
201                0
202            })
203            .timeout(Duration::from_secs(5))
204            .boxed();
205
206            st.try_next()
207                .await
208                .expect_err("timeout after timer elapsed");
209        });
210    }
211}
futures_timeout/lib.rs

futures_timeout/
lib.rs
