use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, watch};
use tokio::task::JoinHandle;
#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum SchedulerError {
#[error("invalid cron expression '{0}': {1}")]
InvalidCron(String, String),
#[error("the scheduler requires an active Tokio runtime")]
RuntimeUnavailable,
#[error("scheduled task '{label}' exceeded its {timeout_ms}ms timeout")]
TaskTimedOut {
label: String,
timeout_ms: u64,
},
#[error("scheduled task '{label}' panicked")]
TaskPanicked {
label: String,
},
#[error("scheduled task '{label}' was cancelled")]
TaskCancelled {
label: String,
},
#[error("scheduled task '{label}' has no future execution")]
ScheduleExhausted {
label: String,
},
#[error("scheduler loop for '{label}' failed: {message}")]
LoopFailed {
label: String,
message: String,
},
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum SchedulerFailurePolicy {
#[default]
Continue,
StopTask,
}
pub type ScheduledHandler =
Arc<Box<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>>;
pub struct ScheduledTask {
label: String,
schedule: cron::Schedule,
handler: ScheduledHandler,
}
pub struct Scheduler {
tasks: Vec<ScheduledTask>,
task_timeout: Duration,
failure_policy: SchedulerFailurePolicy,
}
impl Scheduler {
pub fn new() -> Self {
Self {
tasks: Vec::new(),
task_timeout: Duration::from_secs(300),
failure_policy: SchedulerFailurePolicy::Continue,
}
}
pub fn with_task_timeout(mut self, timeout: Duration) -> Self {
self.task_timeout = timeout;
self
}
pub fn with_failure_policy(mut self, policy: SchedulerFailurePolicy) -> Self {
self.failure_policy = policy;
self
}
pub fn task<F, Fut>(mut self, cron_expr: &str, handler: F) -> Result<Self, SchedulerError>
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let full_expr = format!("0 {cron_expr} *");
let schedule: cron::Schedule = full_expr.parse().map_err(|error: cron::error::Error| {
SchedulerError::InvalidCron(cron_expr.to_string(), error.to_string())
})?;
let boxed: ScheduledHandler = Arc::new(Box::new(move || Box::pin(handler())));
self.tasks.push(ScheduledTask {
label: cron_expr.to_string(),
schedule,
handler: boxed,
});
Ok(self)
}
#[cfg_attr(mutants, mutants::skip)]
pub fn start(self) -> Result<SchedulerHandle, SchedulerError> {
let runtime = tokio::runtime::Handle::try_current()
.map_err(|_| SchedulerError::RuntimeUnavailable)?;
let (shutdown, _) = watch::channel(false);
let (errors_tx, errors) = mpsc::unbounded_channel();
let mut loops = Vec::with_capacity(self.tasks.len());
for task in self.tasks {
let label = task.label.clone();
let shutdown_rx = shutdown.subscribe();
let errors_tx = errors_tx.clone();
let timeout = self.task_timeout;
let policy = self.failure_policy;
let task_loop =
runtime.spawn(run_task_loop(task, timeout, policy, shutdown_rx, errors_tx));
loops.push((label, task_loop));
}
drop(errors_tx);
Ok(SchedulerHandle {
shutdown,
loops,
errors,
})
}
}
impl Default for Scheduler {
fn default() -> Self {
Self::new()
}
}
#[must_use = "dropping the scheduler handle immediately stops all scheduled tasks"]
pub struct SchedulerHandle {
shutdown: watch::Sender<bool>,
loops: Vec<(String, tokio::task::JoinHandle<()>)>,
errors: mpsc::UnboundedReceiver<SchedulerError>,
}
impl SchedulerHandle {
pub async fn next_error(&mut self) -> Option<SchedulerError> {
self.errors.recv().await
}
pub fn try_next_error(&mut self) -> Option<SchedulerError> {
self.errors.try_recv().ok()
}
pub async fn shutdown(mut self) -> Result<(), SchedulerError> {
let _ = self.shutdown.send(true);
let mut first_error = self.try_next_error();
for (label, task_loop) in self.loops.drain(..) {
if let Err(error) = task_loop.await
&& !error.is_cancelled()
&& first_error.is_none()
{
first_error = Some(SchedulerError::LoopFailed {
label,
message: error.to_string(),
});
}
}
first_error = first_error.or_else(|| self.try_next_error());
match first_error {
Some(error) => Err(error),
None => Ok(()),
}
}
fn abort(&mut self) {
let _ = self.shutdown.send(true);
for (_, task_loop) in &self.loops {
task_loop.abort();
}
self.loops.clear();
}
}
impl Drop for SchedulerHandle {
fn drop(&mut self) {
self.abort();
}
}
#[cfg_attr(mutants, mutants::skip)]
async fn run_task_loop(
task: ScheduledTask,
timeout: Duration,
policy: SchedulerFailurePolicy,
mut shutdown: watch::Receiver<bool>,
errors: mpsc::UnboundedSender<SchedulerError>,
) {
loop {
if shutdown_requested(&shutdown) {
break;
}
let now = chrono::Utc::now();
let Some(next) = task.schedule.upcoming(chrono::Utc).next() else {
let _ = errors.send(SchedulerError::ScheduleExhausted {
label: task.label.clone(),
});
break;
};
let wait = (next - now).to_std().unwrap_or(Duration::ZERO);
tokio::select! {
_ = tokio::time::sleep(wait) => {}
_ = wait_for_shutdown(&mut shutdown) => break,
}
match execute_handler(&task, timeout, &mut shutdown).await {
Ok(ExecutionStatus::Completed) => {}
Ok(ExecutionStatus::ShutDown) => break,
Err(error) => {
let _ = errors.send(error);
if policy == SchedulerFailurePolicy::StopTask {
break;
}
}
}
}
}
enum ExecutionStatus {
Completed,
ShutDown,
}
#[cfg_attr(mutants, mutants::skip)]
async fn execute_handler(
task: &ScheduledTask,
timeout: Duration,
shutdown: &mut watch::Receiver<bool>,
) -> Result<ExecutionStatus, SchedulerError> {
let handler = Arc::clone(&task.handler);
let mut execution = AbortOnDrop(tokio::spawn(async move { handler().await }));
let deadline = tokio::time::sleep(timeout);
tokio::pin!(deadline);
tokio::select! {
result = &mut execution.0 => match result {
Ok(()) => Ok(ExecutionStatus::Completed),
Err(error) if error.is_panic() => Err(SchedulerError::TaskPanicked {
label: task.label.clone(),
}),
Err(_) => Err(SchedulerError::TaskCancelled {
label: task.label.clone(),
}),
},
_ = &mut deadline => {
execution.0.abort();
let _ = (&mut execution.0).await;
Err(SchedulerError::TaskTimedOut {
label: task.label.clone(),
timeout_ms: duration_millis_u64(timeout),
})
}
_ = wait_for_shutdown(shutdown) => {
execution.0.abort();
let _ = (&mut execution.0).await;
Ok(ExecutionStatus::ShutDown)
}
}
}
fn shutdown_requested(shutdown: &watch::Receiver<bool>) -> bool {
*shutdown.borrow()
}
async fn wait_for_shutdown(shutdown: &mut watch::Receiver<bool>) {
while !shutdown_requested(shutdown) {
if shutdown.changed().await.is_err() {
break;
}
}
}
fn duration_millis_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
struct AbortOnDrop<T>(JoinHandle<T>);
impl<T> Drop for AbortOnDrop<T> {
fn drop(&mut self) {
self.0.abort();
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
fn task_for_test<F, Fut>(handler: F) -> ScheduledTask
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
Scheduler::new()
.task("* * * * *", handler)
.unwrap()
.tasks
.remove(0)
}
#[test]
fn scheduler_start_without_runtime_is_fallible() {
let result = Scheduler::new().start();
assert!(matches!(result, Err(SchedulerError::RuntimeUnavailable)));
}
#[test]
fn invalid_cron_is_rejected() {
let result = Scheduler::new().task("invalid cron", || async {});
assert!(matches!(result, Err(SchedulerError::InvalidCron(_, _))));
}
#[tokio::test]
async fn handler_timeout_is_typed_and_contained() {
let task = task_for_test(|| async {
tokio::time::sleep(Duration::from_secs(1)).await;
});
let (_shutdown_tx, mut shutdown) = watch::channel(false);
let result = execute_handler(&task, Duration::from_millis(5), &mut shutdown).await;
assert!(matches!(
result,
Err(SchedulerError::TaskTimedOut { timeout_ms: 5, .. })
));
}
#[tokio::test]
async fn handler_panic_is_typed_and_does_not_escape() {
let task = task_for_test(|| async { panic!("test panic") });
let (_shutdown_tx, mut shutdown) = watch::channel(false);
let result = execute_handler(&task, Duration::from_secs(1), &mut shutdown).await;
assert!(matches!(result, Err(SchedulerError::TaskPanicked { .. })));
}
#[tokio::test]
async fn explicit_shutdown_terminates_sleeping_loops() {
let scheduler = Scheduler::new().task("* * * * *", || async {}).unwrap();
let handle = scheduler.start().unwrap();
let result = tokio::time::timeout(Duration::from_millis(100), handle.shutdown()).await;
assert!(result.is_ok());
assert!(result.unwrap().is_ok());
}
#[tokio::test]
async fn shutdown_aborts_current_handler() {
let task = task_for_test(|| async {
std::future::pending::<()>().await;
});
let (shutdown_tx, mut shutdown) = watch::channel(false);
let execution = tokio::spawn(async move {
execute_handler(&task, Duration::from_secs(60), &mut shutdown).await
});
tokio::task::yield_now().await;
shutdown_tx.send(true).unwrap();
assert!(matches!(
execution.await.unwrap().unwrap(),
ExecutionStatus::ShutDown
));
}
#[tokio::test]
async fn task_loop_never_overlaps_slow_executions() {
let active = Arc::new(AtomicUsize::new(0));
let maximum = Arc::new(AtomicUsize::new(0));
let executions = Arc::new(AtomicUsize::new(0));
let active_for_handler = Arc::clone(&active);
let maximum_for_handler = Arc::clone(&maximum);
let executions_for_handler = Arc::clone(&executions);
let handler: ScheduledHandler = Arc::new(Box::new(move || {
let active = Arc::clone(&active_for_handler);
let maximum = Arc::clone(&maximum_for_handler);
let executions = Arc::clone(&executions_for_handler);
Box::pin(async move {
let running = active.fetch_add(1, Ordering::SeqCst) + 1;
maximum.fetch_max(running, Ordering::SeqCst);
executions.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(1_100)).await;
active.fetch_sub(1, Ordering::SeqCst);
})
}));
let task = ScheduledTask {
label: "every second".to_string(),
schedule: "* * * * * * *".parse().unwrap(),
handler,
};
let (shutdown_tx, shutdown) = watch::channel(false);
let (errors_tx, mut errors) = mpsc::unbounded_channel();
let task_loop = tokio::spawn(run_task_loop(
task,
Duration::from_secs(5),
SchedulerFailurePolicy::Continue,
shutdown,
errors_tx,
));
tokio::time::sleep(Duration::from_millis(2_300)).await;
shutdown_tx.send(true).unwrap();
task_loop.await.unwrap();
assert!(executions.load(Ordering::SeqCst) >= 1);
assert_eq!(maximum.load(Ordering::SeqCst), 1);
assert!(errors.try_recv().is_err());
}
#[tokio::test]
async fn aborting_execution_drops_the_handler_future() {
struct Dropped(Arc<AtomicBool>);
impl Drop for Dropped {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
let started = Arc::new(AtomicBool::new(false));
let dropped = Arc::new(AtomicBool::new(false));
let started_for_handler = Arc::clone(&started);
let dropped_for_handler = Arc::clone(&dropped);
let task = task_for_test(move || {
let started = Arc::clone(&started_for_handler);
let dropped = Arc::clone(&dropped_for_handler);
async move {
let _drop_guard = Dropped(dropped);
started.store(true, Ordering::SeqCst);
std::future::pending::<()>().await;
}
});
let (_shutdown_tx, mut shutdown) = watch::channel(false);
let outer = tokio::spawn(async move {
execute_handler(&task, Duration::from_secs(60), &mut shutdown).await
});
while !started.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
outer.abort();
let _ = outer.await;
tokio::time::timeout(Duration::from_millis(100), async {
while !dropped.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
}
#[cfg(test)]
#[path = "scheduler_contract_tests.rs"]
mod contract_tests;