tokio_prompt_orchestrator/
job_scheduler.rs

1//! # Job Scheduler
2//!
3//! Cron-like job scheduler supporting interval, one-shot, and wall-clock
4//! cron scheduling with async handlers.
5//!
6//! ## Example
7//!
8//! ```rust
9//! use std::time::Duration;
10//! use tokio_prompt_orchestrator::job_scheduler::{Schedule, Scheduler};
11//!
12//! # #[tokio::main]
13//! # async fn main() {
14//! let scheduler = Scheduler::new();
15//! let id = scheduler.schedule("my-job", Schedule::Interval(Duration::from_secs(60)), || {
16//!     Box::pin(async { println!("tick!"); })
17//! });
18//! // scheduler.start(); // spawns background loop
19//! # }
20//! ```
21
22use std::{
23    collections::HashMap,
24    sync::{
25        atomic::{AtomicU64, Ordering},
26        Arc,
27    },
28    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
29};
30
31use futures::future::BoxFuture;
32use tokio::{sync::Mutex, task::JoinHandle, time::sleep};
33
34// ── Schedule ──────────────────────────────────────────────────────────────────
35
36/// Defines when a [`Job`] should fire.
37#[derive(Debug, Clone)]
38pub enum Schedule {
39    /// Fire every fixed duration after the previous run.
40    Interval(Duration),
41    /// Fire exactly once at the specified [`Instant`].
42    Once(Instant),
43    /// Fire at a specific wall-clock time.  `None` fields are wildcards.
44    Cron {
45        /// Hour of day (0–23), or `None` to match any hour.
46        hour: Option<u8>,
47        /// Minute of hour (0–59), or `None` to match any minute.
48        minute: Option<u8>,
49        /// Second of minute (0–59), or `None` to match any second.
50        second: Option<u8>,
51    },
52}
53
54impl Schedule {
55    /// Compute the next fire [`Instant`] relative to `now`.
56    ///
57    /// For [`Schedule::Cron`], this performs wall-clock arithmetic using
58    /// [`SystemTime`] to find the next second that matches all non-`None`
59    /// fields.
60    pub fn next_tick(&self, now: Instant) -> Instant {
61        match self {
62            Schedule::Interval(d) => now + *d,
63            Schedule::Once(at) => *at,
64            Schedule::Cron { hour, minute, second } => {
65                cron_next_tick(*hour, *minute, *second, now)
66            }
67        }
68    }
69}
70
71/// Compute the next wall-clock second that satisfies the cron constraints and
72/// map it back to a monotonic [`Instant`].
73fn cron_next_tick(
74    hour: Option<u8>,
75    minute: Option<u8>,
76    second: Option<u8>,
77    now: Instant,
78) -> Instant {
79    // Current wall-clock time (seconds since epoch).
80    let wall_now = SystemTime::now()
81        .duration_since(UNIX_EPOCH)
82        .unwrap_or_default()
83        .as_secs();
84
85    // Search forward one second at a time (up to 24 h ahead).
86    let max_search = 24 * 3600_usize;
87
88    // Start at least one second in the future.
89    for candidate in (wall_now + 1..).take(max_search) {
90        // Decompose candidate into h/m/s.
91        let s = candidate % 60;
92        let m = (candidate / 60) % 60;
93        let h = (candidate / 3600) % 24;
94
95        let h_ok = hour.is_none_or(|hh| hh as u64 == h);
96        let m_ok = minute.is_none_or(|mm| mm as u64 == m);
97        let s_ok = second.is_none_or(|ss| ss as u64 == s);
98
99        if h_ok && m_ok && s_ok {
100            // Convert wall-clock delta back to a monotonic Instant.
101            let delta_secs = candidate.saturating_sub(wall_now);
102            return now + Duration::from_secs(delta_secs);
103        }
104    }
105
106    // Fallback: 24 h from now (should never happen with valid cron fields).
107    now + Duration::from_secs(86_400)
108}
109
110// ── JobId ─────────────────────────────────────────────────────────────────────
111
112/// Opaque, unique identifier for a registered [`Job`].
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
114pub struct JobId(u64);
115
116impl std::fmt::Display for JobId {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        write!(f, "job:{}", self.0)
119    }
120}
121
122// ── JobHandler ────────────────────────────────────────────────────────────────
123
124/// Type alias for an async job handler that can be cloned and sent across threads.
125pub type JobHandler = Arc<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>;
126
127// ── Job ───────────────────────────────────────────────────────────────────────
128
129/// Metadata and scheduling state for a registered job.
130#[derive(Debug, Clone)]
131pub struct Job {
132    /// Unique identifier assigned by the [`Scheduler`].
133    pub id: JobId,
134    /// Human-readable name.
135    pub name: String,
136    /// Schedule that determines when this job fires.
137    pub schedule: Schedule,
138    /// The last time this job started executing, if it has ever run.
139    pub last_run: Option<Instant>,
140    /// Total number of completed executions.
141    pub run_count: u64,
142    /// Whether this job should be considered by the run loop.
143    pub enabled: bool,
144}
145
146impl Job {
147    fn new(id: JobId, name: String, schedule: Schedule) -> Self {
148        Self {
149            id,
150            name,
151            schedule,
152            last_run: None,
153            run_count: 0,
154            enabled: true,
155        }
156    }
157
158    /// Compute the next scheduled tick based on `last_run` or `now`.
159    fn next_tick(&self, now: Instant) -> Option<Instant> {
160        match &self.schedule {
161            Schedule::Once(_) if self.run_count > 0 => None, // already fired
162            _ => Some(self.schedule.next_tick(self.last_run.unwrap_or(now - Duration::from_secs(1)))),
163        }
164    }
165}
166
167// ── JobStats ──────────────────────────────────────────────────────────────────
168
169/// A snapshot of job scheduling statistics.
170#[derive(Debug, Clone)]
171pub struct JobStats {
172    /// The job's unique identifier.
173    pub id: JobId,
174    /// The job's human-readable name.
175    pub name: String,
176    /// Total number of completed executions.
177    pub run_count: u64,
178    /// When the job last ran (monotonic).
179    pub last_run: Option<Instant>,
180    /// When the job will next run (monotonic), if applicable.
181    pub next_run: Option<Instant>,
182    /// Whether the job is currently enabled.
183    pub enabled: bool,
184}
185
186// ── Scheduler ─────────────────────────────────────────────────────────────────
187
188type JobMap = HashMap<JobId, (Job, JobHandler)>;
189
190/// Async job scheduler with interval, one-shot, and cron-style scheduling.
191///
192/// Create a [`Scheduler`], register jobs with [`Scheduler::schedule`], then
193/// call [`Scheduler::start`] to launch the background run loop.
194#[derive(Clone)]
195pub struct Scheduler {
196    jobs: Arc<Mutex<JobMap>>,
197    next_id: Arc<AtomicU64>,
198}
199
200impl Default for Scheduler {
201    fn default() -> Self {
202        Self::new()
203    }
204}
205
206impl Scheduler {
207    /// Create a new, empty scheduler.
208    pub fn new() -> Self {
209        Self {
210            jobs: Arc::new(Mutex::new(HashMap::new())),
211            next_id: Arc::new(AtomicU64::new(1)),
212        }
213    }
214
215    /// Register a new job and return its [`JobId`].
216    ///
217    /// `handler` is called every time the schedule fires.  It must return a
218    /// `BoxFuture<'static, ()>`.
219    pub async fn schedule<F>(&self, name: &str, schedule: Schedule, handler: F) -> JobId
220    where
221        F: Fn() -> BoxFuture<'static, ()> + Send + Sync + 'static,
222    {
223        let id = JobId(self.next_id.fetch_add(1, Ordering::Relaxed));
224        let job = Job::new(id, name.to_string(), schedule);
225        let handler: JobHandler = Arc::new(handler);
226        self.jobs.lock().await.insert(id, (job, handler));
227        id
228    }
229
230    /// Cancel a job by ID.  Returns `true` if the job existed.
231    pub async fn cancel(&self, id: JobId) -> bool {
232        self.jobs.lock().await.remove(&id).is_some()
233    }
234
235    /// Enable or disable a job.  Disabled jobs are skipped by the run loop.
236    pub async fn enable(&self, id: JobId, enabled: bool) {
237        let mut jobs = self.jobs.lock().await;
238        if let Some((job, _)) = jobs.get_mut(&id) {
239            job.enabled = enabled;
240        }
241    }
242
243    /// Return a snapshot of stats for all registered jobs.
244    pub async fn stats(&self) -> Vec<JobStats> {
245        let now = Instant::now();
246        let jobs = self.jobs.lock().await;
247        jobs.values()
248            .map(|(job, _)| JobStats {
249                id: job.id,
250                name: job.name.clone(),
251                run_count: job.run_count,
252                last_run: job.last_run,
253                next_run: job.next_tick(now),
254                enabled: job.enabled,
255            })
256            .collect()
257    }
258
259    /// Background run loop.  Checks all jobs every 100 ms and fires any whose
260    /// `next_tick` has arrived.  Each handler is spawned in its own Tokio task.
261    pub async fn run_loop(&self) {
262        // Poll at least every 100 ms, sooner when a job is due earlier, so
263        // intervals shorter than the poll period still fire on time.
264        const MAX_POLL: Duration = Duration::from_millis(100);
265        const MIN_POLL: Duration = Duration::from_millis(1);
266        let mut wait = MAX_POLL;
267        loop {
268            sleep(wait).await;
269            let now = Instant::now();
270
271            // Collect jobs that should fire this tick.
272            let mut to_fire: Vec<(JobId, JobHandler)> = Vec::new();
273            {
274                let mut jobs = self.jobs.lock().await;
275                for (job, handler) in jobs.values_mut() {
276                    if !job.enabled {
277                        continue;
278                    }
279                    let next = match job.next_tick(now) {
280                        Some(t) => t,
281                        None => continue, // Once job already ran
282                    };
283                    if next <= now {
284                        to_fire.push((job.id, Arc::clone(handler)));
285                        job.last_run = Some(now);
286                        job.run_count += 1;
287                    }
288                }
289            }
290
291            // Spawn each handler independently.
292            for (_id, handler) in to_fire {
293                tokio::spawn(handler());
294            }
295
296            let earliest = {
297                let jobs = self.jobs.lock().await;
298                jobs.values()
299                    .filter(|(job, _)| job.enabled)
300                    .filter_map(|(job, _)| job.next_tick(now))
301                    .min()
302            };
303            wait = earliest
304                .map(|t| t.saturating_duration_since(Instant::now()))
305                .unwrap_or(MAX_POLL)
306                .clamp(MIN_POLL, MAX_POLL);
307        }
308    }
309
310    /// Spawn the run loop as a background Tokio task and return its handle.
311    pub fn start(&self) -> JoinHandle<()> {
312        let scheduler = self.clone();
313        tokio::spawn(async move { scheduler.run_loop().await })
314    }
315}
316
317// ── Tests ─────────────────────────────────────────────────────────────────────
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    use std::sync::atomic::{AtomicUsize, Ordering as AOrdering};
323    use std::sync::Arc;
324    use tokio::time::sleep;
325
326    #[tokio::test]
327    async fn test_interval_fires_repeatedly() {
328        let counter = Arc::new(AtomicUsize::new(0));
329        let scheduler = Scheduler::new();
330
331        let c = Arc::clone(&counter);
332        scheduler
333            .schedule("counter", Schedule::Interval(Duration::from_millis(10)), move || {
334                let cc = Arc::clone(&c);
335                Box::pin(async move {
336                    cc.fetch_add(1, AOrdering::Relaxed);
337                })
338            })
339            .await;
340
341        let handle = scheduler.start();
342        // Wait long enough for multiple firings (at least ~3).
343        sleep(Duration::from_millis(250)).await;
344        handle.abort();
345
346        let count = counter.load(AOrdering::Relaxed);
347        assert!(count >= 3, "expected >= 3 firings, got {}", count);
348    }
349
350    #[tokio::test]
351    async fn test_cancel_stops_firing() {
352        let counter = Arc::new(AtomicUsize::new(0));
353        let scheduler = Scheduler::new();
354
355        let c = Arc::clone(&counter);
356        let id = scheduler
357            .schedule("cancel-me", Schedule::Interval(Duration::from_millis(10)), move || {
358                let cc = Arc::clone(&c);
359                Box::pin(async move {
360                    cc.fetch_add(1, AOrdering::Relaxed);
361                })
362            })
363            .await;
364
365        let handle = scheduler.start();
366        // Let it fire a couple of times.
367        sleep(Duration::from_millis(80)).await;
368        scheduler.cancel(id).await;
369        let count_after_cancel = counter.load(AOrdering::Relaxed);
370        // Wait more time — count should not grow.
371        sleep(Duration::from_millis(150)).await;
372        handle.abort();
373
374        let final_count = counter.load(AOrdering::Relaxed);
375        assert_eq!(count_after_cancel, final_count, "job fired after cancel");
376    }
377
378    #[tokio::test]
379    async fn test_disabled_job_skips() {
380        let counter = Arc::new(AtomicUsize::new(0));
381        let scheduler = Scheduler::new();
382
383        let c = Arc::clone(&counter);
384        let id = scheduler
385            .schedule("disabled", Schedule::Interval(Duration::from_millis(10)), move || {
386                let cc = Arc::clone(&c);
387                Box::pin(async move {
388                    cc.fetch_add(1, AOrdering::Relaxed);
389                })
390            })
391            .await;
392
393        // Disable before starting.
394        scheduler.enable(id, false).await;
395
396        let handle = scheduler.start();
397        sleep(Duration::from_millis(150)).await;
398        handle.abort();
399
400        assert_eq!(counter.load(AOrdering::Relaxed), 0, "disabled job must not fire");
401    }
402
403    #[tokio::test]
404    async fn test_once_fires_once() {
405        let counter = Arc::new(AtomicUsize::new(0));
406        let scheduler = Scheduler::new();
407
408        let c = Arc::clone(&counter);
409        scheduler
410            .schedule(
411                "once",
412                Schedule::Once(Instant::now() + Duration::from_millis(30)),
413                move || {
414                    let cc = Arc::clone(&c);
415                    Box::pin(async move {
416                        cc.fetch_add(1, AOrdering::Relaxed);
417                    })
418                },
419            )
420            .await;
421
422        let handle = scheduler.start();
423        sleep(Duration::from_millis(250)).await;
424        handle.abort();
425
426        assert_eq!(counter.load(AOrdering::Relaxed), 1, "once job should fire exactly once");
427    }
428
429    #[test]
430    fn test_schedule_next_tick_interval() {
431        let now = Instant::now();
432        let sched = Schedule::Interval(Duration::from_secs(5));
433        let next = sched.next_tick(now);
434        assert!(next > now);
435        assert!(next <= now + Duration::from_secs(6));
436    }
437
438    #[test]
439    fn test_schedule_next_tick_once() {
440        let target = Instant::now() + Duration::from_secs(10);
441        let sched = Schedule::Once(target);
442        assert_eq!(sched.next_tick(Instant::now()), target);
443    }
444
445    #[test]
446    fn test_cron_next_tick_is_future() {
447        let now = Instant::now();
448        let sched = Schedule::Cron { hour: None, minute: None, second: None };
449        let next = sched.next_tick(now);
450        assert!(next > now, "cron next tick should be in the future");
451    }
452}
tokio_prompt_orchestrator/job_scheduler.rs

tokio_prompt_orchestrator/
job_scheduler.rs
