1use crate::completion::CompletionBatcher;
2use crate::dispatcher::{
3 shared_rate_limiter, ConcurrencyMode, Dispatcher, OverflowPool, QueueConfig,
4};
5use crate::events::{BoxedUntypedEventHandler, JobEvent, UntypedJobEvent};
6use crate::executor::{BoxedWorker, DlqPolicy, JobError, JobExecutor, JobResult, Worker};
7use crate::heartbeat::HeartbeatService;
8use crate::maintenance::{MaintenanceService, RetentionPolicy};
9use crate::runtime::{InFlightMap, InFlightRegistry};
10use crate::storage::{QueueStorageRuntime, RuntimeStorage};
11use awa_model::admin::{
12 self, JobKindDescriptor, NamedJobKindDescriptor, NamedQueueDescriptor, QueueDescriptor,
13 QueueRuntimeConfigSnapshot, QueueRuntimeMode, QueueRuntimeSnapshot, RateLimitSnapshot,
14 RuntimeSnapshotInput, StorageCapability, TransitionRole,
15};
16use awa_model::{
17 storage as transition, JobArgs, PartitionedQueue, PeriodicJob, QueueStorageConfig,
18};
19use chrono::{DateTime, Utc};
20use serde::de::DeserializeOwned;
21use sqlx::PgPool;
22use std::any::{Any, TypeId};
23use std::collections::{HashMap, HashSet};
24use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
25use std::sync::Arc;
26use std::time::Duration;
27use tokio::sync::{Mutex, RwLock};
28use tokio::task::JoinSet;
29use tokio_util::sync::CancellationToken;
30use tracing::{info, warn};
31use uuid::Uuid;
32
33#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
35pub enum BuildError {
36 #[error("at least one queue must be configured")]
37 NoQueuesConfigured,
38 #[error("queue descriptor declared for unknown queue '{queue}'")]
39 QueueDescriptorWithoutQueue { queue: String },
40 #[error("queue '{queue}' configured more than once")]
41 DuplicateQueue { queue: String },
42 #[error("sum of min_workers ({total_min}) exceeds global_max_workers ({global_max})")]
43 MinWorkersExceedGlobal { total_min: u32, global_max: u32 },
44 #[error("rate_limit max_rate must be > 0.0")]
45 InvalidRateLimit,
46 #[error("queue weight must be > 0")]
47 InvalidWeight,
48 #[error("queue claimers must be > 0")]
49 InvalidClaimers,
50 #[error("queue claim_batch_size must be > 0")]
51 InvalidClaimBatchSize,
52 #[error("cleanup_batch_size must be > 0")]
53 InvalidBatchSize,
54 #[error("dlq_cleanup_batch_size must be > 0")]
55 InvalidDlqBatchSize,
56 #[error("terminal_count_rollup_interval must be > 0")]
57 InvalidTerminalCountRollupInterval,
58 #[error("invalid queue storage config: {0}")]
59 InvalidQueueStorage(String),
60}
61
62#[derive(Debug, Clone)]
64pub struct HealthCheck {
65 pub healthy: bool,
66 pub postgres_connected: bool,
67 pub poll_loop_alive: bool,
68 pub heartbeat_alive: bool,
69 pub maintenance_alive: bool,
70 pub shutting_down: bool,
71 pub leader: bool,
72 pub queues: HashMap<String, QueueHealth>,
73}
74
75#[derive(Debug, Clone)]
77pub struct QueueHealth {
78 pub in_flight: u32,
79 pub available: u64,
80 pub capacity: QueueCapacity,
82}
83
84#[derive(Debug, Clone)]
86pub enum QueueCapacity {
87 HardReserved { max_workers: u32 },
89 Weighted {
91 min_workers: u32,
92 weight: u32,
93 overflow_held: u32,
94 },
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
104pub enum TransitionWorkerRole {
105 #[default]
109 Auto,
110 CanonicalDrain,
112 QueueStorageTarget,
114}
115
116impl From<TransitionWorkerRole> for TransitionRole {
117 fn from(role: TransitionWorkerRole) -> Self {
118 match role {
119 TransitionWorkerRole::Auto => Self::Auto,
120 TransitionWorkerRole::CanonicalDrain => Self::CanonicalDrain,
121 TransitionWorkerRole::QueueStorageTarget => Self::QueueStorageTarget,
122 }
123 }
124}
125
126pub struct ClientBuilder {
128 pool: PgPool,
129 queues: Vec<(String, QueueConfig)>,
130 queue_descriptors: HashMap<String, QueueDescriptor>,
131 job_kind_descriptors: HashMap<String, JobKindDescriptor>,
132 workers: HashMap<String, BoxedWorker>,
133 lifecycle_handlers: HashMap<String, Vec<BoxedUntypedEventHandler>>,
134 enqueue_specs: HashMap<
135 crate::enqueue_specs::Outcome,
136 HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
137 >,
138 state: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
139 heartbeat_interval: Duration,
140 promote_interval: Duration,
141 heartbeat_rescue_interval: Option<Duration>,
142 heartbeat_staleness: Option<Duration>,
143 deadline_rescue_interval: Option<Duration>,
144 callback_rescue_interval: Option<Duration>,
145 periodic_jobs: Vec<PeriodicJob>,
146 global_max_workers: Option<u32>,
147 leader_election_interval: Option<Duration>,
148 leader_check_interval: Option<Duration>,
149 priority_aging_interval: Option<Duration>,
150 terminal_count_rollup_interval: Option<Duration>,
151 completed_retention: Option<Duration>,
152 failed_retention: Option<Duration>,
153 descriptor_retention: Option<Duration>,
154 cleanup_batch_size: Option<i64>,
155 cleanup_interval: Option<Duration>,
156 queue_retention_overrides: HashMap<String, RetentionPolicy>,
157 runtime_snapshot_interval: Duration,
158 queue_stats_interval: Option<Duration>,
159 dlq_enabled_by_default: bool,
160 dlq_retention: Option<Duration>,
161 dlq_cleanup_batch_size: Option<i64>,
162 dlq_overrides: HashMap<String, bool>,
163 storage: RuntimeStorage,
164 transition_role: TransitionWorkerRole,
165 storage_error: Option<BuildError>,
166}
167
168impl ClientBuilder {
169 pub fn new(pool: PgPool) -> Self {
170 let (storage, storage_error) = match QueueStorageRuntime::new(
176 QueueStorageConfig::default(),
177 Duration::from_millis(1_000),
178 Duration::from_millis(1_000),
179 ) {
180 Ok(runtime) => (RuntimeStorage::QueueStorage(runtime), None),
181 Err(err) => (
182 RuntimeStorage::Canonical,
183 Some(BuildError::InvalidQueueStorage(err.to_string())),
184 ),
185 };
186
187 Self {
188 pool,
189 queues: Vec::new(),
190 queue_descriptors: HashMap::new(),
191 job_kind_descriptors: HashMap::new(),
192 workers: HashMap::new(),
193 lifecycle_handlers: HashMap::new(),
194 enqueue_specs: HashMap::new(),
195 state: HashMap::new(),
196 heartbeat_interval: Duration::from_secs(30),
197 promote_interval: Duration::from_millis(250),
198 heartbeat_rescue_interval: None,
199 heartbeat_staleness: None,
200 deadline_rescue_interval: None,
201 callback_rescue_interval: None,
202 periodic_jobs: Vec::new(),
203 global_max_workers: None,
204 leader_election_interval: None,
205 leader_check_interval: None,
206 priority_aging_interval: None,
207 terminal_count_rollup_interval: None,
208 completed_retention: None,
209 failed_retention: None,
210 descriptor_retention: None,
211 cleanup_batch_size: None,
212 cleanup_interval: None,
213 queue_retention_overrides: HashMap::new(),
214 runtime_snapshot_interval: Duration::from_secs(10),
215 queue_stats_interval: None,
216 dlq_enabled_by_default: false,
217 dlq_retention: None,
218 dlq_cleanup_batch_size: None,
219 dlq_overrides: HashMap::new(),
220 storage,
221 transition_role: TransitionWorkerRole::Auto,
222 storage_error,
223 }
224 }
225
226 pub fn queue(mut self, name: impl Into<String>, config: QueueConfig) -> Self {
228 self.queues.push((name.into(), config));
229 self
230 }
231
232 pub fn partitioned_queue(
248 mut self,
249 partitioned_queue: &PartitionedQueue,
250 config: QueueConfig,
251 ) -> Self {
252 for queue in partitioned_queue.physical_queues() {
253 self.queues.push((queue.clone(), config.clone()));
254 }
255 self
256 }
257
258 pub fn queue_descriptor(
267 mut self,
268 name: impl Into<String>,
269 descriptor: QueueDescriptor,
270 ) -> Self {
271 self.queue_descriptors.insert(name.into(), descriptor);
272 self
273 }
274
275 pub fn job_kind_descriptor<T: JobArgs>(mut self, descriptor: JobKindDescriptor) -> Self {
278 self.job_kind_descriptors
279 .insert(T::kind().to_string(), descriptor);
280 self
281 }
282
283 pub fn job_kind_descriptor_kind(
286 mut self,
287 kind: impl Into<String>,
288 descriptor: JobKindDescriptor,
289 ) -> Self {
290 self.job_kind_descriptors.insert(kind.into(), descriptor);
291 self
292 }
293
294 pub fn register<T, F, Fut>(mut self, handler: F) -> Self
299 where
300 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
301 F: Fn(T, &crate::context::JobContext) -> Fut + Send + Sync + 'static,
302 Fut: std::future::Future<Output = Result<JobResult, JobError>> + Send + 'static,
303 {
304 let kind = T::kind().to_string();
305 let worker = TypedWorker {
306 kind: T::kind(),
307 handler: Arc::new(handler),
308 _phantom: std::marker::PhantomData,
309 };
310 self.workers.insert(kind, Box::new(worker));
311 self
312 }
313
314 pub fn on_event<T, F, Fut>(mut self, handler: F) -> Self
324 where
325 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
326 F: Fn(JobEvent<T>) -> Fut + Send + Sync + 'static,
327 Fut: std::future::Future<Output = ()> + Send + 'static,
328 {
329 let kind = T::kind().to_string();
330 let handler = Arc::new(handler);
331 let erased: BoxedUntypedEventHandler = Arc::new(move |event: UntypedJobEvent| {
332 let handler = handler.clone();
333 Box::pin(async move {
334 let args: T = match serde_json::from_value(event.job().args.clone()) {
335 Ok(args) => args,
336 Err(err) => {
337 warn!(
338 job_id = event.job().id,
339 kind = %event.job().kind,
340 error = %err,
341 "Failed to deserialize args for lifecycle event handler"
342 );
343 return;
344 }
345 };
346
347 (handler)(event.into_typed(args)).await;
348 })
349 });
350 self.lifecycle_handlers
351 .entry(kind)
352 .or_default()
353 .push(erased);
354 self
355 }
356
357 pub fn on_event_kind<F, Fut>(mut self, kind: impl Into<String>, handler: F) -> Self
362 where
363 F: Fn(UntypedJobEvent) -> Fut + Send + Sync + 'static,
364 Fut: std::future::Future<Output = ()> + Send + 'static,
365 {
366 let kind = kind.into();
367 let handler = Arc::new(handler);
368 let erased: BoxedUntypedEventHandler = Arc::new(move |event: UntypedJobEvent| {
369 let handler = handler.clone();
370 Box::pin(async move {
371 (handler)(event).await;
372 })
373 });
374 self.lifecycle_handlers
375 .entry(kind)
376 .or_default()
377 .push(erased);
378 self
379 }
380
381 pub fn on_completed_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
410 where
411 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
412 F: JobArgs + Send + Sync + 'static,
413 MakeFn: Fn(T, &awa_model::JobRow) -> F + Send + Sync + 'static,
414 {
415 self.on_completed_enqueue_with::<T, F, _>(move |args, job| {
416 crate::enqueue_specs::EnqueueRequest::new(make(args, job))
417 })
418 }
419
420 pub fn on_completed_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
424 where
425 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
426 F: JobArgs + Send + Sync + 'static,
427 MakeFn: Fn(T, &awa_model::JobRow) -> crate::enqueue_specs::EnqueueRequest<F>
428 + Send
429 + Sync
430 + 'static,
431 {
432 let kind = T::kind().to_string();
433 let spec: crate::enqueue_specs::BoxedEnqueueSpec =
434 Arc::new(crate::enqueue_specs::CompletedFollowUp::<T, F, _> {
435 make,
436 _phantom: std::marker::PhantomData,
437 });
438 self.enqueue_specs
439 .entry(crate::enqueue_specs::Outcome::Completed)
440 .or_default()
441 .entry(kind)
442 .or_default()
443 .push(spec);
444 self
445 }
446
447 pub fn on_retried_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
458 where
459 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
460 F: JobArgs + Send + Sync + 'static,
461 MakeFn: Fn(T, &awa_model::JobRow, &str, i16, chrono::DateTime<chrono::Utc>) -> F
462 + Send
463 + Sync
464 + 'static,
465 {
466 self.on_retried_enqueue_with::<T, F, _>(move |args, job, error, attempt, next_run_at| {
467 crate::enqueue_specs::EnqueueRequest::new(make(args, job, error, attempt, next_run_at))
468 })
469 }
470
471 pub fn on_retried_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
474 where
475 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
476 F: JobArgs + Send + Sync + 'static,
477 MakeFn: Fn(
478 T,
479 &awa_model::JobRow,
480 &str,
481 i16,
482 chrono::DateTime<chrono::Utc>,
483 ) -> crate::enqueue_specs::EnqueueRequest<F>
484 + Send
485 + Sync
486 + 'static,
487 {
488 let kind = T::kind().to_string();
489 let spec: crate::enqueue_specs::BoxedEnqueueSpec =
490 Arc::new(crate::enqueue_specs::RetriedFollowUp::<T, F, _> {
491 make,
492 _phantom: std::marker::PhantomData,
493 });
494 self.enqueue_specs
495 .entry(crate::enqueue_specs::Outcome::Retried)
496 .or_default()
497 .entry(kind)
498 .or_default()
499 .push(spec);
500 self
501 }
502
503 pub fn on_exhausted_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
515 where
516 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
517 F: JobArgs + Send + Sync + 'static,
518 MakeFn: Fn(T, &awa_model::JobRow, &str, i16) -> F + Send + Sync + 'static,
519 {
520 self.on_exhausted_enqueue_with::<T, F, _>(move |args, job, error, attempt| {
521 crate::enqueue_specs::EnqueueRequest::new(make(args, job, error, attempt))
522 })
523 }
524
525 pub fn on_exhausted_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
528 where
529 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
530 F: JobArgs + Send + Sync + 'static,
531 MakeFn: Fn(T, &awa_model::JobRow, &str, i16) -> crate::enqueue_specs::EnqueueRequest<F>
532 + Send
533 + Sync
534 + 'static,
535 {
536 let kind = T::kind().to_string();
537 let spec: crate::enqueue_specs::BoxedEnqueueSpec =
538 Arc::new(crate::enqueue_specs::ExhaustedFollowUp::<T, F, _> {
539 make,
540 _phantom: std::marker::PhantomData,
541 });
542 self.enqueue_specs
543 .entry(crate::enqueue_specs::Outcome::Exhausted)
544 .or_default()
545 .entry(kind)
546 .or_default()
547 .push(spec);
548 self
549 }
550
551 pub fn on_cancelled_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
557 where
558 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
559 F: JobArgs + Send + Sync + 'static,
560 MakeFn: Fn(T, &awa_model::JobRow, &str) -> F + Send + Sync + 'static,
561 {
562 self.on_cancelled_enqueue_with::<T, F, _>(move |args, job, reason| {
563 crate::enqueue_specs::EnqueueRequest::new(make(args, job, reason))
564 })
565 }
566
567 pub fn on_cancelled_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
570 where
571 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
572 F: JobArgs + Send + Sync + 'static,
573 MakeFn: Fn(T, &awa_model::JobRow, &str) -> crate::enqueue_specs::EnqueueRequest<F>
574 + Send
575 + Sync
576 + 'static,
577 {
578 let kind = T::kind().to_string();
579 let spec: crate::enqueue_specs::BoxedEnqueueSpec =
580 Arc::new(crate::enqueue_specs::CancelledFollowUp::<T, F, _> {
581 make,
582 _phantom: std::marker::PhantomData,
583 });
584 self.enqueue_specs
585 .entry(crate::enqueue_specs::Outcome::Cancelled)
586 .or_default()
587 .entry(kind)
588 .or_default()
589 .push(spec);
590 self
591 }
592
593 pub fn on_waiting_for_callback_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
601 where
602 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
603 F: JobArgs + Send + Sync + 'static,
604 MakeFn: Fn(T, &awa_model::JobRow) -> F + Send + Sync + 'static,
605 {
606 self.on_waiting_for_callback_enqueue_with::<T, F, _>(move |args, job| {
607 crate::enqueue_specs::EnqueueRequest::new(make(args, job))
608 })
609 }
610
611 pub fn on_waiting_for_callback_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
614 where
615 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
616 F: JobArgs + Send + Sync + 'static,
617 MakeFn: Fn(T, &awa_model::JobRow) -> crate::enqueue_specs::EnqueueRequest<F>
618 + Send
619 + Sync
620 + 'static,
621 {
622 let kind = T::kind().to_string();
623 let spec: crate::enqueue_specs::BoxedEnqueueSpec = Arc::new(
624 crate::enqueue_specs::WaitingForCallbackFollowUp::<T, F, _> {
625 make,
626 _phantom: std::marker::PhantomData,
627 },
628 );
629 self.enqueue_specs
630 .entry(crate::enqueue_specs::Outcome::WaitingForCallback)
631 .or_default()
632 .entry(kind)
633 .or_default()
634 .push(spec);
635 self
636 }
637
638 pub fn on_rescued_enqueue<T, F, MakeFn>(self, make: MakeFn) -> Self
652 where
653 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
654 F: JobArgs + Send + Sync + 'static,
655 MakeFn: Fn(T, &awa_model::JobRow, crate::events::RescueReason) -> F + Send + Sync + 'static,
656 {
657 self.on_rescued_enqueue_with::<T, F, _>(move |args, job, reason| {
658 crate::enqueue_specs::EnqueueRequest::new(make(args, job, reason))
659 })
660 }
661
662 pub fn on_rescued_enqueue_with<T, F, MakeFn>(mut self, make: MakeFn) -> Self
665 where
666 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
667 F: JobArgs + Send + Sync + 'static,
668 MakeFn: Fn(
669 T,
670 &awa_model::JobRow,
671 crate::events::RescueReason,
672 ) -> crate::enqueue_specs::EnqueueRequest<F>
673 + Send
674 + Sync
675 + 'static,
676 {
677 let kind = T::kind().to_string();
678 let spec: crate::enqueue_specs::BoxedEnqueueSpec =
679 Arc::new(crate::enqueue_specs::RescuedFollowUp::<T, F, _> {
680 make,
681 _phantom: std::marker::PhantomData,
682 });
683 self.enqueue_specs
684 .entry(crate::enqueue_specs::Outcome::Rescued)
685 .or_default()
686 .entry(kind)
687 .or_default()
688 .push(spec);
689 self
690 }
691
692 pub fn register_worker(mut self, worker: impl Worker + 'static) -> Self {
694 let kind = worker.kind().to_string();
695 self.workers.insert(kind, Box::new(worker));
696 self
697 }
698
699 #[cfg(feature = "http-worker")]
706 pub fn http_worker(
707 self,
708 kind: impl Into<String>,
709 config: crate::http_worker::HttpWorkerConfig,
710 ) -> Self {
711 let worker = crate::http_worker::HttpWorker::new(kind.into(), config);
712 self.register_worker(worker)
713 }
714
715 pub fn state<T: Any + Send + Sync + Clone>(mut self, value: T) -> Self {
717 self.state.insert(TypeId::of::<T>(), Box::new(value));
718 self
719 }
720
721 pub fn heartbeat_interval(mut self, interval: Duration) -> Self {
723 self.heartbeat_interval = interval;
724 self
725 }
726
727 pub fn promote_interval(mut self, interval: Duration) -> Self {
729 self.promote_interval = interval;
730 self
731 }
732
733 pub fn heartbeat_rescue_interval(mut self, interval: Duration) -> Self {
735 self.heartbeat_rescue_interval = Some(interval);
736 self
737 }
738
739 pub fn heartbeat_staleness(mut self, staleness: Duration) -> Self {
743 self.heartbeat_staleness = Some(staleness);
744 self
745 }
746
747 pub fn deadline_rescue_interval(mut self, interval: Duration) -> Self {
749 self.deadline_rescue_interval = Some(interval);
750 self
751 }
752
753 pub fn callback_rescue_interval(mut self, interval: Duration) -> Self {
755 self.callback_rescue_interval = Some(interval);
756 self
757 }
758
759 pub fn leader_election_interval(mut self, interval: Duration) -> Self {
764 self.leader_election_interval = Some(interval);
765 self
766 }
767
768 pub fn leader_check_interval(mut self, interval: Duration) -> Self {
770 self.leader_check_interval = Some(interval);
771 self
772 }
773
774 pub fn global_max_workers(mut self, max: u32) -> Self {
779 self.global_max_workers = Some(max);
780 self
781 }
782
783 pub fn completed_retention(mut self, retention: Duration) -> Self {
785 self.completed_retention = Some(retention);
786 self
787 }
788
789 pub fn failed_retention(mut self, retention: Duration) -> Self {
791 self.failed_retention = Some(retention);
792 self
793 }
794
795 pub fn descriptor_retention(mut self, retention: Duration) -> Self {
800 self.descriptor_retention = Some(retention);
801 self
802 }
803
804 pub fn cleanup_batch_size(mut self, batch_size: i64) -> Self {
806 self.cleanup_batch_size = Some(batch_size);
807 self
808 }
809
810 pub fn cleanup_interval(mut self, interval: Duration) -> Self {
812 self.cleanup_interval = Some(interval);
813 self
814 }
815
816 pub fn queue_retention(mut self, queue: impl Into<String>, policy: RetentionPolicy) -> Self {
818 self.queue_retention_overrides.insert(queue.into(), policy);
819 self
820 }
821
822 pub fn runtime_snapshot_interval(mut self, interval: Duration) -> Self {
824 self.runtime_snapshot_interval = interval;
825 self
826 }
827
828 pub fn priority_aging_interval(mut self, interval: Duration) -> Self {
834 self.priority_aging_interval = Some(interval);
835 self
836 }
837
838 pub fn terminal_count_rollup_interval(mut self, interval: Duration) -> Self {
845 self.terminal_count_rollup_interval = Some(interval);
846 self
847 }
848
849 pub fn queue_stats_interval(mut self, interval: Duration) -> Self {
851 self.queue_stats_interval = Some(interval);
852 self
853 }
854
855 pub fn dlq_enabled_by_default(mut self, enabled: bool) -> Self {
857 self.dlq_enabled_by_default = enabled;
858 self
859 }
860
861 pub fn queue_dlq_enabled(mut self, queue: impl Into<String>, enabled: bool) -> Self {
863 self.dlq_overrides.insert(queue.into(), enabled);
864 self
865 }
866
867 pub fn dlq_retention(mut self, retention: Duration) -> Self {
869 self.dlq_retention = Some(retention);
870 self
871 }
872
873 pub fn dlq_cleanup_batch_size(mut self, batch_size: i64) -> Self {
875 self.dlq_cleanup_batch_size = Some(batch_size);
876 self
877 }
878
879 pub fn queue_storage(
885 mut self,
886 config: QueueStorageConfig,
887 queue_rotate_interval: Duration,
888 lease_rotate_interval: Duration,
889 ) -> Self {
890 match QueueStorageRuntime::new(config, queue_rotate_interval, lease_rotate_interval) {
891 Ok(runtime) => {
892 self.storage = RuntimeStorage::QueueStorage(runtime);
893 self.storage_error = None;
894 }
895 Err(err) => {
896 self.storage = RuntimeStorage::Canonical;
897 self.storage_error = Some(BuildError::InvalidQueueStorage(err.to_string()));
898 }
899 }
900 self
901 }
902
903 pub fn claim_rotate_interval(mut self, claim_rotate_interval: Duration) -> Self {
909 if let RuntimeStorage::QueueStorage(runtime) = self.storage {
910 self.storage = RuntimeStorage::QueueStorage(
911 runtime.with_claim_rotate_interval(claim_rotate_interval),
912 );
913 }
914 self
915 }
916
917 pub fn canonical_storage(mut self) -> Self {
923 self.storage = RuntimeStorage::Canonical;
924 self.storage_error = None;
925 self
926 }
927
928 pub fn transition_role(mut self, role: TransitionWorkerRole) -> Self {
930 self.transition_role = role;
931 self
932 }
933
934 pub fn periodic(mut self, job: PeriodicJob) -> Self {
939 self.periodic_jobs.push(job);
940 self
941 }
942
943 pub fn build(self) -> Result<Client, BuildError> {
945 if self.queues.is_empty() {
946 return Err(BuildError::NoQueuesConfigured);
947 }
948
949 if let Some(err) = self.storage_error.clone() {
950 return Err(err);
951 }
952
953 let mut queue_names = HashSet::with_capacity(self.queues.len());
954 for (queue, _) in &self.queues {
955 if !queue_names.insert(queue.as_str()) {
956 return Err(BuildError::DuplicateQueue {
957 queue: queue.clone(),
958 });
959 }
960 }
961
962 for queue in self.queue_descriptors.keys() {
963 if !self.queues.iter().any(|(name, _)| name == queue) {
964 return Err(BuildError::QueueDescriptorWithoutQueue {
965 queue: queue.clone(),
966 });
967 }
968 }
969
970 for (_, config) in &self.queues {
972 if let Some(rl) = &config.rate_limit {
973 if rl.max_rate <= 0.0 {
974 return Err(BuildError::InvalidRateLimit);
975 }
976 }
977 if config.weight == 0 {
978 return Err(BuildError::InvalidWeight);
979 }
980 if config.claimers == 0 {
981 return Err(BuildError::InvalidClaimers);
982 }
983 if config.claim_batch_size == 0 {
984 return Err(BuildError::InvalidClaimBatchSize);
985 }
986 }
987
988 if let Some(bs) = self.cleanup_batch_size {
990 if bs <= 0 {
991 return Err(BuildError::InvalidBatchSize);
992 }
993 }
994 if let Some(bs) = self.dlq_cleanup_batch_size {
995 if bs <= 0 {
996 return Err(BuildError::InvalidDlqBatchSize);
997 }
998 }
999 if self
1000 .terminal_count_rollup_interval
1001 .is_some_and(|interval| interval.is_zero())
1002 {
1003 return Err(BuildError::InvalidTerminalCountRollupInterval);
1004 }
1005
1006 let overflow_pool = if let Some(global_max) = self.global_max_workers {
1008 let total_min: u32 = self.queues.iter().map(|(_, c)| c.min_workers).sum();
1009 if total_min > global_max {
1010 return Err(BuildError::MinWorkersExceedGlobal {
1011 total_min,
1012 global_max,
1013 });
1014 }
1015 let overflow_capacity = global_max - total_min;
1016 let weights: HashMap<String, u32> = self
1017 .queues
1018 .iter()
1019 .map(|(name, c)| (name.clone(), c.weight.max(1)))
1020 .collect();
1021 Some(Arc::new(OverflowPool::new(overflow_capacity, weights)))
1022 } else {
1023 None
1024 };
1025
1026 if let Some(staleness) = self.heartbeat_staleness {
1028 let min_safe = self.heartbeat_interval * 3;
1029 if staleness < min_safe {
1030 tracing::warn!(
1031 heartbeat_staleness_ms = staleness.as_millis() as u64,
1032 heartbeat_interval_ms = self.heartbeat_interval.as_millis() as u64,
1033 recommended_min_ms = min_safe.as_millis() as u64,
1034 "heartbeat_staleness ({:?}) is less than 3× heartbeat_interval ({:?}); \
1035 this may cause false rescues of jobs that are still running",
1036 staleness,
1037 self.heartbeat_interval,
1038 );
1039 }
1040 }
1041
1042 let metrics = crate::metrics::AwaMetrics::from_global();
1043 let queue_in_flight = Arc::new(
1044 self.queues
1045 .iter()
1046 .map(|(name, _)| (name.clone(), Arc::new(AtomicU32::new(0))))
1047 .collect(),
1048 );
1049 let dispatcher_alive = Arc::new(
1050 self.queues
1051 .iter()
1052 .map(|(name, _)| (name.clone(), Arc::new(AtomicBool::new(false))))
1053 .collect(),
1054 );
1055 let dlq_policy = DlqPolicy::new(self.dlq_enabled_by_default, self.dlq_overrides);
1056
1057 Ok(Client {
1058 pool: self.pool,
1059 queues: self.queues,
1060 queue_descriptors: self.queue_descriptors,
1061 job_kind_descriptors: self.job_kind_descriptors,
1062 workers: Arc::new(self.workers),
1063 lifecycle_handlers: Arc::new(self.lifecycle_handlers),
1064 enqueue_specs: Arc::new(self.enqueue_specs),
1065 state: Arc::new(self.state),
1066 heartbeat_interval: self.heartbeat_interval,
1067 promote_interval: self.promote_interval,
1068 heartbeat_rescue_interval: self.heartbeat_rescue_interval,
1069 heartbeat_staleness: self.heartbeat_staleness,
1070 deadline_rescue_interval: self.deadline_rescue_interval,
1071 callback_rescue_interval: self.callback_rescue_interval,
1072 periodic_jobs: Arc::new(self.periodic_jobs),
1073 dispatch_cancel: CancellationToken::new(),
1074 service_cancel: CancellationToken::new(),
1075 dispatcher_handles: RwLock::new(Vec::new()),
1076 service_handles: RwLock::new(Vec::new()),
1077 job_set: Arc::new(Mutex::new(JoinSet::new())),
1078 in_flight: Arc::new(InFlightRegistry::default()),
1079 queue_in_flight,
1080 dispatcher_alive,
1081 heartbeat_alive: Arc::new(AtomicBool::new(false)),
1082 maintenance_alive: Arc::new(AtomicBool::new(false)),
1083 leader: Arc::new(AtomicBool::new(false)),
1084 overflow_pool,
1085 metrics,
1086 leader_election_interval: self.leader_election_interval,
1087 leader_check_interval: self.leader_check_interval,
1088 priority_aging_interval: self.priority_aging_interval,
1089 terminal_count_rollup_interval: self.terminal_count_rollup_interval,
1090 completed_retention: self.completed_retention,
1091 failed_retention: self.failed_retention,
1092 descriptor_retention: self.descriptor_retention,
1093 cleanup_batch_size: self.cleanup_batch_size,
1094 cleanup_interval: self.cleanup_interval,
1095 queue_retention_overrides: self.queue_retention_overrides,
1096 queue_stats_interval: self.queue_stats_interval,
1097 dlq_policy,
1098 dlq_retention: self.dlq_retention,
1099 dlq_cleanup_batch_size: self.dlq_cleanup_batch_size,
1100 effective_storage: Arc::new(RwLock::new(self.storage.clone())),
1101 storage: self.storage,
1102 transition_role: self.transition_role,
1103 global_max_workers: self.global_max_workers,
1104 runtime_snapshot_interval: self.runtime_snapshot_interval,
1105 runtime_instance_id: Uuid::new_v4(),
1106 runtime_started_at: Utc::now(),
1107 runtime_hostname: std::env::var("HOSTNAME").ok(),
1108 runtime_pid: std::process::id() as i32,
1109 runtime_version: env!("CARGO_PKG_VERSION"),
1110 })
1111 }
1112}
1113
1114struct TypedWorker<T, F, Fut>
1116where
1117 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
1118 F: Fn(T, &crate::context::JobContext) -> Fut + Send + Sync + 'static,
1119 Fut: std::future::Future<Output = Result<JobResult, JobError>> + Send + 'static,
1120{
1121 kind: &'static str,
1122 handler: Arc<F>,
1123 _phantom: std::marker::PhantomData<fn() -> (T, Fut)>,
1124}
1125
1126#[async_trait::async_trait]
1127impl<T, F, Fut> Worker for TypedWorker<T, F, Fut>
1128where
1129 T: JobArgs + DeserializeOwned + Send + Sync + 'static,
1130 F: Fn(T, &crate::context::JobContext) -> Fut + Send + Sync + 'static,
1131 Fut: std::future::Future<Output = Result<JobResult, JobError>> + Send + 'static,
1132{
1133 fn kind(&self) -> &'static str {
1134 self.kind
1135 }
1136
1137 async fn perform(&self, ctx: &crate::context::JobContext) -> Result<JobResult, JobError> {
1138 let args: T = serde_json::from_value(ctx.job.args.clone())
1139 .map_err(|err| JobError::Terminal(format!("failed to deserialize args: {}", err)))?;
1140
1141 (self.handler)(args, ctx).await
1142 }
1143}
1144
1145pub struct Client {
1147 pool: PgPool,
1148 queues: Vec<(String, QueueConfig)>,
1149 queue_descriptors: HashMap<String, QueueDescriptor>,
1150 job_kind_descriptors: HashMap<String, JobKindDescriptor>,
1151 workers: Arc<HashMap<String, BoxedWorker>>,
1152 lifecycle_handlers: Arc<HashMap<String, Vec<BoxedUntypedEventHandler>>>,
1153 enqueue_specs: Arc<
1154 HashMap<
1155 crate::enqueue_specs::Outcome,
1156 HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
1157 >,
1158 >,
1159 state: Arc<HashMap<TypeId, Box<dyn Any + Send + Sync>>>,
1160 heartbeat_interval: Duration,
1161 promote_interval: Duration,
1162 heartbeat_rescue_interval: Option<Duration>,
1163 heartbeat_staleness: Option<Duration>,
1164 deadline_rescue_interval: Option<Duration>,
1165 callback_rescue_interval: Option<Duration>,
1166 periodic_jobs: Arc<Vec<PeriodicJob>>,
1167 dispatch_cancel: CancellationToken,
1169 service_cancel: CancellationToken,
1171 dispatcher_handles: RwLock<Vec<tokio::task::JoinHandle<()>>>,
1173 service_handles: RwLock<Vec<tokio::task::JoinHandle<()>>>,
1175 job_set: Arc<Mutex<JoinSet<()>>>,
1177 in_flight: InFlightMap,
1178 queue_in_flight: Arc<HashMap<String, Arc<AtomicU32>>>,
1179 dispatcher_alive: Arc<HashMap<String, Arc<AtomicBool>>>,
1180 heartbeat_alive: Arc<AtomicBool>,
1181 maintenance_alive: Arc<AtomicBool>,
1182 leader: Arc<AtomicBool>,
1183 overflow_pool: Option<Arc<OverflowPool>>,
1185 metrics: crate::metrics::AwaMetrics,
1186 leader_election_interval: Option<Duration>,
1187 leader_check_interval: Option<Duration>,
1188 priority_aging_interval: Option<Duration>,
1189 terminal_count_rollup_interval: Option<Duration>,
1190 completed_retention: Option<Duration>,
1191 failed_retention: Option<Duration>,
1192 descriptor_retention: Option<Duration>,
1193 cleanup_batch_size: Option<i64>,
1194 cleanup_interval: Option<Duration>,
1195 queue_retention_overrides: HashMap<String, RetentionPolicy>,
1196 queue_stats_interval: Option<Duration>,
1197 dlq_policy: DlqPolicy,
1198 dlq_retention: Option<Duration>,
1199 dlq_cleanup_batch_size: Option<i64>,
1200 storage: RuntimeStorage,
1201 transition_role: TransitionWorkerRole,
1202 effective_storage: Arc<RwLock<RuntimeStorage>>,
1203 global_max_workers: Option<u32>,
1204 runtime_snapshot_interval: Duration,
1205 runtime_instance_id: Uuid,
1206 runtime_started_at: DateTime<Utc>,
1207 runtime_hostname: Option<String>,
1208 runtime_pid: i32,
1209 runtime_version: &'static str,
1210}
1211
1212#[derive(Clone)]
1213struct RuntimeReporterState {
1214 pool: PgPool,
1215 queues: Vec<(String, QueueConfig)>,
1216 queue_descriptors: HashMap<String, QueueDescriptor>,
1217 job_kind_descriptors: HashMap<String, JobKindDescriptor>,
1218 worker_kinds: Vec<String>,
1219 queue_in_flight: Arc<HashMap<String, Arc<AtomicU32>>>,
1220 dispatcher_alive: Arc<HashMap<String, Arc<AtomicBool>>>,
1221 heartbeat_alive: Arc<AtomicBool>,
1222 maintenance_alive: Arc<AtomicBool>,
1223 leader: Arc<AtomicBool>,
1224 dispatch_cancel: CancellationToken,
1225 overflow_pool: Option<Arc<OverflowPool>>,
1226 global_max_workers: Option<u32>,
1227 dlq_policy: DlqPolicy,
1228 instance_id: Uuid,
1229 started_at: DateTime<Utc>,
1230 hostname: Option<String>,
1231 pid: i32,
1232 version: &'static str,
1233 snapshot_interval: Duration,
1234 effective_storage: Arc<RwLock<RuntimeStorage>>,
1235 queue_storage_capable: bool,
1236 transition_role: TransitionWorkerRole,
1237 metrics: crate::metrics::AwaMetrics,
1238}
1239
1240fn latest_error_message(job: &awa_model::JobRow) -> String {
1244 job.errors
1245 .as_ref()
1246 .and_then(|errors| errors.last())
1247 .and_then(|entry| entry.get("error"))
1248 .and_then(|value| value.as_str())
1249 .map(str::to_string)
1250 .unwrap_or_default()
1251}
1252
1253impl Client {
1254 pub fn builder(pool: PgPool) -> ClientBuilder {
1256 ClientBuilder::new(pool)
1257 }
1258
1259 fn expected_queue_storage_schema(
1260 status: &transition::StorageStatus,
1261 ) -> Result<Option<String>, awa_model::AwaError> {
1262 let prepared_schema = || {
1263 status
1264 .details
1265 .get("schema")
1266 .and_then(serde_json::Value::as_str)
1267 .unwrap_or("awa")
1268 .to_string()
1269 };
1270
1271 match status.state.as_str() {
1272 "prepared" if status.prepared_engine.as_deref() == Some("queue_storage") => {
1273 Ok(Some(prepared_schema()))
1274 }
1275 "mixed_transition" | "active" if status.active_engine == "queue_storage" => {
1276 Ok(Some(prepared_schema()))
1277 }
1278 "canonical" if status.prepared_engine.as_deref() == Some("queue_storage") => {
1279 Ok(Some(prepared_schema()))
1280 }
1281 "mixed_transition" | "active" => Err(awa_model::AwaError::Validation(format!(
1282 "unsupported active storage engine '{}'",
1283 status.active_engine
1284 ))),
1285 _ => Ok(None),
1286 }
1287 }
1288
1289 async fn resolve_effective_storage(&self) -> Result<RuntimeStorage, awa_model::AwaError> {
1290 let Some(runtime) = self.storage.queue_storage() else {
1291 return Ok(RuntimeStorage::Canonical);
1292 };
1293
1294 let status = transition::status(&self.pool).await?;
1295 let expected_schema = Self::expected_queue_storage_schema(&status)?;
1296 let prepared_schema_ready = if let Some(schema) = expected_schema.as_deref() {
1297 if runtime.store.schema() != schema {
1298 return Err(awa_model::AwaError::Validation(format!(
1299 "queue storage runtime configured for schema '{}' but transition state requires '{}'",
1300 runtime.store.schema(),
1301 schema
1302 )));
1303 }
1304 transition::queue_storage_schema_ready(&self.pool, schema).await?
1305 } else {
1306 false
1307 };
1308
1309 match self.transition_role {
1310 TransitionWorkerRole::CanonicalDrain => Ok(RuntimeStorage::Canonical),
1311 TransitionWorkerRole::QueueStorageTarget => {
1312 let schema = expected_schema.ok_or_else(|| {
1313 awa_model::AwaError::Validation(
1314 "queue_storage_target requires a prepared queue-storage schema".into(),
1315 )
1316 })?;
1317 if !prepared_schema_ready {
1318 return Err(awa_model::AwaError::Validation(format!(
1319 "queue storage schema '{schema}' is not prepared; run schema preparation before starting queue-storage-target runtimes"
1320 )));
1321 }
1322 Ok(RuntimeStorage::QueueStorage(runtime.clone()))
1323 }
1324 TransitionWorkerRole::Auto => {
1325 if let Some(schema) = expected_schema.as_deref() {
1326 if !prepared_schema_ready {
1327 return Err(awa_model::AwaError::Validation(format!(
1328 "queue storage schema '{schema}' is not prepared; run schema preparation before starting 0.6 runtimes"
1329 )));
1330 }
1331 }
1332
1333 if status.state == "canonical" && status.prepared_engine.is_none() {
1345 let configured_schema = runtime.store.schema().to_string();
1346 if !transition::queue_storage_schema_ready(&self.pool, &configured_schema)
1347 .await?
1348 {
1349 runtime.store.prepare_schema(&self.pool).await?;
1350 }
1351 let promoted: bool =
1352 sqlx::query_scalar("SELECT awa.storage_auto_finalize_if_fresh($1)")
1353 .bind(&configured_schema)
1354 .fetch_one(&self.pool)
1355 .await?;
1356 if promoted {
1357 return Ok(RuntimeStorage::QueueStorage(runtime.clone()));
1358 }
1359 let refetched = transition::status(&self.pool).await?;
1363 if matches!(refetched.state.as_str(), "mixed_transition" | "active")
1364 && refetched.active_engine == "queue_storage"
1365 {
1366 return Ok(RuntimeStorage::QueueStorage(runtime.clone()));
1367 }
1368 }
1374
1375 if matches!(status.state.as_str(), "mixed_transition" | "active")
1376 && status.active_engine == "queue_storage"
1377 {
1378 Ok(RuntimeStorage::QueueStorage(runtime.clone()))
1379 } else {
1380 Ok(RuntimeStorage::Canonical)
1381 }
1382 }
1383 }
1384 }
1385
1386 fn declared_queue_descriptors(&self) -> Vec<NamedQueueDescriptor> {
1387 self.queues
1388 .iter()
1389 .map(|(queue, _)| NamedQueueDescriptor {
1390 queue: queue.clone(),
1391 descriptor: self
1392 .queue_descriptors
1393 .get(queue)
1394 .cloned()
1395 .unwrap_or_default(),
1396 })
1397 .collect()
1398 }
1399
1400 fn declared_job_kind_descriptors(&self) -> Vec<NamedJobKindDescriptor> {
1401 let mut kinds: Vec<String> = self.workers.keys().cloned().collect();
1402 for kind in self.job_kind_descriptors.keys() {
1403 if !kinds.iter().any(|existing| existing == kind) {
1404 kinds.push(kind.clone());
1405 }
1406 }
1407 kinds.sort();
1408
1409 kinds
1410 .into_iter()
1411 .map(|kind| NamedJobKindDescriptor {
1412 descriptor: self
1413 .job_kind_descriptors
1414 .get(&kind)
1415 .cloned()
1416 .unwrap_or_default(),
1417 kind,
1418 })
1419 .collect()
1420 }
1421
1422 fn runtime_reporter_state(&self) -> RuntimeReporterState {
1423 RuntimeReporterState {
1424 pool: self.pool.clone(),
1425 queues: self.queues.clone(),
1426 queue_descriptors: self.queue_descriptors.clone(),
1427 job_kind_descriptors: self.job_kind_descriptors.clone(),
1428 worker_kinds: self.workers.keys().cloned().collect(),
1429 queue_in_flight: self.queue_in_flight.clone(),
1430 dispatcher_alive: self.dispatcher_alive.clone(),
1431 heartbeat_alive: self.heartbeat_alive.clone(),
1432 maintenance_alive: self.maintenance_alive.clone(),
1433 leader: self.leader.clone(),
1434 dispatch_cancel: self.dispatch_cancel.clone(),
1435 overflow_pool: self.overflow_pool.clone(),
1436 global_max_workers: self.global_max_workers,
1437 dlq_policy: self.dlq_policy.clone(),
1438 instance_id: self.runtime_instance_id,
1439 started_at: self.runtime_started_at,
1440 hostname: self.runtime_hostname.clone(),
1441 pid: self.runtime_pid,
1442 version: self.runtime_version,
1443 snapshot_interval: self.runtime_snapshot_interval,
1444 effective_storage: self.effective_storage.clone(),
1445 queue_storage_capable: self.storage.queue_storage().is_some(),
1446 transition_role: self.transition_role,
1447 metrics: self.metrics.clone(),
1448 }
1449 }
1450
1451 async fn publish_runtime_snapshot(&self) {
1452 let reporter = self.runtime_reporter_state();
1453 reporter.publish_snapshot().await;
1454 }
1455
1456 async fn log_transition_startup_status(
1457 &self,
1458 effective_storage: &RuntimeStorage,
1459 ) -> Result<(), awa_model::AwaError> {
1460 if self.storage.queue_storage().is_none() {
1461 return Ok(());
1462 }
1463
1464 let report = transition::status_report(&self.pool).await?;
1465 let effective_engine = match effective_storage {
1466 RuntimeStorage::Canonical => "canonical",
1467 RuntimeStorage::QueueStorage(_) => "queue_storage",
1468 };
1469
1470 info!(
1471 transition_role = ?self.transition_role,
1472 state = %report.status.state,
1473 current_engine = %report.status.current_engine,
1474 active_engine = %report.status.active_engine,
1475 prepared_engine = ?report.status.prepared_engine,
1476 effective_engine,
1477 canonical_live_backlog = report.canonical_live_backlog,
1478 "Resolved storage transition state for worker startup"
1479 );
1480
1481 if report.status.state == "prepared" && !report.can_enter_mixed_transition {
1482 warn!(
1483 blockers = %report.enter_mixed_transition_blockers.join("; "),
1484 "Storage transition is prepared but cannot yet enter mixed transition"
1485 );
1486 }
1487
1488 if report.status.state == "mixed_transition" && !report.can_finalize {
1489 warn!(
1490 blockers = %report.finalize_blockers.join("; "),
1491 "Storage transition is in mixed_transition but cannot yet finalize"
1492 );
1493 }
1494
1495 Ok(())
1496 }
1497
1498 pub async fn start(&self) -> Result<(), awa_model::AwaError> {
1500 info!(
1501 queues = self.queues.len(),
1502 workers = self.workers.len(),
1503 "Starting Awa worker runtime"
1504 );
1505
1506 let effective_storage = self.resolve_effective_storage().await?;
1507 {
1508 let mut guard = self.effective_storage.write().await;
1509 *guard = effective_storage.clone();
1510 }
1511
1512 self.log_transition_startup_status(&effective_storage)
1513 .await?;
1514
1515 admin::sync_queue_descriptors(
1516 &self.pool,
1517 &self.declared_queue_descriptors(),
1518 self.runtime_snapshot_interval,
1519 )
1520 .await?;
1521 admin::sync_job_kind_descriptors(
1522 &self.pool,
1523 &self.declared_job_kind_descriptors(),
1524 self.runtime_snapshot_interval,
1525 )
1526 .await?;
1527
1528 let runtime_worker_capacity = self.global_max_workers.unwrap_or_else(|| {
1531 self.queues
1532 .iter()
1533 .map(|(_, config)| config.max_workers)
1534 .sum()
1535 });
1536 let (completion_batcher, completion_handle) = CompletionBatcher::new(
1537 self.pool.clone(),
1538 self.service_cancel.clone(),
1539 self.metrics.clone(),
1540 effective_storage.clone(),
1541 runtime_worker_capacity,
1542 );
1543
1544 let executor = Arc::new(JobExecutor::new(
1546 self.pool.clone(),
1547 self.workers.clone(),
1548 self.lifecycle_handlers.clone(),
1549 self.enqueue_specs.clone(),
1550 self.in_flight.clone(),
1551 self.queue_in_flight.clone(),
1552 self.state.clone(),
1553 self.metrics.clone(),
1554 completion_handle,
1555 effective_storage.clone(),
1556 self.dlq_policy.clone(),
1557 ));
1558
1559 let cancel_listener = crate::cancel_listener::CancelListener::new(
1564 self.pool.clone(),
1565 self.in_flight.clone(),
1566 self.service_cancel.clone(),
1567 );
1568 let cancel_listener_handle = cancel_listener.spawn().await;
1569
1570 let mut service_handles = self.service_handles.write().await;
1571
1572 service_handles.extend(completion_batcher.spawn());
1573 if let Some(handle) = cancel_listener_handle {
1574 service_handles.push(handle);
1575 }
1576
1577 let heartbeat = HeartbeatService::new(
1579 self.pool.clone(),
1580 self.storage.clone(),
1581 self.in_flight.clone(),
1582 self.heartbeat_interval,
1583 self.heartbeat_alive.clone(),
1584 self.service_cancel.clone(),
1585 self.metrics.clone(),
1586 );
1587 service_handles.push(tokio::spawn(async move {
1588 heartbeat.run().await;
1589 }));
1590
1591 let mut maintenance = MaintenanceService::new(
1593 self.pool.clone(),
1594 self.metrics.clone(),
1595 self.leader.clone(),
1596 self.maintenance_alive.clone(),
1597 self.service_cancel.clone(),
1598 self.periodic_jobs.clone(),
1599 self.in_flight.clone(),
1600 effective_storage.clone(),
1601 self.enqueue_specs.clone(),
1602 self.lifecycle_handlers.clone(),
1603 )
1604 .promote_interval(self.promote_interval);
1605 if let Some(interval) = self.heartbeat_rescue_interval {
1606 maintenance = maintenance.heartbeat_rescue_interval(interval);
1607 }
1608 if let Some(staleness) = self.heartbeat_staleness {
1609 maintenance = maintenance.heartbeat_staleness(staleness);
1610 }
1611 if let Some(interval) = self.deadline_rescue_interval {
1612 maintenance = maintenance.deadline_rescue_interval(interval);
1613 }
1614 if let Some(interval) = self.callback_rescue_interval {
1615 maintenance = maintenance.callback_rescue_interval(interval);
1616 }
1617 if let Some(interval) = self.leader_election_interval {
1618 maintenance = maintenance.leader_election_interval(interval);
1619 }
1620 if let Some(interval) = self.leader_check_interval {
1621 maintenance = maintenance.leader_check_interval(interval);
1622 }
1623 if let Some(interval) = self.priority_aging_interval {
1624 maintenance = maintenance.priority_aging_interval(interval);
1625 }
1626 if let Some(interval) = self.terminal_count_rollup_interval {
1627 maintenance = maintenance.terminal_count_rollup_interval(interval);
1628 }
1629 if let Some(retention) = self.completed_retention {
1630 maintenance = maintenance.completed_retention(retention);
1631 }
1632 if let Some(retention) = self.failed_retention {
1633 maintenance = maintenance.failed_retention(retention);
1634 }
1635 if let Some(retention) = self.descriptor_retention {
1636 maintenance = maintenance.descriptor_retention(retention);
1637 }
1638 if let Some(batch_size) = self.cleanup_batch_size {
1639 maintenance = maintenance.cleanup_batch_size(batch_size);
1640 }
1641 if let Some(interval) = self.cleanup_interval {
1642 maintenance = maintenance.cleanup_interval(interval);
1643 }
1644 if !self.queue_retention_overrides.is_empty() {
1645 maintenance =
1646 maintenance.queue_retention_overrides(self.queue_retention_overrides.clone());
1647 }
1648 if let Some(interval) = self.queue_stats_interval {
1649 maintenance = maintenance.queue_stats_interval(interval);
1650 }
1651 if let Some(retention) = self.dlq_retention {
1652 maintenance = maintenance.dlq_retention(retention);
1653 }
1654 if let Some(batch_size) = self.dlq_cleanup_batch_size {
1655 maintenance = maintenance.dlq_cleanup_batch_size(batch_size);
1656 }
1657 maintenance = maintenance.dlq_policy(self.dlq_policy.clone());
1658 service_handles.push(tokio::spawn(async move {
1659 maintenance.run().await;
1660 }));
1661
1662 let mut dispatcher_handles = self.dispatcher_handles.write().await;
1664 for (queue_name, config) in &self.queues {
1665 let alive = self
1666 .dispatcher_alive
1667 .get(queue_name)
1668 .cloned()
1669 .unwrap_or_else(|| Arc::new(AtomicBool::new(false)));
1670 let claimers = usize::from(config.claimers.max(1));
1671 let capacity_wake = Arc::new(tokio::sync::Notify::new());
1672 let rate_limiter = shared_rate_limiter(config);
1673
1674 let hard_reserved = self
1675 .overflow_pool
1676 .is_none()
1677 .then(|| Arc::new(tokio::sync::Semaphore::new(config.max_workers as usize)));
1678 let weighted_local = self
1679 .overflow_pool
1680 .as_ref()
1681 .map(|_| Arc::new(tokio::sync::Semaphore::new(config.min_workers as usize)));
1682
1683 for claimer_idx in 0..claimers {
1684 let concurrency = if let Some(overflow_pool) = &self.overflow_pool {
1685 ConcurrencyMode::Weighted {
1686 local_semaphore: weighted_local
1687 .as_ref()
1688 .expect("weighted local semaphore should exist")
1689 .clone(),
1690 overflow_pool: overflow_pool.clone(),
1691 queue_name: queue_name.clone(),
1692 }
1693 } else {
1694 ConcurrencyMode::HardReserved {
1695 semaphore: hard_reserved
1696 .as_ref()
1697 .expect("hard-reserved semaphore should exist")
1698 .clone(),
1699 }
1700 };
1701 let claimer_owner_id = if claimer_idx == 0 {
1702 self.runtime_instance_id
1703 } else {
1704 Uuid::new_v4()
1708 };
1709 let dispatcher = Dispatcher::with_concurrency(
1710 queue_name.clone(),
1711 self.runtime_instance_id,
1712 config.clone(),
1713 self.pool.clone(),
1714 executor.clone(),
1715 self.metrics.clone(),
1716 self.in_flight.clone(),
1717 alive.clone(),
1718 self.dispatch_cancel.clone(),
1719 self.job_set.clone(),
1720 concurrency,
1721 rate_limiter.clone(),
1722 capacity_wake.clone(),
1723 claimer_owner_id,
1724 effective_storage.clone(),
1725 );
1726 dispatcher_handles.push(tokio::spawn(async move {
1727 dispatcher.run().await;
1728 }));
1729 }
1730 }
1731
1732 self.publish_runtime_snapshot().await;
1733
1734 let reporter = self.runtime_reporter_state();
1735 service_handles.push(tokio::spawn(async move {
1736 reporter.run().await;
1737 }));
1738
1739 info!("Awa worker runtime started");
1740 Ok(())
1741 }
1742
1743 pub async fn shutdown(&self, timeout: Duration) {
1752 info!("Initiating graceful shutdown");
1753
1754 self.dispatch_cancel.cancel();
1756
1757 self.publish_runtime_snapshot().await;
1758
1759 for flag in self.in_flight.flags() {
1761 flag.store(true, Ordering::SeqCst);
1762 }
1763
1764 let dispatcher_handles: Vec<_> = {
1766 let mut guard = self.dispatcher_handles.write().await;
1767 std::mem::take(&mut *guard)
1768 };
1769 for handle in dispatcher_handles {
1770 let _ = handle.await;
1771 }
1772
1773 let drain = async {
1775 let mut set = self.job_set.lock().await;
1776 while set.join_next().await.is_some() {}
1777 };
1778 if tokio::time::timeout(timeout, drain).await.is_err() {
1779 warn!(
1780 timeout_secs = timeout.as_secs(),
1781 "Shutdown drain timeout exceeded, some jobs may not have completed"
1782 );
1783 }
1784
1785 self.service_cancel.cancel();
1787 let service_handles: Vec<_> = {
1788 let mut guard = self.service_handles.write().await;
1789 std::mem::take(&mut *guard)
1790 };
1791 for handle in service_handles {
1792 let _ = handle.await;
1793 }
1794
1795 info!("Awa worker runtime stopped");
1796 }
1797
1798 pub fn pool(&self) -> &PgPool {
1800 &self.pool
1801 }
1802
1803 pub async fn resolve_callback(
1821 &self,
1822 callback_id: Uuid,
1823 payload: Option<serde_json::Value>,
1824 default_action: awa_model::DefaultAction,
1825 run_lease: Option<i64>,
1826 ) -> Result<awa_model::ResolveOutcome, awa_model::AwaError> {
1827 let mut tx = self.pool.begin().await?;
1828 let outcome =
1829 admin::resolve_callback_in_tx(&mut tx, callback_id, payload, default_action, run_lease)
1830 .await?;
1831 let event = match &outcome {
1836 awa_model::ResolveOutcome::Completed { job, .. } => {
1837 self.dispatch_callback_followups_in_tx(
1838 &mut tx,
1839 job,
1840 crate::enqueue_specs::Outcome::Completed,
1841 None,
1842 )
1843 .await?;
1844 Some(UntypedJobEvent::Completed {
1845 job: job.clone(),
1846 duration: Duration::ZERO,
1847 })
1848 }
1849 awa_model::ResolveOutcome::Failed { job } => {
1850 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
1851 error: latest_error_message(job),
1852 attempt: job.attempt,
1853 };
1854 self.dispatch_callback_followups_in_tx(
1855 &mut tx,
1856 job,
1857 crate::enqueue_specs::Outcome::Exhausted,
1858 Some(outcome_ctx),
1859 )
1860 .await?;
1861 Some(UntypedJobEvent::Exhausted {
1862 job: job.clone(),
1863 error: latest_error_message(job),
1864 attempt: job.attempt,
1865 })
1866 }
1867 awa_model::ResolveOutcome::Ignored { .. } => None,
1868 };
1869 tx.commit().await?;
1870 if let Some(event) = event {
1871 self.dispatch_callback_event(event).await;
1872 }
1873 Ok(outcome)
1874 }
1875
1876 pub async fn complete_external(
1882 &self,
1883 callback_id: Uuid,
1884 payload: Option<serde_json::Value>,
1885 run_lease: Option<i64>,
1886 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1887 let mut tx = self.pool.begin().await?;
1888 let job = admin::complete_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
1889 self.dispatch_callback_followups_in_tx(
1890 &mut tx,
1891 &job,
1892 crate::enqueue_specs::Outcome::Completed,
1893 None,
1894 )
1895 .await?;
1896 tx.commit().await?;
1897 self.dispatch_callback_event(UntypedJobEvent::Completed {
1898 job: job.clone(),
1899 duration: Duration::ZERO,
1900 })
1901 .await;
1902 Ok(job)
1903 }
1904
1905 pub async fn fail_external(
1911 &self,
1912 callback_id: Uuid,
1913 error: &str,
1914 run_lease: Option<i64>,
1915 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1916 let mut tx = self.pool.begin().await?;
1917 let job = admin::fail_external_in_tx(&mut tx, callback_id, error, run_lease).await?;
1918 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
1919 error: error.to_string(),
1920 attempt: job.attempt,
1921 };
1922 self.dispatch_callback_followups_in_tx(
1923 &mut tx,
1924 &job,
1925 crate::enqueue_specs::Outcome::Exhausted,
1926 Some(outcome_ctx),
1927 )
1928 .await?;
1929 tx.commit().await?;
1930 self.dispatch_callback_event(UntypedJobEvent::Exhausted {
1931 job: job.clone(),
1932 error: error.to_string(),
1933 attempt: job.attempt,
1934 })
1935 .await;
1936 Ok(job)
1937 }
1938
1939 pub async fn retry_external(
1954 &self,
1955 callback_id: Uuid,
1956 run_lease: Option<i64>,
1957 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1958 let mut tx = self.pool.begin().await?;
1959 let parked_attempt: Option<i16> = sqlx::query_scalar(
1960 "SELECT attempt FROM awa.jobs \
1961 WHERE callback_id = $1 AND state = 'waiting_external' \
1962 AND ($2::bigint IS NULL OR run_lease = $2)",
1963 )
1964 .bind(callback_id)
1965 .bind(run_lease)
1966 .fetch_optional(&mut *tx)
1967 .await?;
1968
1969 let job = admin::retry_external_in_tx(&mut tx, callback_id, run_lease).await?;
1970 let attempt = parked_attempt.unwrap_or(job.attempt);
1971 let error_msg = latest_error_message(&job);
1972 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
1973 error: error_msg.clone(),
1974 attempt,
1975 next_run_at: job.run_at,
1976 };
1977 self.dispatch_callback_followups_in_tx(
1978 &mut tx,
1979 &job,
1980 crate::enqueue_specs::Outcome::Retried,
1981 Some(outcome_ctx),
1982 )
1983 .await?;
1984 tx.commit().await?;
1985 self.dispatch_callback_event(UntypedJobEvent::Retried {
1986 job: job.clone(),
1987 error: error_msg,
1988 attempt,
1989 next_run_at: job.run_at,
1990 })
1991 .await;
1992 Ok(job)
1993 }
1994
1995 async fn dispatch_callback_event(&self, event: UntypedJobEvent) {
1996 let kind = event.job().kind.clone();
1997 crate::executor::dispatch_lifecycle_event(&self.lifecycle_handlers, &kind, event).await;
1998 }
1999
2000 async fn dispatch_callback_followups_in_tx(
2007 &self,
2008 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2009 job: &awa_model::JobRow,
2010 outcome: crate::enqueue_specs::Outcome,
2011 outcome_context: Option<crate::enqueue_specs::OutcomeContext>,
2012 ) -> Result<(), awa_model::AwaError> {
2013 let Some(specs) = self
2014 .enqueue_specs
2015 .get(&outcome)
2016 .and_then(|by_kind| by_kind.get(&job.kind))
2017 .cloned()
2018 else {
2019 return Ok(());
2020 };
2021 if specs.is_empty() {
2022 return Ok(());
2023 }
2024 crate::enqueue_specs::dispatch_specs_in_tx(tx, job, &specs, outcome_context.as_ref()).await
2025 }
2026
2027 pub async fn health_check(&self) -> HealthCheck {
2029 let postgres_connected = sqlx::query("SELECT 1").execute(&self.pool).await.is_ok();
2030 let poll_loop_alive = self
2031 .dispatcher_alive
2032 .values()
2033 .all(|alive| alive.load(Ordering::SeqCst));
2034 let heartbeat_alive = self.heartbeat_alive.load(Ordering::SeqCst);
2035 let maintenance_alive = self.maintenance_alive.load(Ordering::SeqCst);
2036 let shutting_down = self.dispatch_cancel.is_cancelled();
2037 let leader = self.leader.load(Ordering::SeqCst);
2038 let effective_storage = self.effective_storage.read().await.clone();
2039 let available_rows = if let Some(store) = effective_storage.queue_storage_store() {
2040 sqlx::query_as::<_, (String, i64)>(&format!(
2041 r#"
2042 SELECT
2043 enqueues.queue,
2044 COALESCE(
2045 sum(GREATEST(
2046 {}.sequence_next_value(enqueues.seq_name)
2047 - {}.sequence_next_value(claims.seq_name),
2048 0
2049 )),
2050 0
2051 )::bigint AS available
2052 FROM {}.queue_enqueue_heads AS enqueues
2053 JOIN {}.queue_claim_heads AS claims
2054 ON claims.queue = enqueues.queue
2055 AND claims.priority = enqueues.priority
2056 AND claims.enqueue_shard = enqueues.enqueue_shard
2057 GROUP BY enqueues.queue
2058 "#,
2059 store.schema(),
2060 store.schema(),
2061 store.schema(),
2062 store.schema()
2063 ))
2064 .fetch_all(&self.pool)
2065 .await
2066 .unwrap_or_default()
2067 } else {
2068 sqlx::query_as::<_, (String, i64)>(
2069 r#"
2070 SELECT queue, count(*)::bigint AS available
2071 FROM awa.jobs_hot
2072 WHERE state = 'available'
2073 GROUP BY queue
2074 "#,
2075 )
2076 .fetch_all(&self.pool)
2077 .await
2078 .unwrap_or_default()
2079 };
2080 let available_by_queue: HashMap<_, _> = available_rows.into_iter().collect();
2081 let queues = self
2082 .queues
2083 .iter()
2084 .map(|(queue, config)| {
2085 let in_flight = self
2086 .queue_in_flight
2087 .get(queue)
2088 .map(|counter| counter.load(Ordering::SeqCst))
2089 .unwrap_or(0);
2090 let available = available_by_queue.get(queue).copied().unwrap_or(0).max(0) as u64;
2091 let capacity = if let Some(overflow_pool) = &self.overflow_pool {
2092 QueueCapacity::Weighted {
2093 min_workers: config.min_workers,
2094 weight: config.weight,
2095 overflow_held: overflow_pool.held(queue),
2096 }
2097 } else {
2098 QueueCapacity::HardReserved {
2099 max_workers: config.max_workers,
2100 }
2101 };
2102 (
2103 queue.clone(),
2104 QueueHealth {
2105 in_flight,
2106 available,
2107 capacity,
2108 },
2109 )
2110 })
2111 .collect();
2112
2113 HealthCheck {
2114 healthy: postgres_connected
2115 && poll_loop_alive
2116 && heartbeat_alive
2117 && maintenance_alive
2118 && !shutting_down,
2119 postgres_connected,
2120 poll_loop_alive,
2121 heartbeat_alive,
2122 maintenance_alive,
2123 shutting_down,
2124 leader,
2125 queues,
2126 }
2127 }
2128}
2129
2130impl RuntimeReporterState {
2131 async fn storage_capability(&self) -> StorageCapability {
2132 if !self.queue_storage_capable {
2133 return StorageCapability::Canonical;
2134 }
2135
2136 let effective_storage = self.effective_storage.read().await.clone();
2137 if matches!(effective_storage, RuntimeStorage::QueueStorage(_)) {
2138 return StorageCapability::QueueStorage;
2139 }
2140
2141 match transition::status(&self.pool).await {
2142 Ok(status)
2143 if matches!(status.state.as_str(), "mixed_transition" | "active")
2144 && status.active_engine == "queue_storage" =>
2145 {
2146 StorageCapability::CanonicalDrainOnly
2147 }
2148 Ok(_) => StorageCapability::QueueStorage,
2149 Err(err) => {
2150 warn!(
2151 error = %err,
2152 "Failed to resolve storage transition status for runtime snapshot"
2153 );
2154 StorageCapability::QueueStorage
2155 }
2156 }
2157 }
2158
2159 fn queue_descriptor_hashes(&self) -> HashMap<String, String> {
2160 self.declared_queue_descriptors()
2161 .into_iter()
2162 .map(|named| (named.queue, named.descriptor.descriptor_hash()))
2163 .collect()
2164 }
2165
2166 fn job_kind_descriptor_hashes(&self) -> HashMap<String, String> {
2167 self.declared_job_kind_descriptors()
2168 .into_iter()
2169 .map(|named| (named.kind, named.descriptor.descriptor_hash()))
2170 .collect()
2171 }
2172
2173 fn declared_queue_descriptors(&self) -> Vec<NamedQueueDescriptor> {
2174 self.queues
2175 .iter()
2176 .map(|(queue, _)| NamedQueueDescriptor {
2177 queue: queue.clone(),
2178 descriptor: self
2179 .queue_descriptors
2180 .get(queue)
2181 .cloned()
2182 .unwrap_or_default(),
2183 })
2184 .collect()
2185 }
2186
2187 fn declared_job_kind_descriptors(&self) -> Vec<NamedJobKindDescriptor> {
2188 let mut kinds = self.worker_kinds.clone();
2189 for kind in self.job_kind_descriptors.keys() {
2190 if !kinds.iter().any(|existing| existing == kind) {
2191 kinds.push(kind.clone());
2192 }
2193 }
2194 kinds.sort();
2195 kinds.dedup();
2196
2197 kinds
2198 .into_iter()
2199 .map(|kind| NamedJobKindDescriptor {
2200 descriptor: self
2201 .job_kind_descriptors
2202 .get(&kind)
2203 .cloned()
2204 .unwrap_or_default(),
2205 kind,
2206 })
2207 .collect()
2208 }
2209
2210 fn queue_snapshot(&self, queue: &str, config: &QueueConfig) -> QueueRuntimeSnapshot {
2211 let in_flight = self
2212 .queue_in_flight
2213 .get(queue)
2214 .map(|counter| counter.load(Ordering::SeqCst))
2215 .unwrap_or(0);
2216
2217 let (mode, max_workers, min_workers, weight, overflow_held) =
2218 if let Some(overflow_pool) = &self.overflow_pool {
2219 (
2220 QueueRuntimeMode::Weighted,
2221 None,
2222 Some(config.min_workers),
2223 Some(config.weight),
2224 Some(overflow_pool.held(queue)),
2225 )
2226 } else {
2227 (
2228 QueueRuntimeMode::HardReserved,
2229 Some(config.max_workers),
2230 None,
2231 None,
2232 None,
2233 )
2234 };
2235
2236 QueueRuntimeSnapshot {
2237 queue: queue.to_string(),
2238 in_flight,
2239 overflow_held,
2240 config: QueueRuntimeConfigSnapshot {
2241 mode,
2242 max_workers,
2243 min_workers,
2244 weight,
2245 global_max_workers: self.global_max_workers,
2246 poll_interval_ms: config.poll_interval.as_millis() as u64,
2247 deadline_duration_secs: config.deadline_duration.as_secs(),
2248 priority_aging_interval_secs: config.priority_aging_interval.as_secs(),
2249 claimers: Some(config.claimers),
2250 claim_batch_size: Some(config.claim_batch_size),
2251 dlq_enabled: Some(self.dlq_policy.enabled_for(queue)),
2252 rate_limit: config.rate_limit.as_ref().map(|rl| RateLimitSnapshot {
2253 max_rate: rl.max_rate,
2254 burst: rl.burst,
2255 }),
2256 },
2257 }
2258 }
2259
2260 async fn snapshot_input(&self) -> RuntimeSnapshotInput {
2261 let postgres_connected = sqlx::query("SELECT 1").execute(&self.pool).await.is_ok();
2262 let poll_loop_alive = self
2263 .dispatcher_alive
2264 .values()
2265 .all(|alive| alive.load(Ordering::SeqCst));
2266 let heartbeat_alive = self.heartbeat_alive.load(Ordering::SeqCst);
2267 let maintenance_alive = self.maintenance_alive.load(Ordering::SeqCst);
2268 let shutting_down = self.dispatch_cancel.is_cancelled();
2269 let leader = self.leader.load(Ordering::SeqCst);
2270 let healthy = postgres_connected
2271 && poll_loop_alive
2272 && heartbeat_alive
2273 && maintenance_alive
2274 && !shutting_down;
2275 let storage_capability = self.storage_capability().await;
2276 let queues = self
2277 .queues
2278 .iter()
2279 .map(|(queue, config)| self.queue_snapshot(queue, config))
2280 .collect();
2281
2282 RuntimeSnapshotInput {
2283 instance_id: self.instance_id,
2284 hostname: self.hostname.clone(),
2285 pid: self.pid,
2286 version: self.version.to_string(),
2287 storage_capability,
2288 transition_role: TransitionRole::from(self.transition_role),
2289 started_at: self.started_at,
2290 snapshot_interval_ms: self.snapshot_interval.as_millis() as i64,
2291 healthy,
2292 postgres_connected,
2293 poll_loop_alive,
2294 heartbeat_alive,
2295 maintenance_alive,
2296 shutting_down,
2297 leader,
2298 global_max_workers: self.global_max_workers,
2299 queues,
2300 queue_descriptor_hashes: self.queue_descriptor_hashes(),
2301 job_kind_descriptor_hashes: self.job_kind_descriptor_hashes(),
2302 }
2303 }
2304
2305 async fn publish_snapshot(&self) {
2306 let queue_descriptors = self.declared_queue_descriptors();
2307 let kind_descriptors = self.declared_job_kind_descriptors();
2308
2309 if let Err(err) =
2310 admin::sync_queue_descriptors(&self.pool, &queue_descriptors, self.snapshot_interval)
2311 .await
2312 {
2313 warn!(error = %err, "Failed to sync queue descriptors");
2314 }
2315 if let Err(err) =
2316 admin::sync_job_kind_descriptors(&self.pool, &kind_descriptors, self.snapshot_interval)
2317 .await
2318 {
2319 warn!(error = %err, "Failed to sync job kind descriptors");
2320 }
2321
2322 for named in &queue_descriptors {
2327 self.metrics.record_queue_info(
2328 &named.queue,
2329 named.descriptor.display_name.as_deref(),
2330 named.descriptor.description.as_deref(),
2331 named.descriptor.owner.as_deref(),
2332 named.descriptor.docs_url.as_deref(),
2333 &named.descriptor.tags,
2334 );
2335 }
2336 for named in &kind_descriptors {
2337 self.metrics.record_job_kind_info(
2338 &named.kind,
2339 named.descriptor.display_name.as_deref(),
2340 named.descriptor.description.as_deref(),
2341 named.descriptor.owner.as_deref(),
2342 named.descriptor.docs_url.as_deref(),
2343 &named.descriptor.tags,
2344 );
2345 }
2346
2347 let snapshot = self.snapshot_input().await;
2348 if let Err(err) = admin::upsert_runtime_snapshot(&self.pool, &snapshot).await {
2349 warn!(error = %err, "Failed to publish runtime snapshot");
2350 }
2351
2352 if self.queue_storage_capable {
2353 match transition::status_report(&self.pool).await {
2354 Ok(report) => {
2355 self.metrics.record_storage_state(&report.status);
2356 self.metrics.record_storage_transition_ready(
2357 "enter_mixed_transition",
2358 report.can_enter_mixed_transition,
2359 );
2360 self.metrics
2361 .record_storage_transition_ready("finalize", report.can_finalize);
2362 self.metrics
2363 .record_storage_canonical_live_backlog(report.canonical_live_backlog);
2364
2365 for capability in ["canonical", "canonical_drain_only", "queue_storage"] {
2366 let count = report
2367 .live_runtime_capability_counts
2368 .get(capability)
2369 .copied()
2370 .unwrap_or(0) as i64;
2371 self.metrics
2372 .record_storage_live_runtime_capability(capability, count);
2373 }
2374
2375 for (capability, count) in report.live_runtime_capability_counts {
2376 if capability != "canonical"
2377 && capability != "canonical_drain_only"
2378 && capability != "queue_storage"
2379 {
2380 self.metrics
2381 .record_storage_live_runtime_capability(&capability, count as i64);
2382 }
2383 }
2384 }
2385 Err(err) => {
2386 warn!(error = %err, "Failed to publish storage transition metrics");
2387 }
2388 }
2389 }
2390 }
2391
2392 async fn run(self) {
2393 let mut interval = tokio::time::interval(self.snapshot_interval);
2394 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2395 interval.tick().await;
2396 loop {
2397 tokio::select! {
2398 _ = self.dispatch_cancel.cancelled() => {
2399 self.publish_snapshot().await;
2400 break;
2401 }
2402 _ = interval.tick() => {
2403 self.publish_snapshot().await;
2404 }
2405 }
2406 }
2407 }
2408}
2409
2410#[cfg(test)]
2411mod tests {
2412 use super::*;
2413 use awa_model::{migrations, storage, JobArgs, QueueStorage, QueueStorageConfig};
2414 use sqlx::postgres::PgPoolOptions;
2415 use sqlx::PgPool;
2416 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2417 use std::sync::{Arc, Mutex, OnceLock};
2418 use std::time::{Duration, Instant};
2419 use tokio::sync::{oneshot, Notify};
2420
2421 static TEST_MUTEX: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
2422
2423 fn test_mutex() -> &'static tokio::sync::Mutex<()> {
2424 TEST_MUTEX.get_or_init(|| tokio::sync::Mutex::new(()))
2425 }
2426
2427 fn lazy_pool() -> PgPool {
2428 PgPoolOptions::new()
2429 .connect_lazy("postgres://postgres:test@localhost/awa_test")
2430 .expect("lazy pool should build")
2431 }
2432
2433 #[test]
2434 fn queue_config_defaults_use_throughput_oriented_claim_batch() {
2435 let config = QueueConfig::default();
2436 assert_eq!(config.claimers, 1);
2437 assert_eq!(config.claim_batch_size, 512);
2438 }
2439
2440 #[tokio::test]
2441 async fn register_accepts_send_non_sync_handler_future() {
2442 #[derive(serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2443 struct NonSyncFutureJob;
2444
2445 let _client = Client::builder(lazy_pool())
2446 .queue("non_sync_future", QueueConfig::default())
2447 .register::<NonSyncFutureJob, _, _>(|_, _| async move {
2448 let cell = std::cell::Cell::new(0_u8);
2449 tokio::task::yield_now().await;
2450 cell.set(1);
2451 Ok(JobResult::Completed)
2452 })
2453 .build()
2454 .expect("send but non-sync handler future should compile");
2455 }
2456
2457 fn base_database_url() -> String {
2458 std::env::var("DATABASE_URL")
2459 .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string())
2460 }
2461
2462 fn replace_database_name(url: &str, database_name: &str) -> String {
2463 let (without_query, query_suffix) = match url.split_once('?') {
2464 Some((prefix, query)) => (prefix, Some(query)),
2465 None => (url, None),
2466 };
2467 let (base, _) = without_query
2468 .rsplit_once('/')
2469 .expect("database URL should include a database name");
2470 let mut out = format!("{base}/{database_name}");
2471 if let Some(query) = query_suffix {
2472 out.push('?');
2473 out.push_str(query);
2474 }
2475 out
2476 }
2477
2478 fn database_url() -> String {
2479 std::env::var("DATABASE_URL_WORKER_CLIENT").unwrap_or_else(|_| {
2480 replace_database_name(&base_database_url(), "awa_test_worker_client")
2481 })
2482 }
2483
2484 async fn ensure_database_exists(url: &str) {
2485 let database_name = url
2486 .split_once('?')
2487 .map(|(prefix, _)| prefix)
2488 .unwrap_or(url)
2489 .rsplit_once('/')
2490 .map(|(_, database_name)| database_name.to_string())
2491 .expect("database URL should include a database name");
2492 let admin_url = replace_database_name(url, "postgres");
2493 let admin_pool = PgPoolOptions::new()
2494 .max_connections(1)
2495 .connect(&admin_url)
2496 .await
2497 .expect("Failed to connect to admin database for client tests");
2498 let create_sql = format!("CREATE DATABASE {database_name}");
2499 match sqlx::query(&create_sql).execute(&admin_pool).await {
2500 Ok(_) => {}
2501 Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("42P04") => {}
2502 Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => {}
2503 Err(err) => panic!("Failed to create client test database {database_name}: {err}"),
2504 }
2505 }
2506
2507 async fn setup_pool(max_connections: u32) -> PgPool {
2508 let url = database_url();
2509 ensure_database_exists(&url).await;
2510 PgPoolOptions::new()
2511 .max_connections(max_connections)
2512 .acquire_timeout(Duration::from_secs(5))
2513 .connect(&url)
2514 .await
2515 .expect("Failed to connect to client test database")
2516 }
2517
2518 async fn reset_schema(pool: &PgPool) {
2519 sqlx::raw_sql("DROP SCHEMA IF EXISTS awa CASCADE")
2520 .execute(pool)
2521 .await
2522 .expect("Failed to drop awa schema");
2523 }
2524
2525 async fn apply_migrations_through(pool: &PgPool, version: i32) {
2526 for (_version, _desc, sql) in migrations::migration_sql_range(0, version) {
2527 sqlx::raw_sql(&sql).execute(pool).await.unwrap();
2528 }
2529 }
2530
2531 async fn drop_queue_storage_schema(pool: &PgPool, schema: &str) {
2532 let sql = format!("DROP SCHEMA IF EXISTS {schema} CASCADE");
2533 sqlx::query(&sql)
2534 .execute(pool)
2535 .await
2536 .expect("Failed to drop queue storage schema");
2537 }
2538
2539 async fn insert_available_job(pool: &PgPool, kind: &str, queue: &str) -> i64 {
2540 sqlx::query_scalar(
2541 r#"
2542 INSERT INTO awa.jobs (
2543 kind,
2544 queue,
2545 args,
2546 state,
2547 priority,
2548 max_attempts,
2549 run_at,
2550 metadata,
2551 tags
2552 )
2553 VALUES (
2554 $1,
2555 $2,
2556 '{}'::jsonb,
2557 'available'::awa.job_state,
2558 2,
2559 25,
2560 clock_timestamp(),
2561 '{}'::jsonb,
2562 '{}'::text[]
2563 )
2564 RETURNING id
2565 "#,
2566 )
2567 .bind(kind)
2568 .bind(queue)
2569 .fetch_one(pool)
2570 .await
2571 .expect("Failed to insert job")
2572 }
2573
2574 async fn insert_canonical_available_job(pool: &PgPool, kind: &str, queue: &str) -> i64 {
2575 sqlx::query_scalar(
2576 r#"
2577 INSERT INTO awa.jobs_hot (
2578 kind,
2579 queue,
2580 args,
2581 state,
2582 priority,
2583 max_attempts,
2584 run_at,
2585 metadata,
2586 tags
2587 )
2588 VALUES (
2589 $1,
2590 $2,
2591 '{}'::jsonb,
2592 'available'::awa.job_state,
2593 2,
2594 25,
2595 clock_timestamp(),
2596 '{}'::jsonb,
2597 '{}'::text[]
2598 )
2599 RETURNING id
2600 "#,
2601 )
2602 .bind(kind)
2603 .bind(queue)
2604 .fetch_one(pool)
2605 .await
2606 .expect("Failed to insert canonical job")
2607 }
2608
2609 async fn active_queue_storage_schema(pool: &PgPool) -> Option<String> {
2610 sqlx::query_scalar("SELECT awa.active_queue_storage_schema()")
2611 .fetch_one(pool)
2612 .await
2613 .expect("Failed to fetch active queue storage schema")
2614 }
2615
2616 async fn insert_fake_queue_storage_target(pool: &PgPool) {
2621 sqlx::query(
2622 r#"
2623 INSERT INTO awa.runtime_instances (
2624 instance_id, hostname, pid, version,
2625 started_at, last_seen_at, snapshot_interval_ms,
2626 healthy, postgres_connected, poll_loop_alive,
2627 heartbeat_alive, maintenance_alive, shutting_down,
2628 leader, global_max_workers, queues,
2629 storage_capability, transition_role
2630 )
2631 VALUES (
2632 $1, 'fake-target', 7777, '0.6.0-test',
2633 now() - interval '1 minute', now(), 1000,
2634 TRUE, TRUE, TRUE,
2635 TRUE, TRUE, FALSE,
2636 FALSE, NULL, '[]'::jsonb,
2637 'queue_storage', 'queue_storage_target'
2638 )
2639 "#,
2640 )
2641 .bind(Uuid::new_v4())
2642 .execute(pool)
2643 .await
2644 .expect("Failed to insert fake queue_storage_target runtime row");
2645 }
2646
2647 async fn wait_for_runtime_capability(
2648 pool: &PgPool,
2649 instance_id: Uuid,
2650 capability: StorageCapability,
2651 timeout: Duration,
2652 ) {
2653 let start = Instant::now();
2654 loop {
2655 let current: Option<String> = sqlx::query_scalar(
2656 "SELECT storage_capability FROM awa.runtime_instances WHERE instance_id = $1",
2657 )
2658 .bind(instance_id)
2659 .fetch_optional(pool)
2660 .await
2661 .expect("Failed to fetch runtime storage capability");
2662 if current.as_deref() == Some(capability.as_str()) {
2663 return;
2664 }
2665 assert!(
2666 start.elapsed() <= timeout,
2667 "Timed out waiting for runtime {instance_id} to report capability {}; last={current:?}",
2668 capability.as_str()
2669 );
2670 tokio::time::sleep(Duration::from_millis(25)).await;
2671 }
2672 }
2673
2674 async fn expire_runtime_instance(pool: &PgPool, instance_id: Uuid) {
2675 sqlx::query(
2676 "UPDATE awa.runtime_instances SET last_seen_at = now() - interval '1 hour' WHERE instance_id = $1",
2677 )
2678 .bind(instance_id)
2679 .execute(pool)
2680 .await
2681 .expect("Failed to expire runtime instance");
2682 }
2683
2684 async fn prepare_queue_storage_transition_with_config(
2692 pool: &PgPool,
2693 config: QueueStorageConfig,
2694 ) -> QueueStorage {
2695 let schema = config.schema.clone();
2696 let store = QueueStorage::new(config).expect("Failed to build queue storage store");
2697 drop_queue_storage_schema(pool, &schema).await;
2698 store
2699 .prepare_schema(pool)
2700 .await
2701 .expect("Failed to prepare queue storage schema");
2702 storage::prepare(
2703 pool,
2704 "queue_storage",
2705 serde_json::json!({ "schema": schema }),
2706 )
2707 .await
2708 .expect("Failed to prepare queue storage transition");
2709 store
2710 }
2711
2712 async fn wait_for_state(pool: &PgPool, job_id: i64, state: &str, timeout: Duration) {
2713 let start = Instant::now();
2714 loop {
2715 let current: Option<String> = sqlx::query_scalar(
2716 "SELECT state::text FROM awa.jobs_hot WHERE id = $1 UNION ALL SELECT state::text FROM awa.scheduled_jobs WHERE id = $1 LIMIT 1",
2717 )
2718 .bind(job_id)
2719 .fetch_optional(pool)
2720 .await
2721 .expect("Failed to fetch canonical job state");
2722 if current.as_deref() == Some(state) {
2723 return;
2724 }
2725 assert!(
2726 start.elapsed() <= timeout,
2727 "Timed out waiting for job {job_id} to reach state {state}; last_state={current:?}"
2728 );
2729 tokio::time::sleep(Duration::from_millis(25)).await;
2730 }
2731 }
2732
2733 async fn wait_for_queue_storage_done(
2734 pool: &PgPool,
2735 schema: &str,
2736 job_id: i64,
2737 timeout: Duration,
2738 ) {
2739 let sql = format!(
2740 "SELECT EXISTS(SELECT 1 FROM {schema}.terminal_jobs WHERE job_id = $1 AND state = 'completed')"
2741 );
2742 let start = Instant::now();
2743 loop {
2744 let done: bool = sqlx::query_scalar(&sql)
2745 .bind(job_id)
2746 .fetch_one(pool)
2747 .await
2748 .expect("Failed to query queue storage terminal rows");
2749 if done {
2750 return;
2751 }
2752 assert!(
2753 start.elapsed() <= timeout,
2754 "Timed out waiting for queue storage job {job_id} to complete"
2755 );
2756 tokio::time::sleep(Duration::from_millis(25)).await;
2757 }
2758 }
2759
2760 fn force_canonical(mut builder: ClientBuilder) -> ClientBuilder {
2761 builder.storage = RuntimeStorage::Canonical;
2762 builder.storage_error = None;
2763 builder
2764 }
2765
2766 #[tokio::test]
2767 async fn queue_storage_target_requires_prepared_schema() {
2768 let _guard = test_mutex().lock().await;
2769 let pool = setup_pool(4).await;
2770 let queue_storage_schema = "awa_cutover_target_requires_prepare";
2771 reset_schema(&pool).await;
2772 migrations::run(&pool)
2773 .await
2774 .expect("fresh 0.6 schema install should succeed");
2775 drop_queue_storage_schema(&pool, queue_storage_schema).await;
2776
2777 storage::prepare(
2778 &pool,
2779 "queue_storage",
2780 serde_json::json!({ "schema": queue_storage_schema }),
2781 )
2782 .await
2783 .expect("Failed to prepare queue storage transition without schema");
2784
2785 let client = Client::builder(pool.clone())
2786 .queue(
2787 "cutover",
2788 QueueConfig {
2789 max_workers: 1,
2790 poll_interval: Duration::from_millis(25),
2791 ..QueueConfig::default()
2792 },
2793 )
2794 .queue_storage(
2795 QueueStorageConfig {
2796 schema: queue_storage_schema.to_string(),
2797 queue_slot_count: 4,
2798 lease_slot_count: 2,
2799 ..Default::default()
2800 },
2801 Duration::from_millis(1_000),
2802 Duration::from_millis(50),
2803 )
2804 .transition_role(TransitionWorkerRole::QueueStorageTarget)
2805 .register::<CutoverShortJob, _, _>(move |_args, _ctx| async move {
2806 Ok(JobResult::Completed)
2807 })
2808 .build()
2809 .expect("Failed to build queue-storage target client");
2810
2811 let err = client
2812 .start()
2813 .await
2814 .expect_err("queue-storage target should refuse to start without prepared schema");
2815 match err {
2816 awa_model::AwaError::Validation(msg) => {
2817 assert!(
2818 msg.contains("not prepared"),
2819 "unexpected validation message: {msg}"
2820 );
2821 }
2822 other => panic!("expected Validation error, got {other:?}"),
2823 }
2824 }
2825
2826 #[tokio::test]
2827 async fn queue_descriptor_requires_declared_queue() {
2828 let result = Client::builder(lazy_pool())
2829 .queue("default", QueueConfig::default())
2830 .queue_descriptor("billing", QueueDescriptor::new().display_name("Billing"))
2831 .build();
2832
2833 assert!(matches!(
2834 result,
2835 Err(BuildError::QueueDescriptorWithoutQueue { queue }) if queue == "billing"
2836 ));
2837 }
2838
2839 #[tokio::test]
2840 async fn queue_descriptor_allows_declared_queue() {
2841 let result = Client::builder(lazy_pool())
2842 .queue("billing", QueueConfig::default())
2843 .queue_descriptor("billing", QueueDescriptor::new().display_name("Billing"))
2844 .build();
2845
2846 assert!(result.is_ok(), "descriptor for declared queue should build");
2847 }
2848
2849 #[tokio::test]
2850 async fn partitioned_queue_declares_each_physical_queue() {
2851 let partitioned_queue =
2852 PartitionedQueue::new("email", 3).expect("partitioned queue should build");
2853
2854 let client = Client::builder(lazy_pool())
2855 .partitioned_queue(
2856 &partitioned_queue,
2857 QueueConfig {
2858 max_workers: 7,
2859 ..QueueConfig::default()
2860 },
2861 )
2862 .build()
2863 .expect("partitioned queue queues should build");
2864
2865 let queues: Vec<_> = client
2866 .queues
2867 .iter()
2868 .map(|(queue, config)| (queue.as_str(), config.max_workers))
2869 .collect();
2870 assert_eq!(
2871 queues,
2872 vec![("email", 7), ("email__p1", 7), ("email__p2", 7)]
2873 );
2874 }
2875
2876 #[tokio::test]
2877 async fn duplicate_queue_declarations_are_rejected() {
2878 let partitioned_queue =
2879 PartitionedQueue::new("email", 2).expect("partitioned queue should build");
2880
2881 let result = Client::builder(lazy_pool())
2882 .queue("email", QueueConfig::default())
2883 .partitioned_queue(&partitioned_queue, QueueConfig::default())
2884 .build();
2885
2886 assert!(matches!(
2887 result,
2888 Err(BuildError::DuplicateQueue { queue }) if queue == "email"
2889 ));
2890 }
2891
2892 #[tokio::test]
2893 async fn job_kind_descriptor_allows_registered_kind() {
2894 #[derive(serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2895 struct TestJob;
2896
2897 let result = Client::builder(lazy_pool())
2898 .queue("default", QueueConfig::default())
2899 .register::<TestJob, _, _>(|_args, _ctx| async { Ok(JobResult::Completed) })
2900 .job_kind_descriptor::<TestJob>(JobKindDescriptor::new().display_name("Test job"))
2901 .build();
2902
2903 assert!(
2904 result.is_ok(),
2905 "descriptor for registered kind should build"
2906 );
2907 }
2908
2909 #[tokio::test]
2910 async fn dlq_cleanup_batch_size_must_be_positive() {
2911 let result = Client::builder(lazy_pool())
2912 .queue("default", QueueConfig::default())
2913 .dlq_cleanup_batch_size(0)
2914 .build();
2915
2916 assert!(matches!(result, Err(BuildError::InvalidDlqBatchSize)));
2917 }
2918
2919 #[tokio::test]
2920 async fn terminal_count_rollup_interval_must_be_positive() {
2921 let result = Client::builder(lazy_pool())
2922 .queue("default", QueueConfig::default())
2923 .terminal_count_rollup_interval(Duration::ZERO)
2924 .build();
2925
2926 assert!(matches!(
2927 result,
2928 Err(BuildError::InvalidTerminalCountRollupInterval)
2929 ));
2930 }
2931
2932 #[tokio::test]
2933 async fn health_check_reads_available_from_active_queue_storage() {
2934 let _guard = test_mutex().lock().await;
2935 let pool = setup_pool(4).await;
2936 reset_schema(&pool).await;
2937 migrations::run(&pool)
2938 .await
2939 .expect("migrations should succeed");
2940 let queue = "health_queue_storage";
2944 let client = Client::builder(pool.clone())
2945 .queue(queue, QueueConfig::default())
2946 .build()
2947 .expect("queue-storage health client should build");
2948
2949 let store = client
2950 .storage
2951 .queue_storage_store()
2952 .expect("client should default to queue storage");
2953 store
2954 .install(&pool)
2955 .await
2956 .expect("queue storage install should succeed");
2957
2958 insert_available_job(&pool, "cutover_short_job", queue).await;
2959
2960 let health = client.health_check().await;
2961 let queue_health = health
2962 .queues
2963 .get(queue)
2964 .expect("queue should appear in health");
2965 assert_eq!(queue_health.available, 1);
2966 }
2967
2968 #[derive(Clone, serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2969 struct CutoverLongJob {}
2970
2971 #[derive(Clone, serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2972 struct CutoverShortJob {}
2973
2974 #[tokio::test]
2975 async fn fresh_auto_finalize_uses_migrated_default_substrate_without_prepare_schema() {
2976 let _guard = test_mutex().lock().await;
2977 let pool = setup_pool(4).await;
2978 reset_schema(&pool).await;
2979 migrations::run(&pool)
2980 .await
2981 .expect("fresh 0.6 schema install should succeed");
2982
2983 sqlx::raw_sql(
2984 r#"
2985 CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate(
2986 p_schema TEXT,
2987 p_queue_slot_count INT DEFAULT 16,
2988 p_lease_slot_count INT DEFAULT 8,
2989 p_claim_slot_count INT DEFAULT 8,
2990 p_lease_claim_receipts BOOLEAN DEFAULT TRUE
2991 )
2992 RETURNS VOID
2993 LANGUAGE plpgsql
2994 AS $$
2995 BEGIN
2996 RAISE EXCEPTION 'prepare_schema should not run when default queue-storage substrate is already ready'
2997 USING ERRCODE = '55000';
2998 END
2999 $$;
3000 "#,
3001 )
3002 .execute(&pool)
3003 .await
3004 .expect("failed to poison queue-storage helper");
3005
3006 let client = Client::builder(pool.clone())
3007 .queue(
3008 "fresh_auto_finalize",
3009 QueueConfig {
3010 max_workers: 1,
3011 poll_interval: Duration::from_millis(25),
3012 deadline_duration: Duration::ZERO,
3013 ..QueueConfig::default()
3014 },
3015 )
3016 .register::<CutoverShortJob, _, _>(move |_args, _ctx| async move {
3017 Ok(JobResult::Completed)
3018 })
3019 .promote_interval(Duration::from_millis(25))
3020 .leader_election_interval(Duration::from_millis(100))
3021 .leader_check_interval(Duration::from_millis(50))
3022 .runtime_snapshot_interval(Duration::from_millis(100))
3023 .build()
3024 .expect("Failed to build fresh auto-finalize client");
3025
3026 client
3027 .start()
3028 .await
3029 .expect("fresh runtime should auto-finalize without prepare_schema");
3030 assert_eq!(
3031 active_queue_storage_schema(&pool).await.as_deref(),
3032 Some("awa")
3033 );
3034 client.shutdown(Duration::from_secs(5)).await;
3035 }
3036
3037 #[tokio::test]
3038 async fn canonical_runtime_drains_in_flight_jobs_across_schema_upgrade_before_queue_storage_cutover(
3039 ) {
3040 let _guard = test_mutex().lock().await;
3041 let pool = setup_pool(8).await;
3042 let queue_storage_schema = "awa_cutover_runtime";
3043 reset_schema(&pool).await;
3044 drop_queue_storage_schema(&pool, queue_storage_schema).await;
3045 apply_migrations_through(&pool, 9).await;
3046
3047 let long_started_flag = Arc::new(AtomicBool::new(false));
3048 let (long_started_tx_inner, long_started_rx) = oneshot::channel::<()>();
3049 let long_started_tx = Arc::new(Mutex::new(Some(long_started_tx_inner)));
3050 let long_release = Arc::new(Notify::new());
3051 let canonical_short_seen = Arc::new(AtomicUsize::new(0));
3052 let queue_storage_short_seen = Arc::new(AtomicUsize::new(0));
3053
3054 let canonical_client = {
3055 let started = long_started_flag.clone();
3056 let started_tx = long_started_tx.clone();
3057 let release = long_release.clone();
3058 let canonical_short_seen = canonical_short_seen.clone();
3059 let builder = Client::builder(pool.clone())
3060 .queue(
3061 "cutover",
3062 QueueConfig {
3063 max_workers: 2,
3064 poll_interval: Duration::from_millis(25),
3065 ..QueueConfig::default()
3066 },
3067 )
3068 .register::<CutoverLongJob, _, _>(move |_args, _ctx| {
3069 let started = started.clone();
3070 let started_tx = started_tx.clone();
3071 let release = release.clone();
3072 async move {
3073 started.store(true, Ordering::SeqCst);
3074 if let Some(tx) =
3075 started_tx.lock().expect("long-start mutex poisoned").take()
3076 {
3077 let _ = tx.send(());
3078 }
3079 release.notified().await;
3080 Ok(JobResult::Completed)
3081 }
3082 })
3083 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3084 let canonical_short_seen = canonical_short_seen.clone();
3085 async move {
3086 canonical_short_seen.fetch_add(1, Ordering::SeqCst);
3087 Ok(JobResult::Completed)
3088 }
3089 })
3090 .promote_interval(Duration::from_millis(25))
3091 .leader_election_interval(Duration::from_millis(100))
3092 .leader_check_interval(Duration::from_millis(50))
3093 .heartbeat_rescue_interval(Duration::from_millis(100))
3094 .deadline_rescue_interval(Duration::from_millis(100))
3095 .callback_rescue_interval(Duration::from_millis(100));
3096 force_canonical(builder)
3097 .build()
3098 .expect("Failed to build canonical client")
3099 };
3100
3101 canonical_client
3102 .start()
3103 .await
3104 .expect("Failed to start canonical client");
3105
3106 let long_id =
3107 insert_available_job(&pool, <CutoverLongJob as JobArgs>::kind(), "cutover").await;
3108 tokio::time::timeout(Duration::from_secs(5), long_started_rx)
3109 .await
3110 .expect("Timed out waiting for long canonical job to start")
3111 .expect("Long job start signal dropped");
3112 assert!(
3113 long_started_flag.load(Ordering::SeqCst),
3114 "long-running canonical job should be in flight before migration"
3115 );
3116
3117 migrations::run(&pool)
3118 .await
3119 .expect("Schema upgrade from 0.5.x to 0.6 should succeed during canonical runtime");
3120 assert_eq!(
3121 active_queue_storage_schema(&pool).await,
3122 None,
3123 "schema upgrade alone must not activate queue storage"
3124 );
3125
3126 let canonical_short_id =
3127 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3128 let canonical_short_start = Instant::now();
3129 while canonical_short_seen.load(Ordering::SeqCst) == 0 {
3130 assert!(
3131 canonical_short_start.elapsed() <= Duration::from_secs(5),
3132 "canonical worker stopped processing new jobs after schema upgrade"
3133 );
3134 tokio::time::sleep(Duration::from_millis(25)).await;
3135 }
3136 wait_for_state(
3137 &pool,
3138 canonical_short_id,
3139 "completed",
3140 Duration::from_secs(5),
3141 )
3142 .await;
3143
3144 long_release.notify_waiters();
3145 wait_for_state(&pool, long_id, "completed", Duration::from_secs(5)).await;
3146 canonical_client.shutdown(Duration::from_secs(5)).await;
3147 expire_runtime_instance(&pool, canonical_client.runtime_instance_id).await;
3148
3149 let store_config = QueueStorageConfig {
3159 schema: queue_storage_schema.to_string(),
3160 queue_slot_count: 4,
3161 lease_slot_count: 2,
3162 lease_claim_receipts: false,
3163 ..Default::default()
3164 };
3165 let _store =
3166 prepare_queue_storage_transition_with_config(&pool, store_config.clone()).await;
3167 assert_eq!(
3168 active_queue_storage_schema(&pool).await,
3169 None,
3170 "prepare alone must not activate queue storage routing"
3171 );
3172 let drain_only_client = {
3173 let queue_storage_short_seen = queue_storage_short_seen.clone();
3174 Client::builder(pool.clone())
3175 .queue(
3176 "cutover",
3177 QueueConfig {
3178 max_workers: 2,
3179 poll_interval: Duration::from_millis(25),
3180 ..QueueConfig::default()
3181 },
3182 )
3183 .queue_storage(
3184 store_config.clone(),
3185 Duration::from_millis(1_000),
3186 Duration::from_millis(50),
3187 )
3188 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3189 let queue_storage_short_seen = queue_storage_short_seen.clone();
3190 async move {
3191 queue_storage_short_seen.fetch_add(1, Ordering::SeqCst);
3192 Ok(JobResult::Completed)
3193 }
3194 })
3195 .promote_interval(Duration::from_millis(25))
3196 .leader_election_interval(Duration::from_millis(100))
3197 .leader_check_interval(Duration::from_millis(50))
3198 .heartbeat_rescue_interval(Duration::from_millis(100))
3199 .deadline_rescue_interval(Duration::from_millis(100))
3200 .callback_rescue_interval(Duration::from_millis(100))
3201 .runtime_snapshot_interval(Duration::from_millis(100))
3202 .build()
3203 .expect("Failed to build queue storage client")
3204 };
3205
3206 drain_only_client
3207 .start()
3208 .await
3209 .expect("Failed to start queue storage client");
3210 wait_for_runtime_capability(
3211 &pool,
3212 drain_only_client.runtime_instance_id,
3213 StorageCapability::QueueStorage,
3214 Duration::from_secs(5),
3215 )
3216 .await;
3217 assert_eq!(
3218 active_queue_storage_schema(&pool).await,
3219 None,
3220 "prepared queue storage runtime must stay canonical until mixed transition"
3221 );
3222
3223 let prepared_short_id =
3224 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3225 let queue_storage_start = Instant::now();
3226 while queue_storage_short_seen.load(Ordering::SeqCst) == 0 {
3227 assert!(
3228 queue_storage_start.elapsed() <= Duration::from_secs(5),
3229 "queue-storage-capable runtime failed to process canonical work before mixed transition"
3230 );
3231 tokio::time::sleep(Duration::from_millis(25)).await;
3232 }
3233 wait_for_state(
3234 &pool,
3235 prepared_short_id,
3236 "completed",
3237 Duration::from_secs(5),
3238 )
3239 .await;
3240
3241 insert_fake_queue_storage_target(&pool).await;
3248
3249 storage::enter_mixed_transition(&pool)
3250 .await
3251 .expect("enter_mixed_transition should succeed once only 0.6 workers remain");
3252 assert_eq!(
3253 active_queue_storage_schema(&pool).await,
3254 Some(queue_storage_schema.to_string()),
3255 "mixed transition should activate queue storage routing"
3256 );
3257 wait_for_runtime_capability(
3258 &pool,
3259 drain_only_client.runtime_instance_id,
3260 StorageCapability::CanonicalDrainOnly,
3261 Duration::from_secs(5),
3262 )
3263 .await;
3264
3265 let canonical_drain_id =
3266 insert_canonical_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover")
3267 .await;
3268 wait_for_state(
3269 &pool,
3270 canonical_drain_id,
3271 "completed",
3272 Duration::from_secs(5),
3273 )
3274 .await;
3275
3276 drain_only_client.shutdown(Duration::from_secs(5)).await;
3277
3278 let queue_storage_client = {
3279 let queue_storage_short_seen = queue_storage_short_seen.clone();
3280 Client::builder(pool.clone())
3281 .queue(
3282 "cutover",
3283 QueueConfig {
3284 max_workers: 2,
3285 poll_interval: Duration::from_millis(25),
3286 ..QueueConfig::default()
3287 },
3288 )
3289 .queue_storage(
3290 store_config.clone(),
3291 Duration::from_millis(1_000),
3292 Duration::from_millis(50),
3293 )
3294 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3295 let queue_storage_short_seen = queue_storage_short_seen.clone();
3296 async move {
3297 queue_storage_short_seen.fetch_add(1, Ordering::SeqCst);
3298 Ok(JobResult::Completed)
3299 }
3300 })
3301 .promote_interval(Duration::from_millis(25))
3302 .leader_election_interval(Duration::from_millis(100))
3303 .leader_check_interval(Duration::from_millis(50))
3304 .heartbeat_rescue_interval(Duration::from_millis(100))
3305 .deadline_rescue_interval(Duration::from_millis(100))
3306 .callback_rescue_interval(Duration::from_millis(100))
3307 .runtime_snapshot_interval(Duration::from_millis(100))
3308 .build()
3309 .expect("Failed to build post-transition queue storage client")
3310 };
3311
3312 queue_storage_client
3313 .start()
3314 .await
3315 .expect("Failed to start post-transition queue storage client");
3316 wait_for_runtime_capability(
3317 &pool,
3318 queue_storage_client.runtime_instance_id,
3319 StorageCapability::QueueStorage,
3320 Duration::from_secs(5),
3321 )
3322 .await;
3323
3324 let before_queue_storage = queue_storage_short_seen.load(Ordering::SeqCst);
3325 let queue_storage_job_id =
3326 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3327 let queue_storage_start = Instant::now();
3328 while queue_storage_short_seen.load(Ordering::SeqCst) == before_queue_storage {
3329 assert!(
3330 queue_storage_start.elapsed() <= Duration::from_secs(5),
3331 "queue storage runtime failed to process new work after cutover"
3332 );
3333 tokio::time::sleep(Duration::from_millis(25)).await;
3334 }
3335 wait_for_queue_storage_done(
3336 &pool,
3337 queue_storage_schema,
3338 queue_storage_job_id,
3339 Duration::from_secs(5),
3340 )
3341 .await;
3342
3343 queue_storage_client.shutdown(Duration::from_secs(5)).await;
3344 }
3345
3346 #[tokio::test]
3347 async fn queue_storage_target_started_before_mixed_transition_processes_new_work_immediately() {
3348 let _guard = test_mutex().lock().await;
3349 let pool = setup_pool(8).await;
3350 let queue_storage_schema = "awa_cutover_target_runtime";
3351 reset_schema(&pool).await;
3352 migrations::run(&pool)
3353 .await
3354 .expect("fresh 0.6 schema install should succeed");
3355 drop_queue_storage_schema(&pool, queue_storage_schema).await;
3356
3357 let canonical_seen = Arc::new(AtomicUsize::new(0));
3358 let queue_storage_seen = Arc::new(AtomicUsize::new(0));
3359 let store_config = QueueStorageConfig {
3367 schema: queue_storage_schema.to_string(),
3368 queue_slot_count: 4,
3369 lease_slot_count: 2,
3370 lease_claim_receipts: false,
3371 ..Default::default()
3372 };
3373
3374 prepare_queue_storage_transition_with_config(&pool, store_config.clone()).await;
3375 assert_eq!(
3376 active_queue_storage_schema(&pool).await,
3377 None,
3378 "prepare should not activate queue storage routing"
3379 );
3380
3381 let auto_client = {
3382 let canonical_seen = canonical_seen.clone();
3383 Client::builder(pool.clone())
3384 .queue(
3385 "cutover",
3386 QueueConfig {
3387 max_workers: 2,
3388 poll_interval: Duration::from_millis(25),
3389 ..QueueConfig::default()
3390 },
3391 )
3392 .queue_storage(
3393 store_config.clone(),
3394 Duration::from_millis(1_000),
3395 Duration::from_millis(50),
3396 )
3397 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3398 let canonical_seen = canonical_seen.clone();
3399 async move {
3400 canonical_seen.fetch_add(1, Ordering::SeqCst);
3401 Ok(JobResult::Completed)
3402 }
3403 })
3404 .promote_interval(Duration::from_millis(25))
3405 .leader_election_interval(Duration::from_millis(100))
3406 .leader_check_interval(Duration::from_millis(50))
3407 .heartbeat_rescue_interval(Duration::from_millis(100))
3408 .deadline_rescue_interval(Duration::from_millis(100))
3409 .callback_rescue_interval(Duration::from_millis(100))
3410 .runtime_snapshot_interval(Duration::from_millis(100))
3411 .build()
3412 .expect("Failed to build auto cutover client")
3413 };
3414 auto_client
3415 .start()
3416 .await
3417 .expect("Failed to start auto cutover client");
3418
3419 let target_client = {
3420 let queue_storage_seen = queue_storage_seen.clone();
3421 Client::builder(pool.clone())
3422 .queue(
3423 "cutover",
3424 QueueConfig {
3425 max_workers: 2,
3426 poll_interval: Duration::from_millis(25),
3427 ..QueueConfig::default()
3428 },
3429 )
3430 .queue_storage(
3431 store_config.clone(),
3432 Duration::from_millis(1_000),
3433 Duration::from_millis(50),
3434 )
3435 .transition_role(TransitionWorkerRole::QueueStorageTarget)
3436 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3437 let queue_storage_seen = queue_storage_seen.clone();
3438 async move {
3439 queue_storage_seen.fetch_add(1, Ordering::SeqCst);
3440 Ok(JobResult::Completed)
3441 }
3442 })
3443 .promote_interval(Duration::from_millis(25))
3444 .leader_election_interval(Duration::from_millis(100))
3445 .leader_check_interval(Duration::from_millis(50))
3446 .heartbeat_rescue_interval(Duration::from_millis(100))
3447 .deadline_rescue_interval(Duration::from_millis(100))
3448 .callback_rescue_interval(Duration::from_millis(100))
3449 .runtime_snapshot_interval(Duration::from_millis(100))
3450 .build()
3451 .expect("Failed to build queue-storage target client")
3452 };
3453 target_client
3454 .start()
3455 .await
3456 .expect("Failed to start queue-storage target client");
3457
3458 wait_for_runtime_capability(
3459 &pool,
3460 auto_client.runtime_instance_id,
3461 StorageCapability::QueueStorage,
3462 Duration::from_secs(5),
3463 )
3464 .await;
3465 wait_for_runtime_capability(
3466 &pool,
3467 target_client.runtime_instance_id,
3468 StorageCapability::QueueStorage,
3469 Duration::from_secs(5),
3470 )
3471 .await;
3472
3473 let canonical_job_id =
3474 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3475 let canonical_start = Instant::now();
3476 while canonical_seen.load(Ordering::SeqCst) == 0 {
3477 assert!(
3478 canonical_start.elapsed() <= Duration::from_secs(5),
3479 "auto client failed to process canonical work before mixed transition"
3480 );
3481 tokio::time::sleep(Duration::from_millis(25)).await;
3482 }
3483 wait_for_state(&pool, canonical_job_id, "completed", Duration::from_secs(5)).await;
3484 assert_eq!(
3485 queue_storage_seen.load(Ordering::SeqCst),
3486 0,
3487 "queue-storage target should stay idle before routing flips"
3488 );
3489
3490 storage::enter_mixed_transition(&pool)
3491 .await
3492 .expect("enter_mixed_transition should succeed with prepared 0.6 fleet");
3493 wait_for_runtime_capability(
3494 &pool,
3495 auto_client.runtime_instance_id,
3496 StorageCapability::CanonicalDrainOnly,
3497 Duration::from_secs(5),
3498 )
3499 .await;
3500 wait_for_runtime_capability(
3501 &pool,
3502 target_client.runtime_instance_id,
3503 StorageCapability::QueueStorage,
3504 Duration::from_secs(5),
3505 )
3506 .await;
3507 assert_eq!(
3508 active_queue_storage_schema(&pool).await,
3509 Some(queue_storage_schema.to_string()),
3510 "mixed transition should activate queue storage routing"
3511 );
3512
3513 let before_queue_storage = queue_storage_seen.load(Ordering::SeqCst);
3514 let queue_storage_job_id =
3515 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3516 let queue_storage_start = Instant::now();
3517 while queue_storage_seen.load(Ordering::SeqCst) == before_queue_storage {
3518 assert!(
3519 queue_storage_start.elapsed() <= Duration::from_secs(5),
3520 "queue-storage target failed to process new work after routing flip"
3521 );
3522 tokio::time::sleep(Duration::from_millis(25)).await;
3523 }
3524 wait_for_queue_storage_done(
3525 &pool,
3526 queue_storage_schema,
3527 queue_storage_job_id,
3528 Duration::from_secs(5),
3529 )
3530 .await;
3531
3532 let canonical_drain_id =
3533 insert_canonical_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover")
3534 .await;
3535 wait_for_state(
3536 &pool,
3537 canonical_drain_id,
3538 "completed",
3539 Duration::from_secs(5),
3540 )
3541 .await;
3542
3543 target_client.shutdown(Duration::from_secs(5)).await;
3544 auto_client.shutdown(Duration::from_secs(5)).await;
3545 }
3546}