rovenue/polling/
scheduler.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::{Arc, Mutex};
3use std::thread;
4use std::time::{Duration, Instant};
5
6type Tick = Arc<dyn Fn() + Send + Sync>;
7
8struct Registration {
9    interval: Duration,
10    last_fired: Mutex<Option<Instant>>,
11    tick: Tick,
12}
13
14pub struct PollingScheduler {
15    inner: Arc<SchedulerInner>,
16    thread: Mutex<Option<thread::JoinHandle<()>>>,
17}
18
19struct SchedulerInner {
20    registrations: Mutex<Vec<(String, Arc<Registration>)>>,
21    foreground: AtomicBool,
22    running: AtomicBool,
23}
24
25impl PollingScheduler {
26    pub fn new() -> Self {
27        let inner = Arc::new(SchedulerInner {
28            registrations: Mutex::new(Vec::new()),
29            foreground: AtomicBool::new(false),
30            running: AtomicBool::new(true),
31        });
32        let inner_clone = Arc::clone(&inner);
33        let handle = thread::spawn(move || run_loop(inner_clone));
34        Self {
35            inner,
36            thread: Mutex::new(Some(handle)),
37        }
38    }
39
40    pub fn register(
41        &self,
42        name: &str,
43        interval: Duration,
44        tick: impl Fn() + Send + Sync + 'static,
45    ) {
46        let mut regs = self.inner.registrations.lock().expect("regs poisoned");
47        regs.push((
48            name.to_string(),
49            Arc::new(Registration {
50                interval,
51                last_fired: Mutex::new(None),
52                tick: Arc::new(tick),
53            }),
54        ));
55    }
56
57    pub fn set_foreground(&self, foreground: bool) {
58        self.inner.foreground.store(foreground, Ordering::SeqCst);
59    }
60
61    /// Clear every registration's last-fired time so the loop fires each task on
62    /// its next tick. Used on foreground transitions to refresh immediately
63    /// instead of waiting out the remaining interval.
64    pub fn reset_cadence(&self) {
65        let regs = self.inner.registrations.lock().expect("regs poisoned");
66        for (_name, reg) in regs.iter() {
67            *reg.last_fired.lock().expect("last_fired poisoned") = None;
68        }
69    }
70
71    pub fn shutdown(&self) {
72        if self.inner.running.swap(false, Ordering::SeqCst) {
73            // Wake up a sleeping loop quickly.
74            self.inner.foreground.store(false, Ordering::SeqCst);
75        }
76        let mut g = self.thread.lock().expect("thread mutex poisoned");
77        if let Some(h) = g.take() {
78            let _ = h.join();
79        }
80    }
81}
82
83impl Default for PollingScheduler {
84    fn default() -> Self {
85        Self::new()
86    }
87}
88
89impl Drop for PollingScheduler {
90    fn drop(&mut self) {
91        self.shutdown();
92    }
93}
94
95fn run_loop(inner: Arc<SchedulerInner>) {
96    let tick_resolution = Duration::from_millis(10);
97    while inner.running.load(Ordering::SeqCst) {
98        if !inner.foreground.load(Ordering::SeqCst) {
99            thread::sleep(tick_resolution);
100            continue;
101        }
102        let regs = inner.registrations.lock().expect("regs poisoned").clone();
103        let now = Instant::now();
104        for (_name, reg) in regs {
105            let mut last = reg.last_fired.lock().expect("last_fired poisoned");
106            let due = match *last {
107                None => true,
108                Some(t) => now.duration_since(t) >= reg.interval,
109            };
110            if due {
111                *last = Some(now);
112                drop(last);
113                (reg.tick)();
114            }
115        }
116        thread::sleep(tick_resolution);
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn reset_cadence_refires_immediately_after_reforeground() {
126        use std::sync::atomic::{AtomicUsize, Ordering};
127        use std::sync::Arc;
128
129        let sched = PollingScheduler::new();
130        let count = Arc::new(AtomicUsize::new(0));
131        {
132            let count = Arc::clone(&count);
133            sched.register("t", std::time::Duration::from_secs(3600), move || {
134                count.fetch_add(1, Ordering::SeqCst);
135            });
136        }
137        sched.set_foreground(true);
138        std::thread::sleep(std::time::Duration::from_millis(60));
139        assert_eq!(
140            count.load(Ordering::SeqCst),
141            1,
142            "first foreground fires once"
143        );
144
145        sched.reset_cadence();
146        std::thread::sleep(std::time::Duration::from_millis(60));
147        assert_eq!(
148            count.load(Ordering::SeqCst),
149            2,
150            "reset_cadence re-fires immediately"
151        );
152    }
153}