Skip to main content

sz_orm_scheduler/
scheduler.rs

1//! Job handler abstractions used by [`super::CronScheduler`].
2//!
3//! This module is intentionally small and self-contained: it provides the
4//! `JobHandler` trait and a couple of default implementations so the
5//! scheduler can invoke user-supplied logic when a task fires.
6
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9use std::sync::RwLock;
10
11use super::ScheduledTask;
12
13/// Handler invoked by the scheduler when a task's cron expression matches
14/// the current time. Implementations must be `Send + Sync` because they may
15/// be invoked from a background thread.
16pub trait JobHandler: Send + Sync {
17    fn handle(&self, task: &ScheduledTask) -> Result<(), String>;
18}
19
20/// Default `JobHandler` that increments an atomic counter every time it is
21/// invoked. Useful both as a no-op default and for unit tests that need to
22/// observe how many times a job fired.
23pub struct CounterJobHandler {
24    counter: Arc<AtomicU64>,
25}
26
27impl CounterJobHandler {
28    pub fn new() -> Self {
29        Self {
30            counter: Arc::new(AtomicU64::new(0)),
31        }
32    }
33
34    pub fn counter(&self) -> Arc<AtomicU64> {
35        self.counter.clone()
36    }
37
38    pub fn count(&self) -> u64 {
39        self.counter.load(Ordering::SeqCst)
40    }
41}
42
43impl Default for CounterJobHandler {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl JobHandler for CounterJobHandler {
50    fn handle(&self, _task: &ScheduledTask) -> Result<(), String> {
51        self.counter.fetch_add(1, Ordering::SeqCst);
52        Ok(())
53    }
54}
55
56/// `JobHandler` that records the IDs of every task it has handled, in the
57/// order they were handled. Useful for tests that need to assert which tasks
58/// fired and in what order.
59pub struct RecordingJobHandler {
60    handled: Arc<RwLock<Vec<String>>>,
61}
62
63impl RecordingJobHandler {
64    pub fn new() -> Self {
65        Self {
66            handled: Arc::new(RwLock::new(Vec::new())),
67        }
68    }
69
70    pub fn handled_ids(&self) -> Vec<String> {
71        self.handled.read().map(|h| h.clone()).unwrap_or_default()
72    }
73
74    pub fn clear(&self) {
75        if let Ok(mut h) = self.handled.write() {
76            h.clear();
77        }
78    }
79}
80
81impl Default for RecordingJobHandler {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl JobHandler for RecordingJobHandler {
88    fn handle(&self, task: &ScheduledTask) -> Result<(), String> {
89        if let Ok(mut h) = self.handled.write() {
90            h.push(task.id.clone());
91        }
92        Ok(())
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    fn sample_task(id: &str) -> ScheduledTask {
101        ScheduledTask::new(id, id, "* * * * *")
102    }
103
104    #[test]
105    fn test_counter_handler_increments() {
106        let handler = CounterJobHandler::new();
107        assert_eq!(handler.count(), 0);
108        handler.handle(&sample_task("t1")).unwrap();
109        handler.handle(&sample_task("t2")).unwrap();
110        assert_eq!(handler.count(), 2);
111    }
112
113    #[test]
114    fn test_recording_handler_records_ids() {
115        let handler = RecordingJobHandler::new();
116        handler.handle(&sample_task("a")).unwrap();
117        handler.handle(&sample_task("b")).unwrap();
118        handler.handle(&sample_task("c")).unwrap();
119        assert_eq!(handler.handled_ids(), vec!["a", "b", "c"]);
120        handler.clear();
121        assert!(handler.handled_ids().is_empty());
122    }
123
124    #[test]
125    fn test_handler_is_send_sync() {
126        fn assert_send_sync<T: Send + Sync>() {}
127        assert_send_sync::<CounterJobHandler>();
128        assert_send_sync::<RecordingJobHandler>();
129        assert_send_sync::<Arc<dyn JobHandler>>();
130    }
131}