1use crate::error_buffer::{ERROR_BUFFER_CAPACITY, ErrorBuffer, ErrorReporter, error_buffer};
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::time::Duration;
10use tokio::sync::watch;
11use tokio::task::JoinHandle;
12
13#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
15#[non_exhaustive]
16pub enum SchedulerError {
17 #[error("invalid cron expression '{0}': {1}")]
19 InvalidCron(String, String),
20 #[error("the scheduler requires an active Tokio runtime")]
22 RuntimeUnavailable,
23 #[error("scheduled task '{label}' exceeded its {timeout_ms}ms timeout")]
25 TaskTimedOut {
26 label: String,
28 timeout_ms: u64,
30 },
31 #[error("scheduled task '{label}' panicked")]
33 TaskPanicked {
34 label: String,
36 },
37 #[error("scheduled task '{label}' was cancelled")]
39 TaskCancelled {
40 label: String,
42 },
43 #[error("scheduled task '{label}' has no future execution")]
45 ScheduleExhausted {
46 label: String,
48 },
49 #[error("scheduler loop for '{label}' failed: {message}")]
51 LoopFailed {
52 label: String,
54 message: String,
56 },
57}
58
59#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
61#[non_exhaustive]
62pub enum SchedulerFailurePolicy {
63 #[default]
65 Continue,
66 StopTask,
68}
69
70pub type ScheduledHandler =
72 Arc<Box<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>;
73
74pub struct ScheduledTask {
76 label: String,
77 schedule: cron::Schedule,
78 handler: ScheduledHandler,
79}
80
81pub struct Scheduler {
88 tasks: Vec<ScheduledTask>,
89 task_timeout: Duration,
90 failure_policy: SchedulerFailurePolicy,
91}
92
93impl Scheduler {
94 pub fn new() -> Self {
96 Self {
97 tasks: Vec::new(),
98 task_timeout: Duration::from_secs(300),
99 failure_policy: SchedulerFailurePolicy::Continue,
100 }
101 }
102
103 pub fn with_task_timeout(mut self, timeout: Duration) -> Self {
105 self.task_timeout = timeout;
106 self
107 }
108
109 pub fn with_failure_policy(mut self, policy: SchedulerFailurePolicy) -> Self {
111 self.failure_policy = policy;
112 self
113 }
114
115 pub fn task<F, Fut>(mut self, cron_expr: &str, handler: F) -> Result<Self, SchedulerError>
136 where
137 F: Fn() -> Fut + Send + Sync + 'static,
138 Fut: Future<Output = ()> + Send + 'static,
139 {
140 let full_expr = format!("0 {cron_expr} *");
141 let schedule: cron::Schedule = full_expr.parse().map_err(|error: cron::error::Error| {
142 SchedulerError::InvalidCron(cron_expr.to_string(), error.to_string())
143 })?;
144 let boxed: ScheduledHandler = Arc::new(Box::new(move || Box::pin(handler())));
145
146 self.tasks.push(ScheduledTask {
147 label: cron_expr.to_string(),
148 schedule,
149 handler: boxed,
150 });
151 Ok(self)
152 }
153
154 #[cfg_attr(mutants, mutants::skip)]
163 pub fn start(self) -> Result<SchedulerHandle, SchedulerError> {
164 let runtime = tokio::runtime::Handle::try_current()
165 .map_err(|_| SchedulerError::RuntimeUnavailable)?;
166 let (shutdown, _) = watch::channel(false);
167 let (errors_tx, errors) = error_buffer(ERROR_BUFFER_CAPACITY, "scheduler");
168 let mut loops = Vec::with_capacity(self.tasks.len());
169
170 for task in self.tasks {
171 let label = task.label.clone();
172 let shutdown_rx = shutdown.subscribe();
173 let errors_tx = errors_tx.clone();
174 let timeout = self.task_timeout;
175 let policy = self.failure_policy;
176 let task_loop =
177 runtime.spawn(run_task_loop(task, timeout, policy, shutdown_rx, errors_tx));
178 loops.push((label, task_loop));
179 }
180 drop(errors_tx);
181
182 Ok(SchedulerHandle {
183 shutdown,
184 loops,
185 errors,
186 })
187 }
188}
189
190impl Default for Scheduler {
191 fn default() -> Self {
192 Self::new()
193 }
194}
195
196#[must_use = "dropping the scheduler handle immediately stops all scheduled tasks"]
202pub struct SchedulerHandle {
203 shutdown: watch::Sender<bool>,
204 loops: Vec<(String, tokio::task::JoinHandle<()>)>,
205 errors: ErrorBuffer<SchedulerError>,
206}
207
208impl SchedulerHandle {
209 pub async fn next_error(&mut self) -> Option<SchedulerError> {
211 self.errors.next().await
212 }
213
214 pub fn try_next_error(&mut self) -> Option<SchedulerError> {
216 self.errors.try_next()
217 }
218
219 pub fn dropped_errors(&self) -> u64 {
221 self.errors.dropped()
222 }
223
224 pub async fn shutdown(mut self) -> Result<(), SchedulerError> {
229 let _ = self.shutdown.send(true);
230 let mut first_error = self.try_next_error();
231
232 for (label, task_loop) in self.loops.drain(..) {
233 if let Err(error) = task_loop.await
234 && !error.is_cancelled()
235 && first_error.is_none()
236 {
237 first_error = Some(SchedulerError::LoopFailed {
238 label,
239 message: error.to_string(),
240 });
241 }
242 }
243
244 first_error = first_error.or_else(|| self.try_next_error());
245 match first_error {
246 Some(error) => Err(error),
247 None => Ok(()),
248 }
249 }
250
251 fn abort(&mut self) {
252 let _ = self.shutdown.send(true);
253 for (_, task_loop) in &self.loops {
254 task_loop.abort();
255 }
256 self.loops.clear();
257 }
258}
259
260impl Drop for SchedulerHandle {
261 fn drop(&mut self) {
262 self.abort();
263 }
264}
265
266#[cfg_attr(mutants, mutants::skip)]
267async fn run_task_loop(
268 task: ScheduledTask,
269 timeout: Duration,
270 policy: SchedulerFailurePolicy,
271 mut shutdown: watch::Receiver<bool>,
272 errors: ErrorReporter<SchedulerError>,
273) {
274 loop {
275 if shutdown_requested(&shutdown) {
276 break;
277 }
278
279 let now = chrono::Utc::now();
280 let Some(next) = task.schedule.upcoming(chrono::Utc).next() else {
281 errors.report(SchedulerError::ScheduleExhausted {
282 label: task.label.clone(),
283 });
284 break;
285 };
286 let wait = (next - now).to_std().unwrap_or(Duration::ZERO);
287
288 tokio::select! {
289 _ = tokio::time::sleep(wait) => {}
290 _ = wait_for_shutdown(&mut shutdown) => break,
291 }
292
293 match execute_handler(&task, timeout, &mut shutdown).await {
294 Ok(ExecutionStatus::Completed) => {}
295 Ok(ExecutionStatus::ShutDown) => break,
296 Err(error) => {
297 errors.report(error);
298 if policy == SchedulerFailurePolicy::StopTask {
299 break;
300 }
301 }
302 }
303 }
304}
305
306enum ExecutionStatus {
307 Completed,
308 ShutDown,
309}
310
311#[cfg_attr(mutants, mutants::skip)]
312async fn execute_handler(
313 task: &ScheduledTask,
314 timeout: Duration,
315 shutdown: &mut watch::Receiver<bool>,
316) -> Result<ExecutionStatus, SchedulerError> {
317 let handler = Arc::clone(&task.handler);
318 let mut execution = AbortOnDrop(tokio::spawn(async move { handler().await }));
319 let deadline = tokio::time::sleep(timeout);
320 tokio::pin!(deadline);
321
322 tokio::select! {
323 result = &mut execution.0 => match result {
324 Ok(()) => Ok(ExecutionStatus::Completed),
325 Err(error) if error.is_panic() => Err(SchedulerError::TaskPanicked {
326 label: task.label.clone(),
327 }),
328 Err(_) => Err(SchedulerError::TaskCancelled {
329 label: task.label.clone(),
330 }),
331 },
332 _ = &mut deadline => {
333 execution.0.abort();
334 let _ = (&mut execution.0).await;
335 Err(SchedulerError::TaskTimedOut {
336 label: task.label.clone(),
337 timeout_ms: duration_millis_u64(timeout),
338 })
339 }
340 _ = wait_for_shutdown(shutdown) => {
341 execution.0.abort();
342 let _ = (&mut execution.0).await;
343 Ok(ExecutionStatus::ShutDown)
344 }
345 }
346}
347
348fn shutdown_requested(shutdown: &watch::Receiver<bool>) -> bool {
349 *shutdown.borrow()
350}
351
352async fn wait_for_shutdown(shutdown: &mut watch::Receiver<bool>) {
353 while !shutdown_requested(shutdown) {
354 if shutdown.changed().await.is_err() {
355 break;
356 }
357 }
358}
359
360fn duration_millis_u64(duration: Duration) -> u64 {
361 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
362}
363
364struct AbortOnDrop<T>(JoinHandle<T>);
365
366impl<T> Drop for AbortOnDrop<T> {
367 fn drop(&mut self) {
368 self.0.abort();
369 }
370}
371
372#[cfg(test)]
373#[allow(clippy::unwrap_used, clippy::expect_used)]
374mod tests {
375 use super::*;
376 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
377
378 fn task_for_test<F, Fut>(handler: F) -> ScheduledTask
379 where
380 F: Fn() -> Fut + Send + Sync + 'static,
381 Fut: Future<Output = ()> + Send + 'static,
382 {
383 Scheduler::new()
384 .task("* * * * *", handler)
385 .unwrap()
386 .tasks
387 .remove(0)
388 }
389
390 #[test]
391 fn scheduler_start_without_runtime_is_fallible() {
392 let result = Scheduler::new().start();
393 assert!(matches!(result, Err(SchedulerError::RuntimeUnavailable)));
394 }
395
396 #[test]
397 fn invalid_cron_is_rejected() {
398 let result = Scheduler::new().task("invalid cron", || async {});
399 assert!(matches!(result, Err(SchedulerError::InvalidCron(_, _))));
400 }
401
402 #[tokio::test]
403 async fn handler_timeout_is_typed_and_contained() {
404 let task = task_for_test(|| async {
405 tokio::time::sleep(Duration::from_secs(1)).await;
406 });
407 let (_shutdown_tx, mut shutdown) = watch::channel(false);
408 let result = execute_handler(&task, Duration::from_millis(5), &mut shutdown).await;
409
410 assert!(matches!(
411 result,
412 Err(SchedulerError::TaskTimedOut { timeout_ms: 5, .. })
413 ));
414 }
415
416 #[tokio::test]
417 async fn handler_panic_is_typed_and_does_not_escape() {
418 let task = task_for_test(|| async { panic!("test panic") });
419 let (_shutdown_tx, mut shutdown) = watch::channel(false);
420 let result = execute_handler(&task, Duration::from_secs(1), &mut shutdown).await;
421
422 assert!(matches!(result, Err(SchedulerError::TaskPanicked { .. })));
423 }
424
425 #[tokio::test]
426 async fn explicit_shutdown_terminates_sleeping_loops() {
427 let scheduler = Scheduler::new().task("* * * * *", || async {}).unwrap();
428 let handle = scheduler.start().unwrap();
429 let result = tokio::time::timeout(Duration::from_millis(100), handle.shutdown()).await;
430
431 assert!(result.is_ok());
432 assert!(result.unwrap().is_ok());
433 }
434
435 #[tokio::test]
436 async fn shutdown_aborts_current_handler() {
437 let task = task_for_test(|| async {
438 std::future::pending::<()>().await;
439 });
440 let (shutdown_tx, mut shutdown) = watch::channel(false);
441 let execution = tokio::spawn(async move {
442 execute_handler(&task, Duration::from_secs(60), &mut shutdown).await
443 });
444 tokio::task::yield_now().await;
445 shutdown_tx.send(true).unwrap();
446
447 assert!(matches!(
448 execution.await.unwrap().unwrap(),
449 ExecutionStatus::ShutDown
450 ));
451 }
452
453 #[tokio::test]
454 async fn task_loop_never_overlaps_slow_executions() {
455 let active = Arc::new(AtomicUsize::new(0));
456 let maximum = Arc::new(AtomicUsize::new(0));
457 let executions = Arc::new(AtomicUsize::new(0));
458 let active_for_handler = Arc::clone(&active);
459 let maximum_for_handler = Arc::clone(&maximum);
460 let executions_for_handler = Arc::clone(&executions);
461 let handler: ScheduledHandler = Arc::new(Box::new(move || {
462 let active = Arc::clone(&active_for_handler);
463 let maximum = Arc::clone(&maximum_for_handler);
464 let executions = Arc::clone(&executions_for_handler);
465 Box::pin(async move {
466 let running = active.fetch_add(1, Ordering::SeqCst) + 1;
467 maximum.fetch_max(running, Ordering::SeqCst);
468 executions.fetch_add(1, Ordering::SeqCst);
469 tokio::time::sleep(Duration::from_millis(1_100)).await;
470 active.fetch_sub(1, Ordering::SeqCst);
471 })
472 }));
473 let task = ScheduledTask {
474 label: "every second".to_string(),
475 schedule: "* * * * * * *".parse().unwrap(),
476 handler,
477 };
478 let (shutdown_tx, shutdown) = watch::channel(false);
479 let (errors_tx, mut errors) = error_buffer(ERROR_BUFFER_CAPACITY, "test");
480 let task_loop = tokio::spawn(run_task_loop(
481 task,
482 Duration::from_secs(5),
483 SchedulerFailurePolicy::Continue,
484 shutdown,
485 errors_tx,
486 ));
487
488 tokio::time::sleep(Duration::from_millis(2_300)).await;
489 shutdown_tx.send(true).unwrap();
490 task_loop.await.unwrap();
491
492 assert!(executions.load(Ordering::SeqCst) >= 1);
493 assert_eq!(maximum.load(Ordering::SeqCst), 1);
494 assert!(errors.try_next().is_none());
495 }
496
497 #[tokio::test]
498 async fn aborting_execution_drops_the_handler_future() {
499 struct Dropped(Arc<AtomicBool>);
500 impl Drop for Dropped {
501 fn drop(&mut self) {
502 self.0.store(true, Ordering::SeqCst);
503 }
504 }
505
506 let started = Arc::new(AtomicBool::new(false));
507 let dropped = Arc::new(AtomicBool::new(false));
508 let started_for_handler = Arc::clone(&started);
509 let dropped_for_handler = Arc::clone(&dropped);
510 let task = task_for_test(move || {
511 let started = Arc::clone(&started_for_handler);
512 let dropped = Arc::clone(&dropped_for_handler);
513 async move {
514 let _drop_guard = Dropped(dropped);
515 started.store(true, Ordering::SeqCst);
516 std::future::pending::<()>().await;
517 }
518 });
519 let (_shutdown_tx, mut shutdown) = watch::channel(false);
520 let outer = tokio::spawn(async move {
521 execute_handler(&task, Duration::from_secs(60), &mut shutdown).await
522 });
523 while !started.load(Ordering::SeqCst) {
524 tokio::task::yield_now().await;
525 }
526 outer.abort();
527 let _ = outer.await;
528 tokio::time::timeout(Duration::from_millis(100), async {
529 while !dropped.load(Ordering::SeqCst) {
530 tokio::task::yield_now().await;
531 }
532 })
533 .await
534 .unwrap();
535 }
536}
537
538#[cfg(test)]
539#[path = "scheduler_contract_tests.rs"]
540pub(crate) mod contract_tests;