bramble_common/
timer.rs

1//! Timer utils
2
3use crate::{Error, Result};
4use async_io::Timer;
5use futures::{ready, Future};
6use pin_project::pin_project;
7use std::time::Duration;
8use std::{
9    pin::Pin,
10    task::{Context, Poll},
11};
12
13/// Sleeps for the specified amount of time.
14pub async fn sleep(duration: Duration) {
15    let _ = Timer::after(duration).await;
16}
17
18/// Awaits a future or times out after a duration of time.
19pub async fn timeout<F, T>(duration: Duration, f: F) -> Result<T>
20where
21    F: Future<Output = Result<T>>,
22{
23    TimeoutFuture::new(duration, f).await
24}
25
26#[pin_project]
27struct TimeoutFuture<F> {
28    #[pin]
29    future: F,
30    #[pin]
31    delay: Timer,
32}
33
34impl<F> TimeoutFuture<F> {
35    fn new(duration: Duration, future: F) -> TimeoutFuture<F> {
36        TimeoutFuture {
37            future,
38            delay: Timer::after(duration),
39        }
40    }
41}
42
43impl<F, T> Future for TimeoutFuture<F>
44where
45    F: Future<Output = Result<T>>,
46{
47    type Output = Result<T>;
48
49    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
50        let this = self.project();
51        match this.future.poll(cx) {
52            Poll::Pending => {
53                let _ = ready!(this.delay.poll(cx));
54                Poll::Ready(Err(Error::Timeout))
55            }
56            r => r,
57        }
58    }
59}