Merge pull request #348 from k-nasa/add_stream_timeout

Add stream timeout

Author: yoshuawuyts

src/stream/stream/mod.rs

@@ -93,13 +93,16 @@ use std::marker::PhantomData;
 
 cfg_unstable! {
     use std::pin::Pin;
+    use std::time::Duration;
 
     use crate::future::Future;
     use crate::stream::FromStream;
 
     pub use merge::Merge;
+    pub use timeout::{TimeoutError, Timeout};
 
     mod merge;
+    mod timeout;
 }
 
 extension_trait! {
@@ -1084,6 +1087,40 @@ extension_trait! {
             Skip::new(self, n)
         }
 
+        #[doc=r#"
+            Await a stream or times out after a duration of time.
+
+            If you want to await an I/O future consider using
+            [`io::timeout`](../io/fn.timeout.html) instead.
+
+            # Examples
+
+            ```
+            # fn main() -> std::io::Result<()> { async_std::task::block_on(async {
+            #
+            use std::time::Duration;
+
+            use async_std::stream;
+            use async_std::prelude::*;
+
+            let mut s = stream::repeat(1).take(3).timeout(Duration::from_secs(1));
+
+            while let Some(v) = s.next().await {
+                assert_eq!(v, Ok(1));
+            }
+            #
+            # Ok(()) }) }
+            ```
+        "#]
+        #[cfg(any(feature = "unstable", feature = "docs"))]
+        #[cfg_attr(feature = "docs", doc(cfg(unstable)))]
+        fn timeout(self, dur: Duration) -> Timeout<Self>
+        where
+            Self: Stream + Sized,
+        {
+            Timeout::new(self, dur)
+        }
+
         #[doc = r#"
             A combinator that applies a function as long as it returns successfully, producing a single, final value.
             Immediately returns the error when the function returns unsuccessfully.

src/stream/stream/timeout.rs

@@ -0,0 +1,63 @@
+use std::error::Error;
+use std::fmt;
+use std::pin::Pin;
+use std::time::Duration;
+
+use futures_timer::Delay;
+use pin_project_lite::pin_project;
+
+use crate::future::Future;
+use crate::stream::Stream;
+use crate::task::{Context, Poll};
+
+pin_project! {
+    /// A stream with timeout time set
+    #[derive(Debug)]
+    pub struct Timeout<S: Stream> {
+        #[pin]
+        stream: S,
+        #[pin]
+        delay: Delay,
+    }
+}
+
+impl<S: Stream> Timeout<S> {
+    pub fn new(stream: S, dur: Duration) -> Timeout<S> {
+        let delay = Delay::new(dur);
+
+        Timeout { stream, delay }
+    }
+}
+
+impl<S: Stream> Stream for Timeout<S> {
+    type Item = Result<S::Item, TimeoutError>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+
+        match this.stream.poll_next(cx) {
+            Poll::Ready(Some(v)) => Poll::Ready(Some(Ok(v))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => match this.delay.poll(cx) {
+                Poll::Ready(_) => Poll::Ready(Some(Err(TimeoutError { _private: () }))),
+                Poll::Pending => Poll::Pending,
+            },
+        }
+    }
+}
+
+/// An error returned when a stream times out.
+#[cfg_attr(feature = "docs", doc(cfg(unstable)))]
+#[cfg(any(feature = "unstable", feature = "docs"))]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct TimeoutError {
+    _private: (),
+}
+
+impl Error for TimeoutError {}
+
+impl fmt::Display for TimeoutError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        "stream has timed out".fmt(f)
+    }
+}