aws_unlock/
timer.rs

1use std::time::{Duration, Instant};
2
3use anyhow::Result;
4use cancellable_timer::{Canceller, Timer};
5use tokio::task::spawn;
6use tokio::time::sleep_until;
7
8/// A sleep timer that can cancel sleep at any time and observe remaining time periodically.
9pub struct ObservableTimer {
10    timer: Timer,
11}
12
13impl ObservableTimer {
14    pub fn new() -> Result<(Self, Canceller)> {
15        let (timer, canceller) = Timer::new2()?;
16        Ok((Self { timer }, canceller))
17    }
18
19    pub async fn sleep<F>(
20        mut self,
21        total_duration: Duration,
22        inspection_interval: Duration,
23        mut inspect: F,
24    ) -> Result<()>
25    where
26        F: FnMut(Duration),
27    {
28        let start = Instant::now();
29        let entire_sleep = spawn(async move { self.timer.sleep(total_duration) });
30        tokio::pin!(entire_sleep);
31
32        let mut next_inspection = start;
33        loop {
34            let inspection_sleep = sleep_until(next_inspection.into());
35
36            tokio::select! {
37                end = &mut entire_sleep => return Ok(end??),
38                _ = inspection_sleep => {
39                    let elapsed = start.elapsed();
40                    if total_duration > elapsed {
41                        let remaining = total_duration - elapsed;
42                        inspect(remaining)
43                    }
44                },
45            }
46
47            next_inspection += inspection_interval;
48        }
49    }
50}