From d85694384c52c8cd704007d5ffda9657fda55d08 Mon Sep 17 00:00:00 2001 From: JAVA-Lw <1264129499@qq.com> Date: Fri, 10 Jul 2026 17:55:51 +0800 Subject: [PATCH] feat(cron): deduplicate scheduled executions --- crates/aionui-api-types/src/cron.rs | 7 + crates/aionui-app/src/router/state.rs | 4 +- crates/aionui-app/tests/cron_e2e.rs | 1 + crates/aionui-cron/src/artifacts.rs | 1 + crates/aionui-cron/src/events.rs | 1 + crates/aionui-cron/src/executor.rs | 71 +++ crates/aionui-cron/src/scheduler.rs | 95 +++- crates/aionui-cron/src/service.rs | 275 ++++++++-- crates/aionui-cron/src/types.rs | 6 + .../aionui-cron/tests/service_integration.rs | 47 ++ .../migrations/021_cron_queue_enabled.sql | 1 + .../migrations/022_cron_job_runs.sql | 19 + crates/aionui-db/src/lib.rs | 4 +- crates/aionui-db/src/models/cron_job.rs | 3 + crates/aionui-db/src/repository/cron.rs | 67 +++ .../aionui-db/src/repository/sqlite_cron.rs | 468 +++++++++++++++++- crates/aionui-db/tests/cron_repository.rs | 1 + 17 files changed, 1013 insertions(+), 58 deletions(-) create mode 100644 crates/aionui-db/migrations/021_cron_queue_enabled.sql create mode 100644 crates/aionui-db/migrations/022_cron_job_runs.sql diff --git a/crates/aionui-api-types/src/cron.rs b/crates/aionui-api-types/src/cron.rs index 9d366747f..94eb82098 100644 --- a/crates/aionui-api-types/src/cron.rs +++ b/crates/aionui-api-types/src/cron.rs @@ -123,6 +123,7 @@ pub struct CronJobStateDto { pub run_count: i64, pub retry_count: i64, pub max_retries: i64, + pub queue_enabled: bool, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -160,6 +161,8 @@ pub struct CreateCronJobRequest { #[serde(default)] pub execution_mode: Option, #[serde(default)] + pub queue_enabled: bool, + #[serde(default)] pub agent_config: Option, } @@ -209,6 +212,8 @@ pub struct UpdateCronJobRequest { pub conversation_title: Option, #[serde(default)] pub max_retries: Option, + #[serde(default)] + pub queue_enabled: Option, } // --------------------------------------------------------------------------- @@ -609,6 +614,7 @@ mod tests { run_count: 5, retry_count: 0, max_retries: 3, + queue_enabled: false, }, } } @@ -677,6 +683,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: true, }, }; let json = serde_json::to_value(&resp).unwrap(); diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 50dfa76a9..f1854a87f 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -705,11 +705,11 @@ pub fn build_cron_state(services: &AppServices) -> CronRouterState { let tick_service_ref: Arc = Arc::new(CronServiceTickRef::default()); let tick_ref = tick_service_ref.clone(); let scheduler = Arc::new(aionui_cron::scheduler::CronScheduler::new(Arc::new( - move |job_id: String| { + move |tick: aionui_cron::scheduler::ScheduledTick| { let svc = tick_ref.0.lock().unwrap().clone(); tokio::spawn(async move { if let Some(svc) = svc { - svc.tick(&job_id).await; + svc.tick(&tick.job_id, tick.scheduled_at).await; } }); }, diff --git a/crates/aionui-app/tests/cron_e2e.rs b/crates/aionui-app/tests/cron_e2e.rs index cc92820c7..62f152766 100644 --- a/crates/aionui-app/tests/cron_e2e.rs +++ b/crates/aionui-app/tests/cron_e2e.rs @@ -402,6 +402,7 @@ async fn cj5b_run_now_legacy_workspace_with_whitespace_succeeds() { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, }) .await .unwrap(); diff --git a/crates/aionui-cron/src/artifacts.rs b/crates/aionui-cron/src/artifacts.rs index cdcf55c31..f0097f016 100644 --- a/crates/aionui-cron/src/artifacts.rs +++ b/crates/aionui-cron/src/artifacts.rs @@ -125,6 +125,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } } diff --git a/crates/aionui-cron/src/events.rs b/crates/aionui-cron/src/events.rs index 1cb13ab98..48d9701d4 100644 --- a/crates/aionui-cron/src/events.rs +++ b/crates/aionui-cron/src/events.rs @@ -138,6 +138,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, }, } } diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index f4b7c68c8..f899ea43d 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -112,6 +112,24 @@ impl JobExecutor { return Err(ExecutionResult::Error { message: e.to_string() }); } + if job.queue_enabled && job.execution_mode == ExecutionMode::NewConversation { + match self.find_active_cron_conversation(job).await { + Ok(Some(conversation_id)) => { + info!( + job_id = %job.id, + conversation_id, + "Cron queue is enabled and a previous execution is still active; skipping trigger" + ); + return Err(ExecutionResult::Skipped); + } + Ok(None) => {} + Err(e) => { + error!(job_id = %job.id, error = %e, "Failed to inspect previous cron executions"); + return Err(ExecutionResult::Error { message: e.to_string() }); + } + } + } + let conversation_id = match self.resolve_conversation(job, saved_skill.as_ref()).await { Ok(id) => id, Err(e) => { @@ -444,6 +462,11 @@ impl JobExecutor { impl JobExecutor { fn handle_busy(&self, job: &CronJob) -> ExecutionResult { + if job.queue_enabled { + info!(job_id = %job.id, "Cron queue is enabled; skipping overlapping execution"); + return ExecutionResult::Skipped; + } + let max_retries = job.max_retries; let current_retry = job.retry_count; @@ -504,6 +527,15 @@ impl JobExecutor { } } + async fn find_active_cron_conversation(&self, job: &CronJob) -> Result, CronError> { + let user_id = self.resolve_conversation_owner_user_id(job).await?; + let conversations = self.conversation_repo.list_by_cron_job(&user_id, &job.id).await?; + Ok(conversations + .into_iter() + .find(|conversation| self.is_conversation_claimed(&conversation.id)) + .map(|conversation| conversation.id)) + } + async fn create_new_conversation( &self, job: &CronJob, @@ -1243,6 +1275,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } } @@ -1268,6 +1301,7 @@ mod tests { let job = CronJob { retry_count: 1, max_retries: 3, + queue_enabled: false, ..sample_job() }; let result = executor.handle_busy(&job); @@ -1281,6 +1315,7 @@ mod tests { let job = CronJob { retry_count: 3, max_retries: 3, + queue_enabled: false, ..sample_job() }; let result = executor.handle_busy(&job); @@ -1294,6 +1329,7 @@ mod tests { let job = CronJob { retry_count: 5, max_retries: 3, + queue_enabled: false, ..sample_job() }; let result = executor.handle_busy(&job); @@ -1307,12 +1343,26 @@ mod tests { let job = CronJob { retry_count: 0, max_retries: 3, + queue_enabled: false, ..sample_job() }; let result = executor.handle_busy(&job); assert_eq!(result, ExecutionResult::Retrying { attempt: 1 }); } + #[tokio::test] + async fn handle_busy_returns_skipped_when_queue_is_enabled() { + let executor = make_executor_for_busy_tests(); + let job = CronJob { + retry_count: 0, + max_retries: 3, + queue_enabled: true, + ..sample_job() + }; + + assert_eq!(executor.handle_busy(&job), ExecutionResult::Skipped); + } + #[tokio::test] async fn execute_returns_retrying_when_runtime_state_is_already_claimed() { let agent = Arc::new(RecordingAgent::new("conv_1", "default", true)); @@ -1332,6 +1382,27 @@ mod tests { drop(claim); } + #[tokio::test] + async fn execute_returns_skipped_when_queue_is_enabled_and_conversation_is_claimed() { + let agent = Arc::new(RecordingAgent::new("conv_1", "default", true)); + let executor = make_executor_with_agent(AgentInstance::Mock(agent.clone())); + let job = CronJob { + queue_enabled: true, + ..sample_job() + }; + let claim = executor + .conversation_service + .runtime_state() + .try_claim_turn(&job.conversation_id, "turn-existing") + .expect("runtime claim should succeed"); + + let result = executor.execute(&job).await; + + assert_eq!(result, ExecutionResult::Skipped); + assert_eq!(agent.send_calls(), 0, "queue protection should avoid send attempts"); + drop(claim); + } + #[tokio::test] async fn prepare_run_now_returns_active_conversation_when_runtime_state_is_already_claimed() { let agent = Arc::new(RecordingAgent::new("conv_1", "default", true)); diff --git a/crates/aionui-cron/src/scheduler.rs b/crates/aionui-cron/src/scheduler.rs index a9e108061..f0e409df1 100644 --- a/crates/aionui-cron/src/scheduler.rs +++ b/crates/aionui-cron/src/scheduler.rs @@ -55,6 +55,25 @@ pub fn compute_next_run(schedule: &CronSchedule, now: TimestampMs) -> Option Option { + match schedule { + CronSchedule::At { .. } => None, + CronSchedule::Every { every_ms, .. } => { + if *every_ms <= 0 { + return None; + } + let elapsed = now.saturating_sub(scheduled_at); + let steps = elapsed.div_euclid(*every_ms) + 1; + scheduled_at.checked_add(steps.checked_mul(*every_ms)?) + } + CronSchedule::Cron { expr, tz, .. } => compute_cron_next_run(expr, tz.as_deref(), now.max(scheduled_at)), + } +} + fn compute_cron_next_run(expr: &str, tz: Option<&str>, now: TimestampMs) -> Option { let normalized = normalize_cron_expr(expr); let schedule = Schedule::from_str(&normalized).ok()?; @@ -101,7 +120,13 @@ pub fn validate_schedule(schedule: &CronSchedule) -> Result<(), CronError> { // CronScheduler — manages tokio timers for scheduled jobs // --------------------------------------------------------------------------- -pub type TickCallback = Arc; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ScheduledTick { + pub job_id: String, + pub scheduled_at: TimestampMs, +} + +pub type TickCallback = Arc; pub struct CronScheduler { handles: DashMap>, @@ -152,6 +177,17 @@ impl CronScheduler { self.schedule_job(job); } + pub fn schedule_retry(&self, job_id: &str, scheduled_at: TimestampMs, retry_at: TimestampMs) { + self.cancel_job(job_id); + let handle = spawn_occurrence_timer( + job_id.to_owned(), + scheduled_at, + retry_at, + Arc::clone(&self.tick_callback), + ); + self.handles.insert(job_id.to_owned(), handle); + } + pub fn cancel_all(&self) { for entry in self.handles.iter() { entry.value().abort(); @@ -179,12 +215,21 @@ impl Drop for CronScheduler { // --------------------------------------------------------------------------- fn spawn_at_timer(job_id: String, run_at: TimestampMs, callback: TickCallback) -> JoinHandle<()> { + spawn_occurrence_timer(job_id, run_at, run_at, callback) +} + +fn spawn_occurrence_timer( + job_id: String, + scheduled_at: TimestampMs, + run_at: TimestampMs, + callback: TickCallback, +) -> JoinHandle<()> { tokio::spawn(async move { let delay = delay_until(run_at); if delay > 0 { tokio::time::sleep(tokio::time::Duration::from_millis(delay as u64)).await; } - callback(job_id); + callback(ScheduledTick { job_id, scheduled_at }); }) } @@ -199,14 +244,22 @@ fn spawn_every_timer( if initial_delay > 0 { tokio::time::sleep(tokio::time::Duration::from_millis(initial_delay as u64)).await; } - callback(job_id.clone()); + let mut scheduled_at = first_run_at; + callback(ScheduledTick { + job_id: job_id.clone(), + scheduled_at, + }); let interval_duration = tokio::time::Duration::from_millis(every_ms as u64); let mut interval = tokio::time::interval(interval_duration); interval.tick().await; // first tick fires immediately, skip it loop { interval.tick().await; - callback(job_id.clone()); + scheduled_at += every_ms; + callback(ScheduledTick { + job_id: job_id.clone(), + scheduled_at, + }); } }) } @@ -223,7 +276,10 @@ fn spawn_cron_timer( if initial_delay > 0 { tokio::time::sleep(tokio::time::Duration::from_millis(initial_delay as u64)).await; } - callback(job_id.clone()); + callback(ScheduledTick { + job_id: job_id.clone(), + scheduled_at: first_run_at, + }); loop { let now = now_ms(); @@ -235,7 +291,10 @@ fn spawn_cron_timer( if delay > 0 { tokio::time::sleep(tokio::time::Duration::from_millis(delay as u64)).await; } - callback(job_id.clone()); + callback(ScheduledTick { + job_id: job_id.clone(), + scheduled_at: next_at, + }); } }) } @@ -300,6 +359,23 @@ mod tests { assert_eq!(compute_next_run(&schedule, 1000), None); } + #[test] + fn next_run_after_occurrence_keeps_every_schedule_grid() { + let schedule = CronSchedule::Every { + every_ms: 1_000, + description: None, + }; + + assert_eq!( + compute_next_run_after_occurrence(&schedule, 10_000, 10_100), + Some(11_000) + ); + assert_eq!( + compute_next_run_after_occurrence(&schedule, 10_000, 12_500), + Some(13_000) + ); + } + #[test] fn next_run_cron_returns_future_time() { let now = now_ms(); @@ -559,7 +635,7 @@ mod tests { #[tokio::test] async fn scheduler_at_timer_fires_callback() { - let (tx, rx) = tokio::sync::oneshot::channel::(); + let (tx, rx) = tokio::sync::oneshot::channel::(); let tx = Arc::new(std::sync::Mutex::new(Some(tx))); let scheduler = CronScheduler::new(Arc::new(move |id| { if let Some(sender) = tx.lock().unwrap().take() { @@ -579,7 +655,9 @@ mod tests { let result = tokio::time::timeout(tokio::time::Duration::from_secs(2), rx).await; assert!(result.is_ok()); - assert_eq!(result.unwrap().unwrap(), "cron_at"); + let tick = result.unwrap().unwrap(); + assert_eq!(tick.job_id, "cron_at"); + assert_eq!(tick.scheduled_at, job.next_run_at.unwrap()); } #[tokio::test] @@ -637,6 +715,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } } } diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 242423b51..43d7fe580 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -12,8 +12,9 @@ use aionui_common::{ validate_workspace_path_availability, }; use aionui_db::{ - IAgentMetadataRepository, IAssistantDefinitionRepository, IAssistantOverlayRepository, ICronRepository, - UpdateCronJobParams, models::AgentMetadataRow, resolve_agent_binding_from_rows, runtime_backend_for_agent, + ClaimCronRunParams, CronRunClaimResult, FinishCronRunParams, IAgentMetadataRepository, + IAssistantDefinitionRepository, IAssistantOverlayRepository, ICronRepository, UpdateCronJobParams, + models::AgentMetadataRow, resolve_agent_binding_from_rows, runtime_backend_for_agent, }; use tracing::{debug, error, info, warn}; @@ -21,7 +22,7 @@ use crate::events::CronEventEmitter; use crate::error::CronError; use crate::executor::{ExecutionResult, JobExecutor, PreparedRunNow, RETRY_INTERVAL_MS}; -use crate::scheduler::{CronScheduler, compute_next_run, validate_schedule}; +use crate::scheduler::{CronScheduler, compute_next_run, compute_next_run_after_occurrence, validate_schedule}; use crate::skill_file::{delete_skill_file, has_skill_file, write_raw_skill_file, write_skill_file}; use crate::types::{ CreatedBy, CronAgentConfig, CronJob, CronSchedule, ExecutionMode, cron_job_from_row, cron_job_to_response, @@ -40,6 +41,9 @@ const PLACEHOLDER_PATTERNS: &[&str] = &[ "write your", "put your", ]; +const RUN_LEASE_MS: i64 = 60_000; +const RUN_LEASE_HEARTBEAT_MS: u64 = 20_000; +const RUN_HISTORY_RETENTION_MS: i64 = 30 * 24 * 60 * 60 * 1_000; #[derive(Clone)] pub struct CronService { repo: Arc, @@ -50,6 +54,7 @@ pub struct CronService { executor: Arc, emitter: CronEventEmitter, data_dir: PathBuf, + instance_id: String, } pub struct CronServiceDeps { @@ -74,6 +79,7 @@ impl CronService { executor: deps.executor, emitter: deps.emitter, data_dir: deps.data_dir, + instance_id: generate_prefixed_id("cron-owner"), } } @@ -114,6 +120,7 @@ impl CronService { conversation_title, created_by: "agent".to_owned(), execution_mode: Some("existing".to_owned()), + queue_enabled: false, agent_config, }; @@ -197,6 +204,7 @@ impl CronService { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }, ) .await?; @@ -290,6 +298,7 @@ impl CronService { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: req.queue_enabled, }; self.validate_job_workspace(&job).await?; @@ -370,6 +379,9 @@ impl CronService { if let Some(max_retries) = req.max_retries { job.max_retries = max_retries; } + if let Some(queue_enabled) = req.queue_enabled { + job.queue_enabled = queue_enabled; + } if req.schedule.is_some() || req.enabled.is_some() { job.next_run_at = compute_next_run(&job.schedule, now_ms()); @@ -460,6 +472,11 @@ impl CronService { // ----------------------------------------------------------------------- pub async fn init(&self) { + let now = now_ms(); + if let Err(error) = self.repo.cleanup_runs_before(now - RUN_HISTORY_RETENTION_MS).await { + warn!(error = %error, "Failed to clean up old cron run records"); + } + let rows = match self.repo.list_enabled().await { Ok(rows) => rows, Err(e) => { @@ -479,14 +496,29 @@ impl CronService { } }; - self.scheduler.schedule_job(&job); + match self.repo.get_recoverable_run(&job.id, now).await { + Ok(Some(run)) => { + info!( + job_id = %job.id, + scheduled_at = run.scheduled_at, + wake_at = run.wake_at, + "Recovering unfinished cron occurrence" + ); + self.scheduler.schedule_retry(&job.id, run.scheduled_at, run.wake_at); + } + Ok(None) => self.scheduler.schedule_job(&job), + Err(error) => { + warn!(job_id = %job.id, error = %error, "Failed to inspect recoverable cron run"); + self.scheduler.schedule_job(&job); + } + } scheduled += 1; } info!(scheduled, "Cron service initialized"); } - pub async fn tick(&self, job_id: &str) { + pub async fn tick(&self, job_id: &str, scheduled_at: i64) { let row = match self.repo.get_by_id(job_id).await { Ok(Some(r)) => r, Ok(None) => { @@ -520,10 +552,53 @@ impl CronService { return; } + let claim_now = now_ms(); + match self + .repo + .claim_run(&ClaimCronRunParams { + job_id, + scheduled_at, + owner_id: &self.instance_id, + now: claim_now, + lease_until: claim_now + RUN_LEASE_MS, + queue_enabled: job.queue_enabled, + }) + .await + { + Ok(CronRunClaimResult::Claimed) => {} + Ok(CronRunClaimResult::Duplicate) => { + info!(job_id, scheduled_at, "Duplicate cron occurrence ignored"); + match self.repo.get_recoverable_run(job_id, claim_now).await { + Ok(Some(run)) => { + self.scheduler.schedule_retry(job_id, run.scheduled_at, run.wake_at); + } + Ok(None) => self.reschedule_after_execution(&job, scheduled_at).await, + Err(error) => { + warn!(job_id, scheduled_at, error = %error, "Failed to track duplicate cron occurrence"); + self.reschedule_after_execution(&job, scheduled_at).await; + } + } + return; + } + Ok(CronRunClaimResult::QueueBusy) => { + self.record_queue_busy_skip(&job).await; + self.reschedule_after_execution(&job, scheduled_at).await; + self.emitter.emit_job_executed(job_id, "skipped", None); + return; + } + Err(error) => { + error!(job_id, scheduled_at, error = %error, "Failed to claim cron occurrence"); + return; + } + } + + let heartbeat = self.spawn_run_lease_heartbeat(job_id.to_owned(), scheduled_at); + let prepared = match self.executor.prepare_scheduled(&job).await { Ok(prepared) => prepared, Err(result) => { - self.handle_execution_result(job, result).await; + heartbeat.abort(); + self.handle_execution_result(job, scheduled_at, result).await; return; } }; @@ -538,8 +613,10 @@ impl CronService { error = %err, "Failed to bind materialized cron replacement conversation before execution" ); + heartbeat.abort(); self.handle_execution_result( execution_job, + scheduled_at, ExecutionResult::Error { message: err.to_string(), }, @@ -549,7 +626,8 @@ impl CronService { } let result = self.executor.execute_prepared_scheduled(&execution_job, prepared).await; - self.handle_execution_result(execution_job, result).await; + heartbeat.abort(); + self.handle_execution_result(execution_job, scheduled_at, result).await; } pub async fn handle_system_resume(&self) { @@ -572,6 +650,21 @@ impl CronService { } }; + match self.repo.get_recoverable_run(&job.id, now).await { + Ok(Some(run)) => { + self.scheduler.schedule_retry(&job.id, run.scheduled_at, run.wake_at); + continue; + } + Ok(None) => {} + Err(error) => { + warn!( + job_id = %job.id, + error = %error, + "Resume: failed to inspect recoverable cron run" + ); + } + } + if let Some(next_run) = job.next_run_at && next_run < now { @@ -850,16 +943,25 @@ impl CronService { true } - async fn handle_execution_result(&self, job: CronJob, result: ExecutionResult) { + async fn handle_execution_result(&self, job: CronJob, scheduled_at: i64, result: ExecutionResult) { let job_id = &job.id; match result { ExecutionResult::Success { conversation_id } => { + if !self + .finish_scheduled_run(job_id, scheduled_at, "ok", Some(&conversation_id), None) + .await + { + return; + } self.update_job_after_success(job_id, &conversation_id).await; - self.reschedule_after_execution(&job).await; + self.reschedule_after_execution(&job, scheduled_at).await; self.emitter.emit_job_executed(job_id, "ok", None); } ExecutionResult::Retrying { attempt } => { + if !self.schedule_retry(job_id, scheduled_at, attempt).await { + return; + } let params = UpdateCronJobParams { retry_count: Some(attempt), ..Default::default() @@ -867,9 +969,14 @@ impl CronService { if let Err(e) = self.repo.update(job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update retry count"); } - self.schedule_retry(job_id, attempt); } ExecutionResult::Skipped => { + if !self + .finish_scheduled_run(job_id, scheduled_at, "skipped", None, None) + .await + { + return; + } let params = UpdateCronJobParams { last_status: Some(Some("skipped".into())), retry_count: Some(0), @@ -878,12 +985,18 @@ impl CronService { if let Err(e) = self.repo.update(job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update skipped status"); } - self.reschedule_after_execution(&job).await; + self.reschedule_after_execution(&job, scheduled_at).await; self.emitter.emit_job_executed(job_id, "skipped", None); } ExecutionResult::Error { message } => { + if !self + .finish_scheduled_run(job_id, scheduled_at, "error", None, Some(&message)) + .await + { + return; + } self.update_job_after_error(job_id, &message).await; - self.reschedule_after_execution(&job).await; + self.reschedule_after_execution(&job, scheduled_at).await; self.emitter.emit_job_executed(job_id, "error", Some(&message)); } } @@ -1000,7 +1113,18 @@ impl CronService { } } - async fn reschedule_after_execution(&self, job: &CronJob) { + async fn reschedule_after_execution(&self, job: &CronJob, scheduled_at: i64) { + match self.repo.get_recoverable_run(&job.id, now_ms()).await { + Ok(Some(run)) => { + self.scheduler.schedule_retry(&job.id, run.scheduled_at, run.wake_at); + return; + } + Ok(None) => {} + Err(error) => { + warn!(job_id = %job.id, error = %error, "Failed to inspect remaining cron occurrences"); + } + } + let is_at = matches!(job.schedule, CronSchedule::At { .. }); if is_at { let params = UpdateCronJobParams { @@ -1024,8 +1148,7 @@ impl CronService { return; } - let now = now_ms(); - let next = compute_next_run(&job.schedule, now); + let next = compute_next_run_after_occurrence(&job.schedule, scheduled_at, now_ms()); let updated = CronJob { next_run_at: next, ..job.clone() @@ -1125,36 +1248,95 @@ impl CronService { self.scheduler.reschedule_job(&updated); } - fn schedule_retry(&self, job_id: &str, _attempt: i64) { + async fn schedule_retry(&self, job_id: &str, scheduled_at: i64, _attempt: i64) -> bool { let next_run = now_ms() + RETRY_INTERVAL_MS as i64; - let retry_job = CronJob { - id: job_id.to_owned(), - name: String::new(), - enabled: true, - schedule: CronSchedule::At { - at_ms: next_run, - description: None, - }, - message: String::new(), - execution_mode: ExecutionMode::Existing, - agent_config: None, - conversation_id: String::new(), - conversation_title: None, - agent_type: String::new(), - created_by: CreatedBy::User, - skill_content: None, - description: None, - created_at: 0, - updated_at: 0, - next_run_at: Some(next_run), - last_run_at: None, - last_status: None, - last_error: None, - run_count: 0, - retry_count: 0, - max_retries: 0, + match self + .repo + .defer_run(job_id, scheduled_at, &self.instance_id, next_run, now_ms()) + .await + { + Ok(true) => { + self.scheduler.schedule_retry(job_id, scheduled_at, next_run); + true + } + Ok(false) => { + warn!(job_id, scheduled_at, "Cron retry lease was no longer owned"); + false + } + Err(error) => { + error!(job_id, scheduled_at, error = %error, "Failed to defer cron run"); + false + } + } + } + + fn spawn_run_lease_heartbeat(&self, job_id: String, scheduled_at: i64) -> tokio::task::JoinHandle<()> { + let repo = Arc::clone(&self.repo); + let owner_id = self.instance_id.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(RUN_LEASE_HEARTBEAT_MS)); + interval.tick().await; + loop { + interval.tick().await; + let now = now_ms(); + match repo + .renew_run_lease(&job_id, scheduled_at, &owner_id, now + RUN_LEASE_MS, now) + .await + { + Ok(true) => {} + Ok(false) => break, + Err(error) => warn!(job_id, scheduled_at, error = %error, "Failed to renew cron run lease"), + } + } + }) + } + + async fn finish_scheduled_run( + &self, + job_id: &str, + scheduled_at: i64, + status: &str, + conversation_id: Option<&str>, + error: Option<&str>, + ) -> bool { + let finished_at = now_ms(); + match self + .repo + .finish_run(&FinishCronRunParams { + job_id, + scheduled_at, + owner_id: &self.instance_id, + status, + conversation_id, + error, + finished_at, + }) + .await + { + Ok(true) => true, + Ok(false) => { + warn!( + job_id, + scheduled_at, status, "Cron run lease was no longer owned at completion" + ); + false + } + Err(error) => { + error!(job_id, scheduled_at, status, error = %error, "Failed to finish cron run"); + false + } + } + } + + async fn record_queue_busy_skip(&self, job: &CronJob) { + let params = UpdateCronJobParams { + last_status: Some(Some("skipped".into())), + retry_count: Some(0), + ..Default::default() }; - self.scheduler.schedule_job(&retry_job); + if let Err(error) = self.repo.update(&job.id, ¶ms).await { + error!(job_id = %job.id, error = %error, "Failed to record queue-busy cron skip"); + } } pub async fn delete_jobs_by_conversation(&self, conversation_id: &str) { @@ -1835,6 +2017,7 @@ fn build_update_params(job: &CronJob, req: &UpdateCronJobRequest) -> UpdateCronJ last_error: None, run_count: None, retry_count: None, + queue_enabled: req.queue_enabled, } } @@ -2076,6 +2259,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } } @@ -2092,6 +2276,7 @@ mod tests { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let params = build_update_params(&job, &req); assert_eq!(params.name.as_deref(), Some("New Name")); @@ -2125,6 +2310,7 @@ mod tests { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let params = build_update_params(&job, &req); assert_eq!(params.schedule_kind.as_deref(), Some("cron")); @@ -2166,6 +2352,7 @@ mod tests { }), conversation_title: None, max_retries: None, + queue_enabled: None, }; let params = build_update_params(&job, &req); @@ -2228,6 +2415,7 @@ mod tests { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let params = build_update_params(&job, &req); assert_eq!(params.enabled, Some(false)); @@ -2247,6 +2435,7 @@ mod tests { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let params = build_update_params(&job, &req); assert_eq!( diff --git a/crates/aionui-cron/src/types.rs b/crates/aionui-cron/src/types.rs index 36f00c100..b3f742442 100644 --- a/crates/aionui-cron/src/types.rs +++ b/crates/aionui-cron/src/types.rs @@ -174,6 +174,7 @@ pub struct CronJob { pub run_count: i64, pub retry_count: i64, pub max_retries: i64, + pub queue_enabled: bool, } // --------------------------------------------------------------------------- @@ -222,6 +223,7 @@ pub fn cron_job_from_row(row: CronJobRow) -> Result { run_count: row.run_count, retry_count: row.retry_count, max_retries: row.max_retries, + queue_enabled: row.queue_enabled, }) } @@ -293,6 +295,7 @@ pub fn cron_job_to_row(job: &CronJob) -> Result { run_count: job.run_count, retry_count: job.retry_count, max_retries: job.max_retries, + queue_enabled: job.queue_enabled, }) } @@ -379,6 +382,7 @@ pub fn cron_job_to_response(job: &CronJob) -> CronJobResponse { run_count: job.run_count, retry_count: job.retry_count, max_retries: job.max_retries, + queue_enabled: job.queue_enabled, }, } } @@ -575,6 +579,7 @@ mod tests { run_count: 5, retry_count: 0, max_retries: 3, + queue_enabled: false, } } @@ -616,6 +621,7 @@ mod tests { run_count: 5, retry_count: 0, max_retries: 3, + queue_enabled: false, } } diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index ef5873ac5..86501f0d8 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -953,6 +953,7 @@ fn make_create_req(name: &str, schedule: CronScheduleDto) -> CreateCronJobReques conversation_title: Some("Test Conv".into()), created_by: "user".into(), execution_mode: None, + queue_enabled: false, agent_config: Some(aionui_api_types::CronAgentConfigWriteDto { name: "Default Assistant".into(), cli_path: None, @@ -1374,6 +1375,7 @@ async fn list_jobs_allows_legacy_custom_agent_id_without_assistant_id() { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, }) .await .unwrap(); @@ -1449,6 +1451,7 @@ async fn cj8_update_job() { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let updated = svc.update_job(&created.id, req).await.unwrap(); @@ -1489,6 +1492,7 @@ async fn update_existing_conversation_job_rejects_agent_config_changes() { }), conversation_title: None, max_retries: None, + queue_enabled: None, }; let err = svc.update_job(&created.id, req).await.unwrap_err(); @@ -1527,6 +1531,7 @@ async fn update_existing_conversation_job_rejects_agent_config_even_when_switchi }), conversation_title: None, max_retries: None, + queue_enabled: None, }; let err = svc.update_job(&created.id, req).await.unwrap_err(); @@ -1563,6 +1568,7 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }, ) .await @@ -1628,6 +1634,7 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }, ) .await @@ -1671,6 +1678,7 @@ async fn update_existing_job_to_new_conversation_preserves_custom_workspace() { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }, ) .await @@ -1705,6 +1713,7 @@ async fn update_team_conversation_job_rejects_execution_mode_change() { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let err = svc.update_job(&created.id, req).await.unwrap_err(); @@ -1738,6 +1747,7 @@ async fn update_job_strips_legacy_agent_ids_when_assistant_id_present() { }), conversation_title: None, max_retries: None, + queue_enabled: None, }; let updated = svc.update_job(&created.id, req).await.unwrap(); @@ -1774,6 +1784,7 @@ async fn update_job_rejects_when_assistant_id_cannot_resolve() { }), conversation_title: None, max_retries: None, + queue_enabled: None, }; let err = svc @@ -1808,6 +1819,7 @@ async fn cj9_update_schedule_type() { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let updated = svc.update_job(&created.id, req).await.unwrap(); @@ -1833,6 +1845,7 @@ async fn cj10_update_nonexistent() { agent_config: None, conversation_title: None, max_retries: None, + queue_enabled: None, }; let err = svc.update_job("cron_nonexistent", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::JobNotFound(_))); @@ -2762,11 +2775,45 @@ async fn update_max_retries() { agent_config: None, conversation_title: None, max_retries: Some(5), + queue_enabled: None, }; let updated = svc.update_job(&job.id, req).await.unwrap(); assert_eq!(updated.max_retries, 5); } +#[tokio::test] +async fn create_and_update_queue_enabled() { + let (svc, _, _) = setup().await; + let mut create = make_create_req("Queued", every_60s()); + create.queue_enabled = true; + let job = svc.add_job(create).await.unwrap(); + assert!(job.queue_enabled); + assert!(CronService::to_response(&job).state.queue_enabled); + + let updated = svc + .update_job( + &job.id, + UpdateCronJobRequest { + queue_enabled: Some(false), + ..UpdateCronJobRequest { + name: None, + description: None, + enabled: None, + schedule: None, + message: None, + execution_mode: None, + agent_config: None, + conversation_title: None, + max_retries: None, + queue_enabled: None, + } + }, + ) + .await + .unwrap(); + assert!(!updated.queue_enabled); +} + // ── SC-1: At type — future timestamp, nextRunAtMs == atMs ──────── #[tokio::test] diff --git a/crates/aionui-db/migrations/021_cron_queue_enabled.sql b/crates/aionui-db/migrations/021_cron_queue_enabled.sql new file mode 100644 index 000000000..83d99edb3 --- /dev/null +++ b/crates/aionui-db/migrations/021_cron_queue_enabled.sql @@ -0,0 +1 @@ +ALTER TABLE cron_jobs ADD COLUMN queue_enabled INTEGER NOT NULL DEFAULT 0; diff --git a/crates/aionui-db/migrations/022_cron_job_runs.sql b/crates/aionui-db/migrations/022_cron_job_runs.sql new file mode 100644 index 000000000..35618c1e6 --- /dev/null +++ b/crates/aionui-db/migrations/022_cron_job_runs.sql @@ -0,0 +1,19 @@ +CREATE TABLE cron_job_runs ( + job_id TEXT NOT NULL, + scheduled_at INTEGER NOT NULL, + status TEXT NOT NULL CHECK (status IN ('running', 'retrying', 'ok', 'error', 'skipped')), + owner_id TEXT, + lease_until INTEGER, + conversation_id TEXT, + error TEXT, + started_at INTEGER, + finished_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (job_id, scheduled_at), + FOREIGN KEY (job_id) REFERENCES cron_jobs(id) ON DELETE CASCADE +); + +CREATE INDEX idx_cron_job_runs_active + ON cron_job_runs(job_id, lease_until) + WHERE status IN ('running', 'retrying'); diff --git a/crates/aionui-db/src/lib.rs b/crates/aionui-db/src/lib.rs index ee3a776b0..147f364b2 100644 --- a/crates/aionui-db/src/lib.rs +++ b/crates/aionui-db/src/lib.rs @@ -29,7 +29,9 @@ pub use repository::conversation::{ ConversationFilters, ConversationRowUpdate, MessagePageCursor, MessagePageDirection, MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, }; -pub use repository::cron::UpdateCronJobParams; +pub use repository::cron::{ + ClaimCronRunParams, CronRunClaimResult, FinishCronRunParams, RecoverableCronRun, UpdateCronJobParams, +}; pub use repository::mcp_server::{CreateMcpServerParams, UpdateMcpServerParams}; pub use repository::oauth_token::UpsertOAuthTokenParams; pub use repository::provider::{CreateProviderParams, UpdateProviderParams}; diff --git a/crates/aionui-db/src/models/cron_job.rs b/crates/aionui-db/src/models/cron_job.rs index c091f3a60..9f9151785 100644 --- a/crates/aionui-db/src/models/cron_job.rs +++ b/crates/aionui-db/src/models/cron_job.rs @@ -28,6 +28,7 @@ pub struct CronJobRow { pub run_count: i64, pub retry_count: i64, pub max_retries: i64, + pub queue_enabled: bool, } #[cfg(test)] @@ -61,6 +62,7 @@ mod tests { run_count: 5, retry_count: 0, max_retries: 3, + queue_enabled: true, }; let json = serde_json::to_string(&row).expect("serialize"); let restored: CronJobRow = serde_json::from_str(&json).expect("deserialize"); @@ -98,6 +100,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, }; assert!(row.schedule_tz.is_none()); assert!(row.agent_config.is_none()); diff --git a/crates/aionui-db/src/repository/cron.rs b/crates/aionui-db/src/repository/cron.rs index 90e8c52f1..d2c27359c 100644 --- a/crates/aionui-db/src/repository/cron.rs +++ b/crates/aionui-db/src/repository/cron.rs @@ -3,6 +3,40 @@ use aionui_common::TimestampMs; use crate::error::DbError; use crate::models::CronJobRow; +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CronRunClaimResult { + Claimed, + Duplicate, + QueueBusy, +} + +#[derive(Debug, Clone)] +pub struct ClaimCronRunParams<'a> { + pub job_id: &'a str, + pub scheduled_at: TimestampMs, + pub owner_id: &'a str, + pub now: TimestampMs, + pub lease_until: TimestampMs, + pub queue_enabled: bool, +} + +#[derive(Debug, Clone)] +pub struct FinishCronRunParams<'a> { + pub job_id: &'a str, + pub scheduled_at: TimestampMs, + pub owner_id: &'a str, + pub status: &'a str, + pub conversation_id: Option<&'a str>, + pub error: Option<&'a str>, + pub finished_at: TimestampMs, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RecoverableCronRun { + pub scheduled_at: TimestampMs, + pub wake_at: TimestampMs, +} + /// Parameters for updating a cron job. /// /// All fields are optional; `None` means "keep the current value". @@ -27,6 +61,7 @@ pub struct UpdateCronJobParams { pub last_error: Option>, pub run_count: Option, pub retry_count: Option, + pub queue_enabled: Option, } /// Data access abstraction for the `cron_jobs` table. @@ -57,4 +92,36 @@ pub trait ICronRepository: Send + Sync { /// Deletes all cron jobs associated with a conversation. /// Returns the number of deleted rows. async fn delete_by_conversation(&self, conversation_id: &str) -> Result; + + /// Atomically claims one scheduled occurrence across all backend processes. + async fn claim_run(&self, params: &ClaimCronRunParams<'_>) -> Result; + + /// Extends an active run lease owned by this backend instance. + async fn renew_run_lease( + &self, + job_id: &str, + scheduled_at: TimestampMs, + owner_id: &str, + lease_until: TimestampMs, + updated_at: TimestampMs, + ) -> Result; + + /// Releases a claimed occurrence until its scheduled retry time. + async fn defer_run( + &self, + job_id: &str, + scheduled_at: TimestampMs, + owner_id: &str, + retry_at: TimestampMs, + updated_at: TimestampMs, + ) -> Result; + + /// Completes a claimed occurrence and releases its lease. + async fn finish_run(&self, params: &FinishCronRunParams<'_>) -> Result; + + /// Deletes terminal run records older than the retention cutoff. + async fn cleanup_runs_before(&self, cutoff: TimestampMs) -> Result; + + /// Returns the oldest unfinished occurrence that should be resumed for a job. + async fn get_recoverable_run(&self, job_id: &str, now: TimestampMs) -> Result, DbError>; } diff --git a/crates/aionui-db/src/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index d11df02df..a5207aacb 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -1,9 +1,12 @@ -use aionui_common::now_ms; +use aionui_common::{TimestampMs, now_ms}; use sqlx::SqlitePool; use crate::error::DbError; use crate::models::CronJobRow; -use crate::repository::cron::{ICronRepository, UpdateCronJobParams}; +use crate::repository::cron::{ + ClaimCronRunParams, CronRunClaimResult, FinishCronRunParams, ICronRepository, RecoverableCronRun, + UpdateCronJobParams, +}; #[derive(Clone, Debug)] pub struct SqliteCronRepository { @@ -25,9 +28,9 @@ impl ICronRepository for SqliteCronRepository { schedule_description, payload_message, execution_mode, agent_config, \ conversation_id, conversation_title, created_by, \ skill_content, description, created_at, updated_at, next_run_at, last_run_at, \ - last_status, last_error, run_count, retry_count, max_retries\ + last_status, last_error, run_count, retry_count, max_retries, queue_enabled\ ) VALUES (\ - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?\ + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?\ )", ) .bind(&row.id) @@ -54,6 +57,7 @@ impl ICronRepository for SqliteCronRepository { .bind(row.run_count) .bind(row.retry_count) .bind(row.max_retries) + .bind(row.queue_enabled) .execute(&self.pool) .await?; Ok(()) @@ -103,6 +107,10 @@ impl ICronRepository for SqliteCronRepository { set_parts.push("enabled = ?".to_string()); binds.push(BindValue::Bool(v)); } + if let Some(v) = params.queue_enabled { + set_parts.push("queue_enabled = ?".to_string()); + binds.push(BindValue::Bool(v)); + } push_str!(name); push_str!(schedule_kind); @@ -195,6 +203,188 @@ impl ICronRepository for SqliteCronRepository { .await?; Ok(result.rows_affected()) } + + async fn claim_run(&self, params: &ClaimCronRunParams<'_>) -> Result { + let mut connection = self.pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE").execute(&mut *connection).await?; + + let result = async { + let existing = sqlx::query_as::<_, (String, Option, Option)>( + "SELECT status, owner_id, lease_until FROM cron_job_runs WHERE job_id = ? AND scheduled_at = ?", + ) + .bind(params.job_id) + .bind(params.scheduled_at) + .fetch_optional(&mut *connection) + .await?; + + if let Some((status, _owner_id, lease_until)) = existing { + if matches!(status.as_str(), "running" | "retrying") + && lease_until.is_some_and(|lease| lease <= params.now) + { + let updated = sqlx::query( + "UPDATE cron_job_runs SET status = 'running', owner_id = ?, lease_until = ?, \ + started_at = COALESCE(started_at, ?), finished_at = NULL, error = NULL, updated_at = ? \ + WHERE job_id = ? AND scheduled_at = ? AND status IN ('running', 'retrying') AND lease_until <= ?", + ) + .bind(params.owner_id) + .bind(params.lease_until) + .bind(params.now) + .bind(params.now) + .bind(params.job_id) + .bind(params.scheduled_at) + .bind(params.now) + .execute(&mut *connection) + .await?; + return Ok(if updated.rows_affected() == 1 { + CronRunClaimResult::Claimed + } else { + CronRunClaimResult::Duplicate + }); + } + return Ok(CronRunClaimResult::Duplicate); + } + + if params.queue_enabled { + let has_active_run: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM cron_job_runs \ + WHERE job_id = ? AND status IN ('running', 'retrying') AND lease_until > ?)", + ) + .bind(params.job_id) + .bind(params.now) + .fetch_one(&mut *connection) + .await?; + + if has_active_run { + sqlx::query( + "INSERT INTO cron_job_runs (job_id, scheduled_at, status, created_at, updated_at, finished_at) \ + VALUES (?, ?, 'skipped', ?, ?, ?)", + ) + .bind(params.job_id) + .bind(params.scheduled_at) + .bind(params.now) + .bind(params.now) + .bind(params.now) + .execute(&mut *connection) + .await?; + return Ok(CronRunClaimResult::QueueBusy); + } + } + + sqlx::query( + "INSERT INTO cron_job_runs (job_id, scheduled_at, status, owner_id, lease_until, started_at, created_at, updated_at) \ + VALUES (?, ?, 'running', ?, ?, ?, ?, ?)", + ) + .bind(params.job_id) + .bind(params.scheduled_at) + .bind(params.owner_id) + .bind(params.lease_until) + .bind(params.now) + .bind(params.now) + .bind(params.now) + .execute(&mut *connection) + .await?; + + Ok(CronRunClaimResult::Claimed) + } + .await; + + match result { + Ok(result) => { + sqlx::query("COMMIT").execute(&mut *connection).await?; + Ok(result) + } + Err(error) => { + let _ = sqlx::query("ROLLBACK").execute(&mut *connection).await; + Err(error) + } + } + } + + async fn renew_run_lease( + &self, + job_id: &str, + scheduled_at: TimestampMs, + owner_id: &str, + lease_until: TimestampMs, + updated_at: TimestampMs, + ) -> Result { + let result = sqlx::query( + "UPDATE cron_job_runs SET lease_until = ?, updated_at = ? \ + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + ) + .bind(lease_until) + .bind(updated_at) + .bind(job_id) + .bind(scheduled_at) + .bind(owner_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() == 1) + } + + async fn defer_run( + &self, + job_id: &str, + scheduled_at: TimestampMs, + owner_id: &str, + retry_at: TimestampMs, + updated_at: TimestampMs, + ) -> Result { + let result = sqlx::query( + "UPDATE cron_job_runs SET status = 'retrying', owner_id = NULL, lease_until = ?, updated_at = ? \ + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + ) + .bind(retry_at) + .bind(updated_at) + .bind(job_id) + .bind(scheduled_at) + .bind(owner_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() == 1) + } + + async fn finish_run(&self, params: &FinishCronRunParams<'_>) -> Result { + let result = sqlx::query( + "UPDATE cron_job_runs SET status = ?, conversation_id = ?, error = ?, lease_until = NULL, \ + finished_at = ?, updated_at = ? \ + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + ) + .bind(params.status) + .bind(params.conversation_id) + .bind(params.error) + .bind(params.finished_at) + .bind(params.finished_at) + .bind(params.job_id) + .bind(params.scheduled_at) + .bind(params.owner_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() == 1) + } + + async fn cleanup_runs_before(&self, cutoff: TimestampMs) -> Result { + let result = + sqlx::query("DELETE FROM cron_job_runs WHERE status IN ('ok', 'error', 'skipped') AND finished_at < ?") + .bind(cutoff) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + async fn get_recoverable_run(&self, job_id: &str, now: TimestampMs) -> Result, DbError> { + let row = sqlx::query_as::<_, (TimestampMs, TimestampMs)>( + "SELECT scheduled_at, MAX(lease_until, ?) AS wake_at FROM cron_job_runs \ + WHERE job_id = ? AND status IN ('running', 'retrying') AND lease_until IS NOT NULL \ + ORDER BY lease_until ASC LIMIT 1", + ) + .bind(now) + .bind(job_id) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|(scheduled_at, wake_at)| RecoverableCronRun { scheduled_at, wake_at })) + } } // ── Dynamic bind helpers ──────────────────────────────────────────── @@ -276,6 +466,7 @@ mod tests { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } } @@ -369,6 +560,275 @@ mod tests { assert!(updated.updated_at >= updated.created_at); } + #[tokio::test] + async fn update_queue_enabled() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_queue")).await.unwrap(); + + repo.update( + "cron_queue", + &UpdateCronJobParams { + queue_enabled: Some(true), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert!(repo.get_by_id("cron_queue").await.unwrap().unwrap().queue_enabled); + } + + #[tokio::test] + async fn claim_run_deduplicates_same_occurrence_across_repository_instances() { + let (repo, db) = setup().await; + repo.insert(&make_row("cron_claim")).await.unwrap(); + let other_repo = SqliteCronRepository::new(db.pool().clone()); + + let first = ClaimCronRunParams { + job_id: "cron_claim", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: false, + }; + let second = ClaimCronRunParams { + owner_id: "owner-b", + ..first.clone() + }; + let (first, second) = tokio::join!(repo.claim_run(&first), other_repo.claim_run(&second)); + let results = [first.unwrap(), second.unwrap()]; + + assert_eq!( + results + .iter() + .filter(|result| **result == CronRunClaimResult::Claimed) + .count(), + 1 + ); + assert_eq!( + results + .iter() + .filter(|result| **result == CronRunClaimResult::Duplicate) + .count(), + 1 + ); + } + + #[tokio::test] + async fn claim_run_queue_busy_records_skipped_occurrence() { + let (repo, db) = setup().await; + repo.insert(&make_row("cron_queue_claim")).await.unwrap(); + + assert_eq!( + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_queue_claim", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: true, + }) + .await + .unwrap(), + CronRunClaimResult::Claimed + ); + assert_eq!( + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_queue_claim", + scheduled_at: 20_000, + owner_id: "owner-b", + now: 20_000, + lease_until: 80_000, + queue_enabled: true, + }) + .await + .unwrap(), + CronRunClaimResult::QueueBusy + ); + + let status: String = sqlx::query_scalar( + "SELECT status FROM cron_job_runs WHERE job_id = 'cron_queue_claim' AND scheduled_at = 20000", + ) + .fetch_one(db.pool()) + .await + .unwrap(); + assert_eq!(status, "skipped"); + } + + #[tokio::test] + async fn expired_run_lease_can_be_reclaimed() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_expired")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_expired", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 20_000, + queue_enabled: false, + }) + .await + .unwrap(); + + assert_eq!( + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_expired", + scheduled_at: 10_000, + owner_id: "owner-b", + now: 20_001, + lease_until: 80_001, + queue_enabled: false, + }) + .await + .unwrap(), + CronRunClaimResult::Claimed + ); + } + + #[tokio::test] + async fn deferred_run_reuses_original_occurrence_key() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_retry_claim")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_retry_claim", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: false, + }) + .await + .unwrap(); + assert!( + repo.defer_run("cron_retry_claim", 10_000, "owner-a", 20_000, 11_000) + .await + .unwrap() + ); + + assert_eq!( + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_retry_claim", + scheduled_at: 10_000, + owner_id: "owner-b", + now: 20_000, + lease_until: 80_000, + queue_enabled: false, + }) + .await + .unwrap(), + CronRunClaimResult::Claimed + ); + } + + #[tokio::test] + async fn queue_enabled_treats_deferred_run_as_active() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_retry_queue")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_retry_queue", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: true, + }) + .await + .unwrap(); + repo.defer_run("cron_retry_queue", 10_000, "owner-a", 40_000, 11_000) + .await + .unwrap(); + + assert_eq!( + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_retry_queue", + scheduled_at: 20_000, + owner_id: "owner-b", + now: 20_000, + lease_until: 80_000, + queue_enabled: true, + }) + .await + .unwrap(), + CronRunClaimResult::QueueBusy + ); + } + + #[tokio::test] + async fn recoverable_running_run_waits_for_unexpired_lease() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_recover_running")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_recover_running", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: false, + }) + .await + .unwrap(); + + assert_eq!( + repo.get_recoverable_run("cron_recover_running", 20_000).await.unwrap(), + Some(RecoverableCronRun { + scheduled_at: 10_000, + wake_at: 70_000, + }) + ); + } + + #[tokio::test] + async fn recoverable_running_run_wakes_immediately_after_expiry() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_recover_expired")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_recover_expired", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 20_000, + queue_enabled: false, + }) + .await + .unwrap(); + + assert_eq!( + repo.get_recoverable_run("cron_recover_expired", 30_000).await.unwrap(), + Some(RecoverableCronRun { + scheduled_at: 10_000, + wake_at: 30_000, + }) + ); + } + + #[tokio::test] + async fn recoverable_retry_preserves_original_occurrence() { + let (repo, _db) = setup().await; + repo.insert(&make_row("cron_recover_retry")).await.unwrap(); + repo.claim_run(&ClaimCronRunParams { + job_id: "cron_recover_retry", + scheduled_at: 10_000, + owner_id: "owner-a", + now: 10_000, + lease_until: 70_000, + queue_enabled: false, + }) + .await + .unwrap(); + repo.defer_run("cron_recover_retry", 10_000, "owner-a", 40_000, 11_000) + .await + .unwrap(); + + assert_eq!( + repo.get_recoverable_run("cron_recover_retry", 20_000).await.unwrap(), + Some(RecoverableCronRun { + scheduled_at: 10_000, + wake_at: 40_000, + }) + ); + } + #[tokio::test] async fn update_optional_nullable_fields() { let (repo, _db) = setup().await; diff --git a/crates/aionui-db/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index 316e1c793..4ea61f749 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -64,6 +64,7 @@ fn make_job(id: &str) -> CronJobRow { run_count: 0, retry_count: 0, max_retries: 3, + queue_enabled: false, } }