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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use std::fmt::Display;
use std::sync::atomic::{AtomicBool, Ordering};

use std::{
    sync::Arc,
    time::{Duration, Instant},
};

use tracing::{debug, error, info};

use crate::sync::Mutex;
use crate::task::JoinHandle;

#[derive(Clone)]
/// DoomsdayTimer will configurably panic or exit if it is not
/// `reset()` at least every `duration`
pub struct DoomsdayTimer {
    time_to_explode: Arc<Mutex<Instant>>,
    duration: Duration,
    defused: Arc<AtomicBool>,
    aggressive_mode: bool,
}

impl Display for DoomsdayTimer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let fail_mode = if self.aggressive_mode {
            "Exits"
        } else {
            "Panics"
        };
        write!(
            f,
            "DoomsdayTimer(Duration: {:?}, {})",
            self.duration, fail_mode,
        )
    }
}

impl DoomsdayTimer {
    /// Spawn a new doomsday timer.
    /// If `exit_on_explode` is true, it will terminate process with `exit(1)` if it explodes.
    /// Otherwise it will call `panic()`. Note that `awaiting` on the jh will panic if the `DoomsdayTimer` panicked
    pub fn spawn(duration: Duration, exit_on_explode: bool) -> (Self, JoinHandle<()>) {
        let s = Self {
            time_to_explode: Arc::new(Mutex::new(Instant::now() + duration)),
            duration,
            defused: Default::default(),
            aggressive_mode: exit_on_explode,
        };

        let cloned = s.clone();
        let jh = crate::task::spawn(async move {
            cloned.main_loop().await;
        });
        (s, jh)
    }

    /// Reset the timer to it's full duration
    pub async fn reset(&self) {
        let new_time_to_explode = Instant::now() + self.duration;
        *self.time_to_explode.lock().await = new_time_to_explode;
        debug!("{} has been reset", self);
    }

    async fn main_loop(&self) {
        loop {
            if self.defused.load(Ordering::Relaxed) {
                debug!("{} has been defused, terminating main loop", self);
                return;
            }
            let now = Instant::now();
            let time_to_explode = *self.time_to_explode.lock().await;
            if now > time_to_explode {
                error!("{} exploded due to timeout", self);
                self.explode_inner();
            } else {
                let time_to_sleep = time_to_explode - now;
                crate::timer::sleep(time_to_sleep).await;
            }
        }
    }

    /// Force the timer to explode
    pub fn explode(&self) {
        error!("{} was exploded manually", self);
        self.explode_inner();
    }

    fn explode_inner(&self) {
        if self.aggressive_mode {
            error!("{} exiting", self);
            std::process::exit(1);
        } else {
            error!("{} panicking", self);
            panic!("DoomsdayTimer with Duration {:?} exploded", self.duration);
        }
    }
    /// Defuse the timer. Cannot be undone and will no longer `explode`
    pub fn defuse(&self) {
        self.defused.store(true, Ordering::Relaxed);
        info!("{} has been defused", self);
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::DoomsdayTimer;
    use crate::task::run_block_on;
    use crate::test_async;
    use std::io::Error;

    #[test_async(should_panic)]
    async fn test_explode() -> Result<(), Error> {
        let (_, jh) = DoomsdayTimer::spawn(Duration::from_millis(1), false);
        crate::timer::sleep(Duration::from_millis(2)).await;
        jh.await;
        Ok(())
    }

    #[test_async]
    async fn test_do_not_explode() -> Result<(), Error> {
        let (bomb, jh) = DoomsdayTimer::spawn(Duration::from_millis(10), false);
        crate::timer::sleep(Duration::from_millis(5)).await;
        bomb.reset().await;
        crate::timer::sleep(Duration::from_millis(5)).await;
        bomb.reset().await;
        crate::timer::sleep(Duration::from_millis(5)).await;
        bomb.defuse();
        run_block_on(jh);
        Ok(())
    }
}