diff --git a/examples/stdin-timeout.rs b/examples/stdin-timeout.rs index dde49e5b..04b860a2 100644 --- a/examples/stdin-timeout.rs +++ b/examples/stdin-timeout.rs @@ -9,22 +9,13 @@ use async_std::prelude::*; use async_std::task; fn main() -> io::Result<()> { - task::block_on(async { + task::block_on(io::timeout(Duration::from_secs(5), async { let stdin = io::stdin(); - let mut line = String::new(); - match stdin - .read_line(&mut line) - .timeout(Duration::from_secs(5)) - .await - { - Ok(res) => { - res?; - print!("Got line: {}", line); - } - Err(_) => println!("You have only 5 seconds to enter a line. Try again :)"), - } + let mut line = String::new(); + stdin.read_line(&mut line).await?; + print!("Got line: {}", line); Ok(()) - }) + })) } diff --git a/src/io/mod.rs b/src/io/mod.rs index 0355031e..13b91d6c 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -34,6 +34,7 @@ pub use sink::{sink, Sink}; 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; @@ -46,4 +47,5 @@ mod sink; 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 00000000..b35a8030 --- /dev/null +++ b/src/io/timeout.rs @@ -0,0 +1,71 @@ +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; +/// +/// io::timeout(Duration::from_secs(5), async { +/// let stdin = io::stdin(); +/// let mut line = String::new(); +/// let n = stdin.read_line(&mut line).await?; +/// }) +/// .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, + "I/O operation has timed out", + ))), + Poll::Pending => Poll::Pending, + }, + } + } +}