diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs index 501ece1b..a18b4b67 100644 --- a/src/stream/stream/mod.rs +++ b/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 + 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. diff --git a/src/stream/stream/timeout.rs b/src/stream/stream/timeout.rs new file mode 100644 index 00000000..3c14811f --- /dev/null +++ b/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 { + #[pin] + stream: S, + #[pin] + delay: Delay, + } +} + +impl Timeout { + pub fn new(stream: S, dur: Duration) -> Timeout { + let delay = Delay::new(dur); + + Timeout { stream, delay } + } +} + +impl Stream for Timeout { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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) + } +}