Skip to main content

cliquenet/
time.rs

1use std::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use tokio::time::{Duration, Instant, Sleep, sleep};
8
9/// A countdown timer that can be reset.
10#[derive(Debug)]
11pub struct Countdown {
12    // The actual future to await.
13    sleep: Pin<Box<Sleep>>,
14
15    // Is this countdown running?
16    stopped: bool,
17}
18
19impl Default for Countdown {
20    fn default() -> Self {
21        Self::new()
22    }
23}
24
25impl Countdown {
26    /// Create a new countdown.
27    ///
28    /// When ready, use `Countdown::start` to begin.
29    pub fn new() -> Self {
30        Self {
31            sleep: Box::pin(sleep(Duration::from_secs(1))),
32            stopped: true,
33        }
34    }
35
36    /// Start the countdown.
37    ///
38    /// Once started, a countdown can not be started again, unless
39    /// `Countdown::stop` is invoked first.
40    pub fn start(&mut self, timeout: Duration) {
41        if !self.stopped {
42            // The countdown is already running.
43            return;
44        }
45        self.stopped = false;
46        self.sleep.as_mut().reset(Instant::now() + timeout);
47    }
48
49    /// Stop this countdown.
50    pub fn stop(&mut self) {
51        self.stopped = true
52    }
53}
54
55impl Future for Countdown {
56    type Output = ();
57
58    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
59        if self.stopped {
60            return Poll::Pending;
61        }
62        self.as_mut().sleep.as_mut().poll(cx)
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use tokio::time::{Duration, Instant, timeout};
69
70    use super::Countdown;
71
72    #[tokio::test]
73    async fn countdown() {
74        let mut c = Countdown::new();
75
76        let now = Instant::now();
77        c.start(Duration::from_secs(1));
78        (&mut c).await;
79        assert!(now.elapsed() >= Duration::from_secs(1));
80
81        // Once finished, the countdown stays finished:
82        let now = Instant::now();
83        (&mut c).await;
84        assert!(now.elapsed() < Duration::from_millis(1));
85
86        // If stopped it does not end:
87        c.start(Duration::from_secs(1));
88        c.stop();
89        assert!(timeout(Duration::from_secs(2), &mut c).await.is_err());
90
91        // until started again:
92        c.start(Duration::from_secs(1));
93        let now = Instant::now();
94        (&mut c).await;
95        assert!(now.elapsed() >= Duration::from_secs(1));
96    }
97}