1use crate::runtime::InFlightMap;
2use awa_model::cron::{atomic_enqueue, list_cron_jobs, upsert_cron_job, CronJobRow};
3use awa_model::{JobRow, PeriodicJob};
4use chrono::Utc;
5use croner::Cron;
6use sqlx::pool::PoolConnection;
7use sqlx::{PgPool, Postgres};
8use std::collections::{HashMap, HashSet};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::Arc;
11use std::time::Duration;
12use tokio_util::sync::CancellationToken;
13use tracing::{debug, error, info, warn};
14
15#[derive(Debug, Clone)]
17pub struct RetentionPolicy {
18 pub completed: Duration,
20 pub failed: Duration,
22}
23
24impl Default for RetentionPolicy {
25 fn default() -> Self {
26 Self {
27 completed: Duration::from_secs(86400), failed: Duration::from_secs(259200), }
30 }
31}
32
33pub struct MaintenanceService {
38 pool: PgPool,
39 metrics: crate::metrics::AwaMetrics,
40 cancel: CancellationToken,
41 leader: Arc<AtomicBool>,
42 periodic_jobs: Arc<Vec<PeriodicJob>>,
43 in_flight: InFlightMap,
46 heartbeat_rescue_interval: Duration,
47 deadline_rescue_interval: Duration,
48 callback_rescue_interval: Duration,
49 promote_interval: Duration,
50 cleanup_interval: Duration,
51 cron_sync_interval: Duration,
52 cron_eval_interval: Duration,
53 leader_check_interval: Duration,
54 leader_election_interval: Duration,
55 heartbeat_staleness: Duration,
56 completed_retention: Duration,
57 failed_retention: Duration,
58 cleanup_batch_size: i64,
59 queue_retention_overrides: HashMap<String, RetentionPolicy>,
60}
61
62const PROMOTE_BATCH_SIZE: i64 = 4_096;
63const PROMOTE_MAX_BATCHES_PER_TICK: usize = 32;
64
65impl MaintenanceService {
66 pub(crate) fn new(
67 pool: PgPool,
68 metrics: crate::metrics::AwaMetrics,
69 leader: Arc<AtomicBool>,
70 cancel: CancellationToken,
71 periodic_jobs: Arc<Vec<PeriodicJob>>,
72 in_flight: InFlightMap,
73 ) -> Self {
74 Self {
75 pool,
76 metrics,
77 cancel,
78 leader,
79 periodic_jobs,
80 in_flight,
81 heartbeat_rescue_interval: Duration::from_secs(30),
82 deadline_rescue_interval: Duration::from_secs(30),
83 callback_rescue_interval: Duration::from_secs(30),
84 promote_interval: Duration::from_millis(250),
85 cleanup_interval: Duration::from_secs(60),
86 cron_sync_interval: Duration::from_secs(60),
87 cron_eval_interval: Duration::from_secs(1),
88 leader_check_interval: Duration::from_secs(30),
89 leader_election_interval: Duration::from_secs(10),
90 heartbeat_staleness: Duration::from_secs(90),
91 completed_retention: Duration::from_secs(86400), failed_retention: Duration::from_secs(259200), cleanup_batch_size: 1000,
94 queue_retention_overrides: HashMap::new(),
95 }
96 }
97
98 pub fn leader_election_interval(mut self, interval: Duration) -> Self {
103 self.leader_election_interval = interval;
104 self
105 }
106
107 pub fn promote_interval(mut self, interval: Duration) -> Self {
109 self.promote_interval = interval;
110 self
111 }
112
113 pub fn cleanup_interval(mut self, interval: Duration) -> Self {
115 self.cleanup_interval = interval;
116 self
117 }
118
119 pub fn completed_retention(mut self, retention: Duration) -> Self {
121 self.completed_retention = retention;
122 self
123 }
124
125 pub fn failed_retention(mut self, retention: Duration) -> Self {
127 self.failed_retention = retention;
128 self
129 }
130
131 pub fn cleanup_batch_size(mut self, batch_size: i64) -> Self {
133 self.cleanup_batch_size = batch_size;
134 self
135 }
136
137 pub fn queue_retention_overrides(
139 mut self,
140 overrides: HashMap<String, RetentionPolicy>,
141 ) -> Self {
142 self.queue_retention_overrides = overrides;
143 self
144 }
145
146 pub async fn run(&self) {
148 info!("Maintenance service starting");
149 self.leader.store(false, Ordering::SeqCst);
150
151 loop {
152 let mut leader_conn = match self.try_become_leader().await {
155 Ok(Some(conn)) => conn,
156 Ok(None) => {
157 tokio::select! {
159 _ = self.cancel.cancelled() => {
160 debug!("Maintenance service shutting down (not leader)");
161 self.leader.store(false, Ordering::SeqCst);
162 return;
163 }
164 _ = tokio::time::sleep(self.leader_election_interval) => continue,
165 }
166 }
167 Err(err) => {
168 warn!(error = %err, "Failed to check leader status");
169 tokio::select! {
170 _ = self.cancel.cancelled() => {
171 debug!("Maintenance service shutting down (leader check failed)");
172 self.leader.store(false, Ordering::SeqCst);
173 return;
174 }
175 _ = tokio::time::sleep(self.leader_election_interval) => continue,
176 }
177 }
178 };
179
180 debug!("Elected as maintenance leader");
181 self.leader.store(true, Ordering::SeqCst);
182
183 let mut heartbeat_rescue_timer = tokio::time::interval(self.heartbeat_rescue_interval);
185 let mut deadline_rescue_timer = tokio::time::interval(self.deadline_rescue_interval);
186 let mut callback_rescue_timer = tokio::time::interval(self.callback_rescue_interval);
187 let mut promote_timer = tokio::time::interval(self.promote_interval);
188 let mut cleanup_timer = tokio::time::interval(self.cleanup_interval);
189 let mut cron_sync_timer = tokio::time::interval(self.cron_sync_interval);
190 let mut cron_eval_timer = tokio::time::interval(self.cron_eval_interval);
191 let mut leader_check_timer = tokio::time::interval(self.leader_check_interval);
192
193 heartbeat_rescue_timer.tick().await;
195 deadline_rescue_timer.tick().await;
196 callback_rescue_timer.tick().await;
197 promote_timer.tick().await;
198 cleanup_timer.tick().await;
199 cron_sync_timer.tick().await;
200 cron_eval_timer.tick().await;
201 leader_check_timer.tick().await;
202
203 self.sync_periodic_jobs_to_db().await;
205
206 loop {
207 tokio::select! {
208 _ = self.cancel.cancelled() => {
209 debug!("Maintenance service shutting down");
210 self.leader.store(false, Ordering::SeqCst);
211 let _ = Self::release_leader(&mut leader_conn).await;
214 return;
215 }
216 _ = heartbeat_rescue_timer.tick() => {
217 self.rescue_stale_heartbeats().await;
218 }
219 _ = deadline_rescue_timer.tick() => {
220 self.rescue_expired_deadlines().await;
221 }
222 _ = callback_rescue_timer.tick() => {
223 self.rescue_expired_callbacks().await;
224 }
225 _ = promote_timer.tick() => {
226 self.promote_scheduled().await;
227 }
228 _ = cleanup_timer.tick() => {
229 self.cleanup_completed().await;
230 }
231 _ = cron_sync_timer.tick() => {
232 self.sync_periodic_jobs_to_db().await;
233 }
234 _ = cron_eval_timer.tick() => {
235 self.evaluate_cron_schedules().await;
236 }
237 _ = leader_check_timer.tick() => {
238 if sqlx::query("SELECT 1").execute(&mut *leader_conn).await.is_err() {
242 warn!("Leader connection lost, re-entering election loop");
243 self.leader.store(false, Ordering::SeqCst);
244 break;
245 }
246 }
247 }
248 }
249 }
250 }
251
252 const LOCK_KEY: i64 = 0x_4157_415f_4d41_494e; async fn try_become_leader(&self) -> Result<Option<PoolConnection<Postgres>>, sqlx::Error> {
261 let mut conn = self.pool.acquire().await?;
262 let result: (bool,) = sqlx::query_as("SELECT pg_try_advisory_lock($1)")
263 .bind(Self::LOCK_KEY)
264 .fetch_one(&mut *conn)
265 .await?;
266 if result.0 {
267 Ok(Some(conn))
268 } else {
269 Ok(None)
270 }
271 }
272
273 async fn release_leader(conn: &mut PoolConnection<Postgres>) -> Result<(), sqlx::Error> {
278 sqlx::query("SELECT pg_advisory_unlock($1)")
279 .bind(Self::LOCK_KEY)
280 .execute(&mut **conn)
281 .await?;
282 Ok(())
283 }
284
285 #[tracing::instrument(skip(self), name = "maintenance.cron_sync")]
289 async fn sync_periodic_jobs_to_db(&self) {
290 if self.periodic_jobs.is_empty() {
291 return;
292 }
293
294 for job in self.periodic_jobs.iter() {
295 if let Err(err) = upsert_cron_job(&self.pool, job).await {
296 error!(name = %job.name, error = %err, "Failed to sync periodic job");
297 }
298 }
299
300 debug!(
301 count = self.periodic_jobs.len(),
302 "Synced periodic jobs to database"
303 );
304 }
305
306 #[tracing::instrument(skip(self), name = "maintenance.cron_eval")]
312 async fn evaluate_cron_schedules(&self) {
313 let cron_rows = match list_cron_jobs(&self.pool).await {
314 Ok(rows) => rows,
315 Err(err) => {
316 error!(error = %err, "Failed to load cron jobs for evaluation");
317 return;
318 }
319 };
320
321 if cron_rows.is_empty() {
322 return;
323 }
324
325 let now = Utc::now();
326
327 for row in &cron_rows {
328 let fire_time = match compute_fire_time(row, now) {
329 Some(time) => time,
330 None => continue,
331 };
332
333 match atomic_enqueue(&self.pool, &row.name, fire_time, row.last_enqueued_at).await {
334 Ok(Some(job)) => {
335 info!(
336 cron_name = %row.name,
337 job_id = job.id,
338 fire_time = %fire_time,
339 "Enqueued periodic job"
340 );
341 }
342 Ok(None) => {
343 debug!(cron_name = %row.name, "Cron fire already claimed");
345 }
346 Err(err) => {
347 error!(
348 cron_name = %row.name,
349 error = %err,
350 "Failed to enqueue periodic job"
351 );
352 }
353 }
354 }
355 }
356
357 #[tracing::instrument(skip(self), name = "maintenance.rescue_stale")]
359 async fn rescue_stale_heartbeats(&self) {
360 let staleness_str = format!("{} seconds", self.heartbeat_staleness.as_secs());
361 match sqlx::query_as::<_, JobRow>(
362 r#"
363 UPDATE awa.jobs
364 SET state = 'retryable',
365 finalized_at = now(),
366 heartbeat_at = NULL,
367 deadline_at = NULL,
368 callback_id = NULL,
369 callback_timeout_at = NULL,
370 callback_filter = NULL,
371 callback_on_complete = NULL,
372 callback_on_fail = NULL,
373 callback_transform = NULL,
374 errors = errors || jsonb_build_object(
375 'error', 'heartbeat stale: worker presumed dead',
376 'attempt', attempt,
377 'at', now()
378 )::jsonb
379 WHERE id IN (
380 SELECT id FROM awa.jobs
381 WHERE state = 'running'
382 AND heartbeat_at < now() - $1::interval
383 LIMIT 500
384 FOR UPDATE SKIP LOCKED
385 )
386 RETURNING *
387 "#,
388 )
389 .bind(&staleness_str)
390 .fetch_all(&self.pool)
391 .await
392 {
393 Ok(rescued) if !rescued.is_empty() => {
394 warn!(count = rescued.len(), "Rescued stale heartbeat jobs");
395 self.signal_cancellation(&rescued).await;
397 }
398 Err(err) => {
399 error!(error = %err, "Failed to rescue stale heartbeat jobs");
400 }
401 _ => {}
402 }
403 }
404
405 #[tracing::instrument(skip(self), name = "maintenance.rescue_deadline")]
407 async fn rescue_expired_deadlines(&self) {
408 match sqlx::query_as::<_, JobRow>(
409 r#"
410 UPDATE awa.jobs
411 SET state = 'retryable',
412 finalized_at = now(),
413 heartbeat_at = NULL,
414 deadline_at = NULL,
415 callback_id = NULL,
416 callback_timeout_at = NULL,
417 callback_filter = NULL,
418 callback_on_complete = NULL,
419 callback_on_fail = NULL,
420 callback_transform = NULL,
421 errors = errors || jsonb_build_object(
422 'error', 'hard deadline exceeded',
423 'attempt', attempt,
424 'at', now()
425 )::jsonb
426 WHERE id IN (
427 SELECT id FROM awa.jobs
428 WHERE state = 'running'
429 AND deadline_at IS NOT NULL
430 AND deadline_at < now()
431 LIMIT 500
432 FOR UPDATE SKIP LOCKED
433 )
434 RETURNING *
435 "#,
436 )
437 .fetch_all(&self.pool)
438 .await
439 {
440 Ok(rescued) if !rescued.is_empty() => {
441 warn!(count = rescued.len(), "Rescued deadline-expired jobs");
442 self.signal_cancellation(&rescued).await;
444 }
445 Err(err) => {
446 error!(error = %err, "Failed to rescue deadline-expired jobs");
447 }
448 _ => {}
449 }
450 }
451
452 #[tracing::instrument(skip(self), name = "maintenance.rescue_callback_timeout")]
454 async fn rescue_expired_callbacks(&self) {
455 match sqlx::query_as::<_, JobRow>(
456 r#"
457 UPDATE awa.jobs
458 SET state = CASE WHEN attempt >= max_attempts THEN 'failed'::awa.job_state ELSE 'retryable'::awa.job_state END,
459 finalized_at = now(),
460 callback_id = NULL,
461 callback_timeout_at = NULL,
462 callback_filter = NULL,
463 callback_on_complete = NULL,
464 callback_on_fail = NULL,
465 callback_transform = NULL,
466 run_at = CASE WHEN attempt >= max_attempts THEN run_at
467 ELSE now() + awa.backoff_duration(attempt, max_attempts) END,
468 errors = errors || jsonb_build_object(
469 'error', 'callback timed out',
470 'attempt', attempt,
471 'at', now()
472 )::jsonb
473 WHERE id IN (
474 SELECT id FROM awa.jobs
475 WHERE state = 'waiting_external'
476 AND callback_timeout_at IS NOT NULL
477 AND callback_timeout_at < now()
478 LIMIT 500
479 FOR UPDATE SKIP LOCKED
480 )
481 RETURNING *
482 "#,
483 )
484 .fetch_all(&self.pool)
485 .await
486 {
487 Ok(rescued) if !rescued.is_empty() => {
488 warn!(count = rescued.len(), "Rescued callback-timed-out jobs");
489 }
490 Err(err) => {
491 error!(error = %err, "Failed to rescue callback-timed-out jobs");
492 }
493 _ => {}
494 }
495 }
496
497 async fn signal_cancellation(&self, rescued_jobs: &[JobRow]) {
499 for job in rescued_jobs {
500 if let Some(flag) = self.in_flight.get_cancel((job.id, job.run_lease)) {
501 flag.store(true, Ordering::SeqCst);
502 debug!(job_id = job.id, "Signalled cancellation for rescued job");
503 }
504 }
505 }
506
507 #[tracing::instrument(skip(self), name = "maintenance.promote")]
509 async fn promote_scheduled(&self) {
510 if let Err(err) = self.promote_due_state("scheduled", "scheduled jobs").await {
511 error!(error = %err, "Failed to promote scheduled jobs");
512 }
513 if let Err(err) = self
514 .promote_due_state("retryable", "retryable jobs (backoff elapsed)")
515 .await
516 {
517 error!(error = %err, "Failed to promote retryable jobs");
518 }
519 }
520
521 async fn promote_due_state(
522 &self,
523 state: &'static str,
524 label: &'static str,
525 ) -> Result<(), sqlx::Error> {
526 let mut promoted_total = 0usize;
527 let mut notified_queues = HashSet::new();
528
529 for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK {
530 if self.cancel.is_cancelled() {
531 break;
532 }
533
534 let (promoted, queues) = self.promote_due_batch(state).await?;
535 if promoted == 0 {
536 break;
537 }
538
539 promoted_total += promoted;
540 notified_queues.extend(queues);
541
542 if promoted < PROMOTE_BATCH_SIZE as usize {
543 break;
544 }
545 }
546
547 if promoted_total > 0 {
548 debug!(
549 count = promoted_total,
550 queues = notified_queues.len(),
551 state,
552 "Promoted {label}"
553 );
554 }
555
556 Ok(())
557 }
558
559 async fn promote_due_batch(
560 &self,
561 state: &'static str,
562 ) -> Result<(usize, HashSet<String>), sqlx::Error> {
563 let mut tx = self.pool.begin().await?;
564 let promote_start = std::time::Instant::now();
565 let promoted_rows: Vec<(String,)> = sqlx::query_as(
566 r#"
567 WITH due AS (
568 DELETE FROM awa.scheduled_jobs
569 WHERE id IN (
570 SELECT id
571 FROM awa.scheduled_jobs
572 WHERE state = $1::awa.job_state
573 AND run_at <= now()
574 ORDER BY run_at ASC, id ASC
575 LIMIT $2
576 FOR UPDATE SKIP LOCKED
577 )
578 RETURNING *
579 ),
580 promoted AS (
581 INSERT INTO awa.jobs_hot (
582 id, kind, queue, args, state, priority, attempt, max_attempts,
583 run_at, heartbeat_at, deadline_at, attempted_at, finalized_at,
584 created_at, errors, metadata, tags, unique_key, unique_states,
585 callback_id, callback_timeout_at, callback_filter, callback_on_complete,
586 callback_on_fail, callback_transform, run_lease, progress
587 )
588 SELECT
589 id,
590 kind,
591 queue,
592 args,
593 'available'::awa.job_state,
594 priority,
595 attempt,
596 max_attempts,
597 now(),
598 NULL,
599 NULL,
600 attempted_at,
601 finalized_at,
602 created_at,
603 errors,
604 metadata,
605 tags,
606 unique_key,
607 unique_states,
608 NULL,
609 NULL,
610 NULL,
611 NULL,
612 NULL,
613 NULL,
614 run_lease,
615 progress
616 FROM due
617 RETURNING queue
618 )
619 SELECT queue FROM promoted
620 "#,
621 )
622 .bind(state)
623 .bind(PROMOTE_BATCH_SIZE)
624 .fetch_all(&mut *tx)
625 .await?;
626
627 let promoted = promoted_rows.len();
628 self.metrics
629 .record_promotion_batch(state, promoted as u64, promote_start.elapsed());
630 if promoted == 0 {
631 tx.commit().await?;
632 return Ok((0, HashSet::new()));
633 }
634
635 let queues: HashSet<String> = promoted_rows.into_iter().map(|(queue,)| queue).collect();
636
637 tx.commit().await?;
638 Ok((promoted, queues))
639 }
640
641 #[tracing::instrument(skip(self), name = "maintenance.cleanup")]
648 async fn cleanup_completed(&self) {
649 let mut total_deleted: u64 = 0;
650
651 let override_queues: Vec<String> = self.queue_retention_overrides.keys().cloned().collect();
653
654 let completed_retention = format!("{} seconds", self.completed_retention.as_secs());
656 let failed_retention = format!("{} seconds", self.failed_retention.as_secs());
657
658 let global_result = if override_queues.is_empty() {
659 sqlx::query(
660 r#"
661 DELETE FROM awa.jobs_hot
662 WHERE id IN (
663 SELECT id FROM awa.jobs_hot
664 WHERE (state = 'completed' AND finalized_at < now() - $1::interval)
665 OR (state IN ('failed', 'cancelled') AND finalized_at < now() - $2::interval)
666 LIMIT $3
667 )
668 "#,
669 )
670 .bind(&completed_retention)
671 .bind(&failed_retention)
672 .bind(self.cleanup_batch_size)
673 .execute(&self.pool)
674 .await
675 } else {
676 sqlx::query(
677 r#"
678 DELETE FROM awa.jobs_hot
679 WHERE id IN (
680 SELECT id FROM awa.jobs_hot
681 WHERE ((state = 'completed' AND finalized_at < now() - $1::interval)
682 OR (state IN ('failed', 'cancelled') AND finalized_at < now() - $2::interval))
683 AND queue != ALL($4::text[])
684 LIMIT $3
685 )
686 "#,
687 )
688 .bind(&completed_retention)
689 .bind(&failed_retention)
690 .bind(self.cleanup_batch_size)
691 .bind(&override_queues)
692 .execute(&self.pool)
693 .await
694 };
695
696 match global_result {
697 Ok(result) if result.rows_affected() > 0 => {
698 total_deleted += result.rows_affected();
699 }
700 Err(err) => {
701 error!(error = %err, "Failed to clean up old jobs (global pass)");
702 }
703 _ => {}
704 }
705
706 for (queue_name, policy) in &self.queue_retention_overrides {
708 let queue_completed = format!("{} seconds", policy.completed.as_secs());
709 let queue_failed = format!("{} seconds", policy.failed.as_secs());
710
711 match sqlx::query(
712 r#"
713 DELETE FROM awa.jobs_hot
714 WHERE id IN (
715 SELECT id FROM awa.jobs_hot
716 WHERE queue = $4
717 AND ((state = 'completed' AND finalized_at < now() - $1::interval)
718 OR (state IN ('failed', 'cancelled') AND finalized_at < now() - $2::interval))
719 LIMIT $3
720 )
721 "#,
722 )
723 .bind(&queue_completed)
724 .bind(&queue_failed)
725 .bind(self.cleanup_batch_size)
726 .bind(queue_name)
727 .execute(&self.pool)
728 .await
729 {
730 Ok(result) if result.rows_affected() > 0 => {
731 total_deleted += result.rows_affected();
732 debug!(
733 queue = %queue_name,
734 count = result.rows_affected(),
735 "Cleaned up old jobs (queue override)"
736 );
737 }
738 Err(err) => {
739 error!(
740 queue = %queue_name,
741 error = %err,
742 "Failed to clean up old jobs (queue override)"
743 );
744 }
745 _ => {}
746 }
747 }
748
749 if total_deleted > 0 {
750 info!(count = total_deleted, "Cleaned up old jobs");
751 }
752 }
753}
754
755fn compute_fire_time(
759 row: &CronJobRow,
760 now: chrono::DateTime<Utc>,
761) -> Option<chrono::DateTime<Utc>> {
762 let cron = match Cron::new(&row.cron_expr).parse() {
763 Ok(c) => c,
764 Err(err) => {
765 error!(cron_name = %row.name, error = %err, "Invalid cron expression in database");
766 return None;
767 }
768 };
769
770 let tz: chrono_tz::Tz = match row.timezone.parse() {
771 Ok(tz) => tz,
772 Err(err) => {
773 error!(cron_name = %row.name, error = %err, "Invalid timezone in database");
774 return None;
775 }
776 };
777
778 let search_start = match row.last_enqueued_at {
779 Some(last) => last.with_timezone(&tz),
780 None => (row.created_at - chrono::Duration::minutes(1)).with_timezone(&tz),
785 };
786
787 let mut latest_fire: Option<chrono::DateTime<Utc>> = None;
788
789 for fire_time in cron.iter_from(search_start) {
790 let fire_utc = fire_time.with_timezone(&Utc);
791
792 if fire_utc > now {
793 break;
794 }
795
796 if let Some(last) = row.last_enqueued_at {
797 if fire_utc <= last {
798 continue;
799 }
800 }
801
802 latest_fire = Some(fire_utc);
803 }
804
805 latest_fire
806}