1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

use core::future::Future;
use core::pin::Pin;
use core::task::{ready, Context, Poll};

/// Suspends a future until the specified deadline.
///
/// This `struct` is created by the [`wait_until`] method on [`FutureExt`]. See its
/// documentation for more.
///
/// [`wait_until`]: crate::future::FutureExt::wait_until
/// [`FutureExt`]: crate::future::FutureExt
#[derive(Debug)]
#[pin_project::pin_project]
#[must_use = "futures do nothing unless polled or .awaited"]
pub struct WaitUntil<F, D> {
    #[pin]
    future: F,
    #[pin]
    deadline: D,
    state: State,
}

/// The internal state
#[derive(Debug)]
enum State {
    Started,
    PollFuture,
    Completed,
}

impl<F, D> WaitUntil<F, D> {
    pub(super) fn new(future: F, deadline: D) -> Self {
        Self {
            future,
            deadline,
            state: State::Started,
        }
    }
}

impl<F: Future, D: Future> Future for WaitUntil<F, D> {
    type Output = F::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        loop {
            match this.state {
                State::Started => {
                    ready!(this.deadline.as_mut().poll(cx));
                    *this.state = State::PollFuture;
                }
                State::PollFuture => {
                    let value = ready!(this.future.as_mut().poll(cx));
                    *this.state = State::Completed;
                    return Poll::Ready(value);
                }
                State::Completed => panic!("future polled after completing"),
            }
        }
    }
}