1use std::{
2 future::Future,
3 pin::Pin,
4 task::{Context, Poll},
5};
6
7use tokio::time::{Duration, Instant, Sleep, sleep};
8
9#[derive(Debug)]
11pub struct Countdown {
12 sleep: Pin<Box<Sleep>>,
14
15 stopped: bool,
17}
18
19impl Default for Countdown {
20 fn default() -> Self {
21 Self::new()
22 }
23}
24
25impl Countdown {
26 pub fn new() -> Self {
30 Self {
31 sleep: Box::pin(sleep(Duration::from_secs(1))),
32 stopped: true,
33 }
34 }
35
36 pub fn start(&mut self, timeout: Duration) {
41 if !self.stopped {
42 return;
44 }
45 self.stopped = false;
46 self.sleep.as_mut().reset(Instant::now() + timeout);
47 }
48
49 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 let now = Instant::now();
83 (&mut c).await;
84 assert!(now.elapsed() < Duration::from_millis(1));
85
86 c.start(Duration::from_secs(1));
88 c.stop();
89 assert!(timeout(Duration::from_secs(2), &mut c).await.is_err());
90
91 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}