diff --git a/src/io/mod.rs b/src/io/mod.rs index daf70b8..aae9cb5 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -32,6 +32,7 @@ pub use seek::Seek; pub use stderr::{stderr, Stderr}; pub use stdin::{stdin, Stdin}; pub use stdout::{stdout, Stdout}; +pub use timeout::timeout; pub use write::Write; mod buf_read; @@ -42,4 +43,5 @@ mod seek; mod stderr; mod stdin; mod stdout; +mod timeout; mod write; diff --git a/src/io/timeout.rs b/src/io/timeout.rs new file mode 100644 index 0000000..5902989 --- /dev/null +++ b/src/io/timeout.rs @@ -0,0 +1,70 @@ +use std::pin::Pin; +use std::time::Duration; + +use futures_timer::Delay; +use pin_utils::unsafe_pinned; + +use crate::future::Future; +use crate::io; +use crate::task::{Context, Poll}; + +/// Awaits an I/O future or times out after a duration of time. +/// +/// # Examples +/// +/// ```no_run +/// # #![feature(async_await)] +/// # fn main() -> std::io::Result<()> { async_std::task::block_on(async { +/// # +/// use std::time::Duration; +/// +/// use async_std::io; +/// +/// let stdin = io::stdin(); +/// let mut line = String::new(); +/// +/// let dur = Duration::from_secs(5); +/// let n = io::timeout(dur, stdin.read_line(&mut line)).await?; +/// # +/// # Ok(()) }) } +/// ``` +pub async fn timeout(dur: Duration, f: F) -> io::Result +where + F: Future>, +{ + let f = TimeoutFuture { + future: f, + delay: Delay::new(dur), + }; + f.await +} + +struct TimeoutFuture { + future: F, + delay: Delay, +} + +impl TimeoutFuture { + unsafe_pinned!(future: F); + unsafe_pinned!(delay: Delay); +} + +impl Future for TimeoutFuture +where + F: Future>, +{ + type Output = F::Output; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match self.as_mut().future().poll(cx) { + Poll::Ready(v) => Poll::Ready(v), + Poll::Pending => match self.delay().poll(cx) { + Poll::Ready(_) => Poll::Ready(Err(io::Error::new( + io::ErrorKind::TimedOut, + "future has timed out", + ))), + Poll::Pending => Poll::Pending, + }, + } + } +}