use std::time::Duration;
use std::str::FromStr;
use std::sync::Arc;
use chrono::{DateTime, Utc};
use cron::Schedule;
use serde::{Serialize, Deserialize};
use thiserror::Error;
use crate::{Job, JobEntry, Priority, QueueResult, QueueError};
#[derive(Error, Debug)]
pub enum ScheduleError {
#[error("Invalid cron expression: {0}")]
InvalidCron(String),
#[error("Schedule not found: {0}")]
ScheduleNotFound(String),
#[error("Invalid retry configuration: {0}")]
InvalidRetryConfig(String),
#[error("Queue error: {0}")]
Queue(#[from] QueueError),
}
pub type ScheduleResult<T> = Result<T, ScheduleError>;
#[derive(Debug, Clone, Serialize)]
pub struct CronExpression {
expression: String,
#[serde(skip)]
schedule: Option<Schedule>,
}
impl CronExpression {
pub fn new(expression: &str) -> ScheduleResult<Self> {
let schedule = Schedule::from_str(expression)
.map_err(|e| ScheduleError::InvalidCron(format!("{}: {}", expression, e)))?;
Ok(CronExpression {
expression: expression.to_string(),
schedule: Some(schedule),
})
}
pub fn expression(&self) -> &str {
&self.expression
}
pub fn next_run_time(&self, after: DateTime<Utc>) -> Option<DateTime<Utc>> {
self.schedule.as_ref()?.after(&after).next()
}
pub fn next_run_times(&self, after: DateTime<Utc>, count: usize) -> Vec<DateTime<Utc>> {
self.schedule.as_ref()
.map(|s| s.after(&after).take(count).collect())
.unwrap_or_default()
}
pub fn should_run(&self, at: DateTime<Utc>) -> bool {
if let Some(next) = self.next_run_time(at - chrono::Duration::seconds(1)) {
(next - at).num_seconds().abs() < 30 } else {
false
}
}
}
impl<'de> Deserialize<'de> for CronExpression {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct CronExpressionData {
expression: String,
}
let data = CronExpressionData::deserialize(deserializer)?;
CronExpression::new(&data.expression)
.map_err(|e| serde::de::Error::custom(e.to_string()))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RetryStrategy {
Fixed {
delay: Duration,
max_attempts: u32,
},
Exponential {
initial_delay: Duration,
multiplier: f64,
max_delay: Duration,
max_attempts: u32,
jitter: bool,
},
Linear {
initial_delay: Duration,
increment: Duration,
max_delay: Duration,
max_attempts: u32,
},
Custom {
delays: Vec<Duration>,
},
}
impl Default for RetryStrategy {
fn default() -> Self {
RetryStrategy::Exponential {
initial_delay: Duration::from_secs(1),
multiplier: 2.0,
max_delay: Duration::from_secs(300), max_attempts: 3,
jitter: true,
}
}
}
impl RetryStrategy {
pub fn delay_for_attempt(&self, attempt: u32) -> Option<Duration> {
use rand::Rng;
match self {
RetryStrategy::Fixed { delay, max_attempts } => {
if attempt < *max_attempts {
Some(*delay)
} else {
None
}
}
RetryStrategy::Exponential {
initial_delay,
multiplier,
max_delay,
max_attempts,
jitter,
} => {
if attempt >= *max_attempts {
return None;
}
let delay = initial_delay.as_secs_f64() * multiplier.powi(attempt as i32);
let delay = delay.min(max_delay.as_secs_f64());
let delay = if *jitter {
let mut rng = rand::thread_rng();
let jitter_factor = rng.gen_range(0.75..1.25);
delay * jitter_factor
} else {
delay
};
Some(Duration::from_secs_f64(delay))
}
RetryStrategy::Linear {
initial_delay,
increment,
max_delay,
max_attempts,
} => {
if attempt >= *max_attempts {
return None;
}
let delay = initial_delay.as_secs() + (increment.as_secs() * attempt as u64);
let delay = delay.min(max_delay.as_secs());
Some(Duration::from_secs(delay))
}
RetryStrategy::Custom { delays } => {
delays.get(attempt as usize).copied()
}
}
}
pub fn max_attempts(&self) -> u32 {
match self {
RetryStrategy::Fixed { max_attempts, .. } => *max_attempts,
RetryStrategy::Exponential { max_attempts, .. } => *max_attempts,
RetryStrategy::Linear { max_attempts, .. } => *max_attempts,
RetryStrategy::Custom { delays } => delays.len() as u32,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScheduledJob {
pub id: String,
pub cron: CronExpression,
pub job_type: String,
pub payload: serde_json::Value,
pub priority: Priority,
pub retry_strategy: RetryStrategy,
pub timeout: Duration,
pub enabled: bool,
pub description: Option<String>,
pub next_run: Option<DateTime<Utc>>,
pub last_run: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
}
impl ScheduledJob {
pub fn new<T: Job>(
id: String,
cron_expr: &str,
job: T,
priority: Option<Priority>,
retry_strategy: Option<RetryStrategy>,
) -> ScheduleResult<Self> {
let cron = CronExpression::new(cron_expr)?;
let now = Utc::now();
let next_run = cron.next_run_time(now);
let job_type = job.job_type().to_string();
let timeout = job.timeout();
let payload = serde_json::to_value(job)
.map_err(|e| ScheduleError::Queue(QueueError::Serialization(e)))?;
Ok(ScheduledJob {
id,
cron,
job_type,
payload,
priority: priority.unwrap_or_default(),
retry_strategy: retry_strategy.unwrap_or_default(),
timeout,
enabled: true,
description: None,
next_run,
last_run: None,
created_at: now,
})
}
pub fn update_next_run(&mut self) {
let after = self.last_run.unwrap_or_else(Utc::now);
self.next_run = self.cron.next_run_time(after);
}
pub fn should_run(&self) -> bool {
if !self.enabled {
return false;
}
if let Some(next_run) = self.next_run {
next_run <= Utc::now()
} else {
false
}
}
pub fn mark_executed(&mut self) {
self.last_run = Some(Utc::now());
self.update_next_run();
}
pub fn create_job_entry(&self) -> QueueResult<JobEntry> {
JobEntry::new_with_job_type(
self.job_type.clone(),
self.payload.clone(),
Some(self.priority),
None, self.retry_strategy.max_attempts(),
)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ScheduledJobWrapper {
job_type: String,
payload: serde_json::Value,
max_retries: u32,
timeout: Duration,
}
#[async_trait::async_trait]
impl Job for ScheduledJobWrapper {
async fn execute(&self) -> crate::JobResult<()> {
Ok(())
}
fn job_type(&self) -> &'static str {
"scheduled_job_wrapper"
}
fn max_retries(&self) -> u32 {
self.max_retries
}
fn timeout(&self) -> Duration {
self.timeout
}
}
pub struct JobScheduler<B: crate::QueueBackend> {
backend: std::sync::Arc<B>,
schedules: std::sync::Arc<parking_lot::RwLock<std::collections::HashMap<String, ScheduledJob>>>,
running: std::sync::Arc<std::sync::atomic::AtomicBool>,
}
impl<B: crate::QueueBackend + 'static> JobScheduler<B> {
pub fn new(backend: std::sync::Arc<B>) -> Self {
Self {
backend,
schedules: std::sync::Arc::new(parking_lot::RwLock::new(std::collections::HashMap::new())),
running: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
}
}
pub fn add_schedule(&self, schedule: ScheduledJob) -> ScheduleResult<()> {
let mut schedules = self.schedules.write();
schedules.insert(schedule.id.clone(), schedule);
Ok(())
}
pub fn remove_schedule(&self, id: &str) -> ScheduleResult<bool> {
let mut schedules = self.schedules.write();
Ok(schedules.remove(id).is_some())
}
pub fn get_schedule(&self, id: &str) -> Option<ScheduledJob> {
let schedules = self.schedules.read();
schedules.get(id).cloned()
}
pub fn list_schedules(&self) -> Vec<ScheduledJob> {
let schedules = self.schedules.read();
schedules.values().cloned().collect()
}
pub fn set_schedule_enabled(&self, id: &str, enabled: bool) -> ScheduleResult<bool> {
let mut schedules = self.schedules.write();
if let Some(schedule) = schedules.get_mut(id) {
schedule.enabled = enabled;
Ok(true)
} else {
Ok(false)
}
}
pub async fn start(&self) -> ScheduleResult<()> {
self.running.store(true, std::sync::atomic::Ordering::SeqCst);
let backend = self.backend.clone();
let schedules = self.schedules.clone();
let running = self.running.clone();
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(30));
while running.load(std::sync::atomic::Ordering::SeqCst) {
interval.tick().await;
let mut due_schedules = Vec::new();
{
let mut schedules_guard = schedules.write();
for schedule in schedules_guard.values_mut() {
if schedule.should_run() {
schedule.mark_executed();
due_schedules.push(schedule.clone());
}
}
}
for schedule in due_schedules {
if let Ok(job_entry) = schedule.create_job_entry() {
if let Err(e) = backend.enqueue(job_entry).await {
tracing::error!("Failed to enqueue scheduled job {}: {}", schedule.id, e);
} else {
tracing::info!("Enqueued scheduled job: {}", schedule.id);
}
}
}
}
});
Ok(())
}
pub fn stop(&self) {
self.running.store(false, std::sync::atomic::Ordering::SeqCst);
}
pub fn is_running(&self) -> bool {
self.running.load(std::sync::atomic::Ordering::SeqCst)
}
pub async fn get_dead_jobs(&self, limit: Option<usize>) -> QueueResult<Vec<crate::JobEntry>> {
self.backend.get_jobs_by_state(crate::JobState::Dead, limit).await
}
pub async fn requeue_dead_job(&self, job_id: crate::JobId) -> QueueResult<bool> {
if let Some(job) = self.backend.get_job(job_id).await? {
if job.state() == &crate::JobState::Dead {
self.backend.requeue_job(job_id, job).await
} else {
Ok(false)
}
} else {
Ok(false)
}
}
pub async fn clear_dead_jobs(&self) -> QueueResult<u64> {
self.backend.clear_jobs_by_state(crate::JobState::Dead).await
}
}
pub mod cron_presets {
use super::CronExpression;
pub fn every_minute() -> CronExpression {
CronExpression::new("0 * * * * *").expect("Invalid 'every_minute' cron preset")
}
pub fn every_5_minutes() -> CronExpression {
CronExpression::new("0 */5 * * * *").unwrap()
}
pub fn every_15_minutes() -> CronExpression {
CronExpression::new("0 */15 * * * *").unwrap()
}
pub fn every_30_minutes() -> CronExpression {
CronExpression::new("0 */30 * * * *").unwrap()
}
pub fn hourly() -> CronExpression {
CronExpression::new("0 0 * * * *").unwrap()
}
pub fn daily() -> CronExpression {
CronExpression::new("0 0 0 * * *").unwrap()
}
pub fn weekly() -> CronExpression {
CronExpression::new("0 0 0 * * SUN").unwrap()
}
pub fn monthly() -> CronExpression {
CronExpression::new("0 0 0 1 * *").unwrap()
}
pub fn weekdays_at_9am() -> CronExpression {
CronExpression::new("0 0 9 * * 1-5").unwrap()
}
pub fn custom(expression: &str) -> Result<CronExpression, super::ScheduleError> {
CronExpression::new(expression)
}
}
#[derive(Debug, Clone)]
pub struct CancellationToken {
cancelled: Arc<std::sync::atomic::AtomicBool>,
notify: Arc<tokio::sync::Notify>,
}
impl CancellationToken {
pub fn new() -> Self {
Self {
cancelled: Arc::new(std::sync::atomic::AtomicBool::new(false)),
notify: Arc::new(tokio::sync::Notify::new()),
}
}
pub fn cancel(&self) {
self.cancelled.store(true, std::sync::atomic::Ordering::SeqCst);
self.notify.notify_waiters();
}
pub fn is_cancelled(&self) -> bool {
self.cancelled.load(std::sync::atomic::Ordering::SeqCst)
}
pub async fn wait_for_cancellation(&self) {
if self.is_cancelled() {
return;
}
self.notify.notified().await;
}
pub async fn cancelled(&self) {
self.wait_for_cancellation().await;
}
}
impl Default for CancellationToken {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct JobCancellationManager {
active_tokens: Arc<parking_lot::RwLock<std::collections::HashMap<crate::JobId, CancellationToken>>>,
}
impl JobCancellationManager {
pub fn new() -> Self {
Self {
active_tokens: Arc::new(parking_lot::RwLock::new(std::collections::HashMap::new())),
}
}
pub fn register_job(&self, job_id: crate::JobId) -> CancellationToken {
let token = CancellationToken::new();
self.active_tokens.write().insert(job_id, token.clone());
token
}
pub fn cancel_job(&self, job_id: crate::JobId) -> bool {
if let Some(token) = self.active_tokens.read().get(&job_id) {
token.cancel();
true
} else {
false
}
}
pub fn cancel_all(&self) {
let tokens = self.active_tokens.read();
for token in tokens.values() {
token.cancel();
}
}
pub fn unregister_job(&self, job_id: crate::JobId) {
self.active_tokens.write().remove(&job_id);
}
pub fn active_job_count(&self) -> usize {
self.active_tokens.read().len()
}
pub fn active_jobs(&self) -> Vec<crate::JobId> {
self.active_tokens.read().keys().cloned().collect()
}
}
impl Default for JobCancellationManager {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
pub trait CancellableJob: Job {
async fn execute_with_cancellation(&self, token: &CancellationToken) -> crate::JobResult<()>;
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct JobMetrics {
pub total_scheduled: u64,
pub total_executed: u64,
pub successful_jobs: u64,
pub failed_jobs: u64,
pub retried_jobs: u64,
pub timeout_jobs: u64,
pub cancelled_jobs: u64,
pub avg_execution_time_ms: f64,
pub min_execution_time_ms: u64,
pub max_execution_time_ms: u64,
pub jobs_by_priority: std::collections::HashMap<String, u64>,
pub jobs_by_type: std::collections::HashMap<String, u64>,
pub success_rate: f64,
pub avg_retry_attempts: f64,
pub last_reset: DateTime<Utc>,
}
impl JobMetrics {
pub fn new() -> Self {
Self {
last_reset: Utc::now(),
..Default::default()
}
}
pub fn record_scheduled(&mut self, job_type: &str, priority: Priority) {
self.total_scheduled += 1;
*self.jobs_by_type.entry(job_type.to_string()).or_insert(0) += 1;
*self.jobs_by_priority.entry(format!("{:?}", priority)).or_insert(0) += 1;
}
pub fn record_execution_start(&mut self) {
self.total_executed += 1;
}
pub fn record_success(&mut self, execution_time_ms: u64) {
self.successful_jobs += 1;
self.update_execution_time(execution_time_ms);
self.update_success_rate();
}
pub fn record_failure(&mut self, execution_time_ms: u64, retry_attempts: u32) {
self.failed_jobs += 1;
self.update_execution_time(execution_time_ms);
self.update_success_rate();
self.update_retry_attempts(retry_attempts);
}
pub fn record_retry(&mut self) {
self.retried_jobs += 1;
}
pub fn record_timeout(&mut self, execution_time_ms: u64) {
self.timeout_jobs += 1;
self.update_execution_time(execution_time_ms);
}
pub fn record_cancellation(&mut self, execution_time_ms: u64) {
self.cancelled_jobs += 1;
self.update_execution_time(execution_time_ms);
}
pub fn reset(&mut self) {
*self = Self::new();
}
fn update_execution_time(&mut self, execution_time_ms: u64) {
if self.min_execution_time_ms == 0 || execution_time_ms < self.min_execution_time_ms {
self.min_execution_time_ms = execution_time_ms;
}
if execution_time_ms > self.max_execution_time_ms {
self.max_execution_time_ms = execution_time_ms;
}
let completed_jobs = self.successful_jobs + self.failed_jobs + self.timeout_jobs + self.cancelled_jobs;
if completed_jobs > 0 {
let new_sample = execution_time_ms as f64;
self.avg_execution_time_ms += (new_sample - self.avg_execution_time_ms) / completed_jobs as f64;
}
}
fn update_success_rate(&mut self) {
let total_completed = self.successful_jobs + self.failed_jobs + self.timeout_jobs + self.cancelled_jobs;
if total_completed > 0 {
self.success_rate = self.successful_jobs as f64 / total_completed as f64;
}
}
fn update_retry_attempts(&mut self, attempts: u32) {
if self.failed_jobs > 0 {
let new_sample = attempts as f64;
self.avg_retry_attempts += (new_sample - self.avg_retry_attempts) / self.failed_jobs as f64;
}
}
}
#[derive(Debug)]
pub struct JobMetricsCollector {
metrics: Arc<parking_lot::RwLock<JobMetrics>>,
active_executions: Arc<parking_lot::RwLock<std::collections::HashMap<crate::JobId, std::time::Instant>>>,
}
impl JobMetricsCollector {
pub fn new() -> Self {
Self {
metrics: Arc::new(parking_lot::RwLock::new(JobMetrics::new())),
active_executions: Arc::new(parking_lot::RwLock::new(std::collections::HashMap::new())),
}
}
pub fn record_job_scheduled(&self, job_type: &str, priority: Priority) {
let mut metrics = self.metrics.write();
metrics.record_scheduled(job_type, priority);
}
pub fn record_execution_start(&self, job_id: crate::JobId) {
let mut metrics = self.metrics.write();
let mut executions = self.active_executions.write();
metrics.record_execution_start();
executions.insert(job_id, std::time::Instant::now());
}
pub fn record_job_success(&self, job_id: crate::JobId) {
let execution_time = self.get_and_remove_execution_time(job_id);
let mut metrics = self.metrics.write();
metrics.record_success(execution_time);
}
pub fn record_job_failure(&self, job_id: crate::JobId, retry_attempts: u32) {
let execution_time = self.get_and_remove_execution_time(job_id);
let mut metrics = self.metrics.write();
metrics.record_failure(execution_time, retry_attempts);
}
pub fn record_job_retry(&self, _job_id: crate::JobId) {
let mut metrics = self.metrics.write();
metrics.record_retry();
}
pub fn record_job_timeout(&self, job_id: crate::JobId) {
let execution_time = self.get_and_remove_execution_time(job_id);
let mut metrics = self.metrics.write();
metrics.record_timeout(execution_time);
}
pub fn record_job_cancellation(&self, job_id: crate::JobId) {
let execution_time = self.get_and_remove_execution_time(job_id);
let mut metrics = self.metrics.write();
metrics.record_cancellation(execution_time);
}
pub fn get_metrics(&self) -> JobMetrics {
self.metrics.read().clone()
}
pub fn reset_metrics(&self) {
let mut metrics = self.metrics.write();
let mut executions = self.active_executions.write();
metrics.reset();
executions.clear();
}
fn get_and_remove_execution_time(&self, job_id: crate::JobId) -> u64 {
let mut executions = self.active_executions.write();
if let Some(start_time) = executions.remove(&job_id) {
start_time.elapsed().as_millis() as u64
} else {
0
}
}
pub fn active_executions_count(&self) -> usize {
self.active_executions.read().len()
}
}
impl Default for JobMetricsCollector {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use crate::{MemoryBackend, QueueConfig};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct TestJob {
message: String,
}
#[async_trait::async_trait]
impl Job for TestJob {
async fn execute(&self) -> crate::JobResult<()> {
println!("Executing test job: {}", self.message);
Ok(())
}
fn job_type(&self) -> &'static str {
"test_job"
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct CancellableTestJob {
message: String,
sleep_duration: Duration,
}
#[async_trait::async_trait]
impl Job for CancellableTestJob {
async fn execute(&self) -> crate::JobResult<()> {
tokio::time::sleep(self.sleep_duration).await;
Ok(())
}
fn job_type(&self) -> &'static str {
"cancellable_test_job"
}
}
#[async_trait::async_trait]
impl CancellableJob for CancellableTestJob {
async fn execute_with_cancellation(&self, token: &CancellationToken) -> crate::JobResult<()> {
tokio::select! {
_ = tokio::time::sleep(self.sleep_duration) => {
println!("Job completed: {}", self.message);
Ok(())
}
_ = token.cancelled() => {
println!("Job cancelled: {}", self.message);
Err("Job was cancelled".into())
}
}
}
}
#[test]
fn test_cron_expression_validation() {
assert!(CronExpression::new("0 0 0 * * *").is_ok()); assert!(CronExpression::new("0 */5 * * * *").is_ok()); assert!(CronExpression::new("0 0 9-17 * * 1-5").is_ok());
assert!(CronExpression::new("invalid").is_err());
assert!(CronExpression::new("* * * * *").is_err()); }
#[test]
fn test_cron_next_run_time() {
let cron = CronExpression::new("0 0 0 * * *").unwrap(); let now = Utc::now();
let next = cron.next_run_time(now);
assert!(next.is_some());
assert!(next.unwrap() > now);
}
#[test]
fn test_cron_presets() {
assert!(cron_presets::every_minute().next_run_time(Utc::now()).is_some());
assert!(cron_presets::hourly().next_run_time(Utc::now()).is_some());
assert!(cron_presets::daily().next_run_time(Utc::now()).is_some());
assert!(cron_presets::weekly().next_run_time(Utc::now()).is_some());
assert!(cron_presets::monthly().next_run_time(Utc::now()).is_some());
assert!(cron_presets::weekdays_at_9am().next_run_time(Utc::now()).is_some());
}
#[test]
fn test_retry_strategy_exponential() {
let strategy = RetryStrategy::Exponential {
initial_delay: Duration::from_secs(1),
multiplier: 2.0,
max_delay: Duration::from_secs(60),
max_attempts: 3,
jitter: false,
};
assert_eq!(strategy.delay_for_attempt(0), Some(Duration::from_secs(1)));
assert_eq!(strategy.delay_for_attempt(1), Some(Duration::from_secs(2)));
assert_eq!(strategy.delay_for_attempt(2), Some(Duration::from_secs(4)));
assert_eq!(strategy.delay_for_attempt(3), None); assert_eq!(strategy.max_attempts(), 3);
}
#[test]
fn test_retry_strategy_linear() {
let strategy = RetryStrategy::Linear {
initial_delay: Duration::from_secs(5),
increment: Duration::from_secs(10),
max_delay: Duration::from_secs(60),
max_attempts: 4,
};
assert_eq!(strategy.delay_for_attempt(0), Some(Duration::from_secs(5)));
assert_eq!(strategy.delay_for_attempt(1), Some(Duration::from_secs(15)));
assert_eq!(strategy.delay_for_attempt(2), Some(Duration::from_secs(25)));
assert_eq!(strategy.delay_for_attempt(3), Some(Duration::from_secs(35)));
assert_eq!(strategy.delay_for_attempt(4), None); assert_eq!(strategy.max_attempts(), 4);
}
#[test]
fn test_retry_strategy_fixed() {
let strategy = RetryStrategy::Fixed {
delay: Duration::from_secs(10),
max_attempts: 2,
};
assert_eq!(strategy.delay_for_attempt(0), Some(Duration::from_secs(10)));
assert_eq!(strategy.delay_for_attempt(1), Some(Duration::from_secs(10)));
assert_eq!(strategy.delay_for_attempt(2), None); assert_eq!(strategy.max_attempts(), 2);
}
#[test]
fn test_retry_strategy_custom() {
let strategy = RetryStrategy::Custom {
delays: vec![
Duration::from_secs(1),
Duration::from_secs(5),
Duration::from_secs(30),
],
};
assert_eq!(strategy.delay_for_attempt(0), Some(Duration::from_secs(1)));
assert_eq!(strategy.delay_for_attempt(1), Some(Duration::from_secs(5)));
assert_eq!(strategy.delay_for_attempt(2), Some(Duration::from_secs(30)));
assert_eq!(strategy.delay_for_attempt(3), None); assert_eq!(strategy.max_attempts(), 3);
}
#[test]
fn test_scheduled_job_creation() {
let job = TestJob {
message: "Hello, World!".to_string(),
};
let scheduled = ScheduledJob::new(
"test_schedule".to_string(),
"0 0 0 * * *", job,
Some(Priority::High),
None,
).unwrap();
assert_eq!(scheduled.id, "test_schedule");
assert_eq!(scheduled.job_type, "test_job");
assert_eq!(scheduled.priority, Priority::High);
assert!(scheduled.enabled);
assert!(scheduled.next_run.is_some());
}
#[tokio::test]
async fn test_job_scheduler_basic() {
let backend = std::sync::Arc::new(MemoryBackend::new(QueueConfig::default()));
let scheduler = JobScheduler::new(backend);
let job = TestJob {
message: "Scheduled job".to_string(),
};
let scheduled = ScheduledJob::new(
"test_schedule".to_string(),
"0 * * * * *", job,
Some(Priority::Normal),
None,
).unwrap();
scheduler.add_schedule(scheduled).unwrap();
let schedules = scheduler.list_schedules();
assert_eq!(schedules.len(), 1);
assert_eq!(schedules[0].id, "test_schedule");
let retrieved = scheduler.get_schedule("test_schedule");
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap().id, "test_schedule");
assert!(scheduler.remove_schedule("test_schedule").unwrap());
assert!(scheduler.get_schedule("test_schedule").is_none());
}
#[tokio::test]
async fn test_cancellation_token() {
let token = CancellationToken::new();
assert!(!token.is_cancelled());
token.cancel();
assert!(token.is_cancelled());
token.wait_for_cancellation().await;
let cloned = token.clone();
assert!(cloned.is_cancelled());
cloned.wait_for_cancellation().await; }
#[tokio::test]
async fn test_cancellation_token_async_notification() {
let token = CancellationToken::new();
let token_clone = token.clone();
let wait_task = tokio::spawn(async move {
token_clone.wait_for_cancellation().await;
"cancelled"
});
tokio::time::sleep(Duration::from_millis(1)).await;
assert!(!wait_task.is_finished());
token.cancel();
let result = tokio::time::timeout(Duration::from_millis(100), wait_task).await;
assert!(result.is_ok());
assert_eq!(result.unwrap().unwrap(), "cancelled");
}
#[test]
fn test_job_cancellation_manager() {
let manager = JobCancellationManager::new();
let job_id = crate::JobId::new_v4();
let token = manager.register_job(job_id);
assert_eq!(manager.active_job_count(), 1);
assert!(manager.active_jobs().contains(&job_id));
assert!(!token.is_cancelled());
assert!(manager.cancel_job(job_id));
assert!(token.is_cancelled());
manager.unregister_job(job_id);
assert_eq!(manager.active_job_count(), 0);
assert!(!manager.cancel_job(job_id));
}
#[test]
fn test_job_metrics() {
let mut metrics = JobMetrics::new();
metrics.record_scheduled("test_job", Priority::High);
metrics.record_scheduled("test_job", Priority::Normal);
metrics.record_scheduled("email_job", Priority::High);
assert_eq!(metrics.total_scheduled, 3);
assert_eq!(*metrics.jobs_by_type.get("test_job").unwrap(), 2);
assert_eq!(*metrics.jobs_by_type.get("email_job").unwrap(), 1);
assert_eq!(*metrics.jobs_by_priority.get("High").unwrap(), 2);
assert_eq!(*metrics.jobs_by_priority.get("Normal").unwrap(), 1);
metrics.record_execution_start();
metrics.record_execution_start();
assert_eq!(metrics.total_executed, 2);
metrics.record_success(100); assert_eq!(metrics.successful_jobs, 1);
assert_eq!(metrics.min_execution_time_ms, 100);
assert_eq!(metrics.max_execution_time_ms, 100);
assert_eq!(metrics.avg_execution_time_ms, 100.0);
metrics.record_failure(200, 2); assert_eq!(metrics.failed_jobs, 1);
assert_eq!(metrics.avg_retry_attempts, 2.0);
assert_eq!(metrics.max_execution_time_ms, 200);
assert_eq!(metrics.avg_execution_time_ms, 150.0);
assert_eq!(metrics.success_rate, 0.5); }
#[test]
fn test_job_metrics_collector() {
let collector = JobMetricsCollector::new();
let job_id = crate::JobId::new_v4();
collector.record_job_scheduled("test_job", Priority::High);
collector.record_execution_start(job_id);
assert_eq!(collector.active_executions_count(), 1);
std::thread::sleep(Duration::from_millis(10)); collector.record_job_success(job_id);
assert_eq!(collector.active_executions_count(), 0);
let metrics = collector.get_metrics();
assert_eq!(metrics.total_scheduled, 1);
assert_eq!(metrics.total_executed, 1);
assert_eq!(metrics.successful_jobs, 1);
assert_eq!(metrics.success_rate, 1.0);
assert!(metrics.min_execution_time_ms >= 10);
}
#[tokio::test]
async fn test_atomic_clear_dead_jobs() {
use crate::{MemoryBackend, QueueConfig, JobEntry, Priority, QueueBackend};
let backend = std::sync::Arc::new(MemoryBackend::new(QueueConfig::default()));
let scheduler = JobScheduler::new(backend.clone());
let job1 = TestJob { message: "job1".to_string() };
let job2 = TestJob { message: "job2".to_string() };
let job3 = TestJob { message: "job3".to_string() };
let mut entry1 = JobEntry::new(job1, Some(Priority::Normal), None).unwrap();
let mut entry2 = JobEntry::new(job2, Some(Priority::High), None).unwrap();
let mut entry3 = JobEntry::new(job3, Some(Priority::Low), None).unwrap();
for _ in 0..=3 {
entry1.mark_failed("Test failure".to_string());
entry2.mark_failed("Test failure".to_string());
entry3.mark_failed("Test failure".to_string());
}
backend.enqueue(entry1).await.unwrap();
backend.enqueue(entry2).await.unwrap();
backend.enqueue(entry3).await.unwrap();
let dead_jobs_before = scheduler.get_dead_jobs(None).await.unwrap();
assert_eq!(dead_jobs_before.len(), 3);
let cleared_count = scheduler.clear_dead_jobs().await.unwrap();
assert_eq!(cleared_count, 3);
let dead_jobs_after = scheduler.get_dead_jobs(None).await.unwrap();
assert_eq!(dead_jobs_after.len(), 0);
let stats = backend.stats().await.unwrap();
assert_eq!(stats.dead_jobs, 0);
}
#[tokio::test]
async fn test_atomic_requeue_dead_job() {
use crate::{MemoryBackend, QueueConfig, JobEntry, JobState, Priority, QueueBackend};
let backend = std::sync::Arc::new(MemoryBackend::new(QueueConfig::default()));
let scheduler = JobScheduler::new(backend.clone());
let job = TestJob { message: "dead job".to_string() };
let mut entry = JobEntry::new(job, Some(Priority::Normal), None).unwrap();
let job_id = entry.id();
for _ in 0..=3 {
entry.mark_failed("Test failure".to_string());
}
backend.enqueue(entry).await.unwrap();
let stats_before = backend.stats().await.unwrap();
assert_eq!(stats_before.dead_jobs, 1);
assert_eq!(stats_before.pending_jobs, 0);
let requeued = scheduler.requeue_dead_job(job_id).await.unwrap();
assert!(requeued);
let stats_after = backend.stats().await.unwrap();
assert_eq!(stats_after.dead_jobs, 0);
assert_eq!(stats_after.pending_jobs, 1);
let job_entry = backend.get_job(job_id).await.unwrap().unwrap();
assert_eq!(job_entry.state(), &JobState::Pending);
assert_eq!(job_entry.attempts(), 0);
}
}