futures_time/stream/
timeout.rs

1use std::io;
2use std::pin::Pin;
3
4use pin_project_lite::pin_project;
5
6use core::task::{Context, Poll};
7use futures_core::stream::Stream;
8
9use crate::{future::Timer, utils};
10
11pin_project! {
12    /// A stream with timeout time set
13    ///
14    /// This `struct` is created by the [`timeout`] method on [`StreamExt`]. See its
15    /// documentation for more.
16    ///
17    /// [`timeout`]: crate::stream::StreamExt::timeout
18    /// [`StreamExt`]: crate::stream::StreamExt
19    #[derive(Debug)]
20    #[must_use = "streams do nothing unless polled or .awaited"]
21    pub struct Timeout<S, D> {
22        #[pin]
23        stream: S,
24        #[pin]
25        deadline: D,
26    }
27}
28
29impl<S, D> Timeout<S, D> {
30    pub(crate) fn new(stream: S, deadline: D) -> Self {
31        Self { stream, deadline }
32    }
33}
34
35impl<S: Stream, D: Timer> Stream for Timeout<S, D> {
36    type Item = io::Result<S::Item>;
37
38    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
39        let mut this = self.project();
40
41        let r = match this.stream.poll_next(cx) {
42            Poll::Ready(Some(v)) => Poll::Ready(Some(Ok(v))),
43            Poll::Ready(None) => Poll::Ready(None),
44            Poll::Pending => match this.deadline.as_mut().poll(cx) {
45                Poll::Ready(_) => Poll::Ready(Some(Err(utils::timeout_err("stream timed out")))),
46                Poll::Pending => return Poll::Pending,
47            },
48        };
49
50        this.deadline.as_mut().reset_timer();
51
52        r
53    }
54}