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 .standby_queue_storage(
1605 match (&effective_storage, &self.storage) {
1610 (RuntimeStorage::Canonical, RuntimeStorage::QueueStorage(runtime)) => {
1611 Some(runtime.clone())
1612 }
1613 _ => None,
1614 },
1615 )
1616 .promote_interval(self.promote_interval);
1617 if let Some(interval) = self.heartbeat_rescue_interval {
1618 maintenance = maintenance.heartbeat_rescue_interval(interval);
1619 }
1620 if let Some(staleness) = self.heartbeat_staleness {
1621 maintenance = maintenance.heartbeat_staleness(staleness);
1622 }
1623 if let Some(interval) = self.deadline_rescue_interval {
1624 maintenance = maintenance.deadline_rescue_interval(interval);
1625 }
1626 if let Some(interval) = self.callback_rescue_interval {
1627 maintenance = maintenance.callback_rescue_interval(interval);
1628 }
1629 if let Some(interval) = self.leader_election_interval {
1630 maintenance = maintenance.leader_election_interval(interval);
1631 }
1632 if let Some(interval) = self.leader_check_interval {
1633 maintenance = maintenance.leader_check_interval(interval);
1634 }
1635 if let Some(interval) = self.priority_aging_interval {
1636 maintenance = maintenance.priority_aging_interval(interval);
1637 }
1638 if let Some(interval) = self.terminal_count_rollup_interval {
1639 maintenance = maintenance.terminal_count_rollup_interval(interval);
1640 }
1641 if let Some(retention) = self.completed_retention {
1642 maintenance = maintenance.completed_retention(retention);
1643 }
1644 if let Some(retention) = self.failed_retention {
1645 maintenance = maintenance.failed_retention(retention);
1646 }
1647 if let Some(retention) = self.descriptor_retention {
1648 maintenance = maintenance.descriptor_retention(retention);
1649 }
1650 if let Some(batch_size) = self.cleanup_batch_size {
1651 maintenance = maintenance.cleanup_batch_size(batch_size);
1652 }
1653 if let Some(interval) = self.cleanup_interval {
1654 maintenance = maintenance.cleanup_interval(interval);
1655 }
1656 if !self.queue_retention_overrides.is_empty() {
1657 maintenance =
1658 maintenance.queue_retention_overrides(self.queue_retention_overrides.clone());
1659 }
1660 if let Some(interval) = self.queue_stats_interval {
1661 maintenance = maintenance.queue_stats_interval(interval);
1662 }
1663 if let Some(retention) = self.dlq_retention {
1664 maintenance = maintenance.dlq_retention(retention);
1665 }
1666 if let Some(batch_size) = self.dlq_cleanup_batch_size {
1667 maintenance = maintenance.dlq_cleanup_batch_size(batch_size);
1668 }
1669 maintenance = maintenance.dlq_policy(self.dlq_policy.clone());
1670 service_handles.push(tokio::spawn(async move {
1671 maintenance.run().await;
1672 }));
1673
1674 let mut dispatcher_handles = self.dispatcher_handles.write().await;
1676 for (queue_name, config) in &self.queues {
1677 let alive = self
1678 .dispatcher_alive
1679 .get(queue_name)
1680 .cloned()
1681 .unwrap_or_else(|| Arc::new(AtomicBool::new(false)));
1682 let claimers = usize::from(config.claimers.max(1));
1683 let capacity_wake = Arc::new(tokio::sync::Notify::new());
1684 let rate_limiter = shared_rate_limiter(config);
1685
1686 let hard_reserved = self
1687 .overflow_pool
1688 .is_none()
1689 .then(|| Arc::new(tokio::sync::Semaphore::new(config.max_workers as usize)));
1690 let weighted_local = self
1691 .overflow_pool
1692 .as_ref()
1693 .map(|_| Arc::new(tokio::sync::Semaphore::new(config.min_workers as usize)));
1694
1695 for claimer_idx in 0..claimers {
1696 let concurrency = if let Some(overflow_pool) = &self.overflow_pool {
1697 ConcurrencyMode::Weighted {
1698 local_semaphore: weighted_local
1699 .as_ref()
1700 .expect("weighted local semaphore should exist")
1701 .clone(),
1702 overflow_pool: overflow_pool.clone(),
1703 queue_name: queue_name.clone(),
1704 }
1705 } else {
1706 ConcurrencyMode::HardReserved {
1707 semaphore: hard_reserved
1708 .as_ref()
1709 .expect("hard-reserved semaphore should exist")
1710 .clone(),
1711 }
1712 };
1713 let claimer_owner_id = if claimer_idx == 0 {
1714 self.runtime_instance_id
1715 } else {
1716 Uuid::new_v4()
1720 };
1721 let dispatcher = Dispatcher::with_concurrency(
1722 queue_name.clone(),
1723 self.runtime_instance_id,
1724 config.clone(),
1725 self.pool.clone(),
1726 executor.clone(),
1727 self.metrics.clone(),
1728 self.in_flight.clone(),
1729 alive.clone(),
1730 self.dispatch_cancel.clone(),
1731 self.job_set.clone(),
1732 concurrency,
1733 rate_limiter.clone(),
1734 capacity_wake.clone(),
1735 claimer_owner_id,
1736 effective_storage.clone(),
1737 );
1738 dispatcher_handles.push(tokio::spawn(async move {
1739 dispatcher.run().await;
1740 }));
1741 }
1742 }
1743
1744 self.publish_runtime_snapshot().await;
1745
1746 let reporter = self.runtime_reporter_state();
1747 service_handles.push(tokio::spawn(async move {
1748 reporter.run().await;
1749 }));
1750
1751 info!("Awa worker runtime started");
1752 Ok(())
1753 }
1754
1755 pub async fn shutdown(&self, timeout: Duration) {
1764 info!("Initiating graceful shutdown");
1765
1766 self.dispatch_cancel.cancel();
1768
1769 self.publish_runtime_snapshot().await;
1770
1771 for flag in self.in_flight.flags() {
1773 flag.store(true, Ordering::SeqCst);
1774 }
1775
1776 let dispatcher_handles: Vec<_> = {
1778 let mut guard = self.dispatcher_handles.write().await;
1779 std::mem::take(&mut *guard)
1780 };
1781 for handle in dispatcher_handles {
1782 let _ = handle.await;
1783 }
1784
1785 let drain = async {
1787 let mut set = self.job_set.lock().await;
1788 while set.join_next().await.is_some() {}
1789 };
1790 if tokio::time::timeout(timeout, drain).await.is_err() {
1791 warn!(
1792 timeout_secs = timeout.as_secs(),
1793 "Shutdown drain timeout exceeded, some jobs may not have completed"
1794 );
1795 }
1796
1797 self.service_cancel.cancel();
1799 let service_handles: Vec<_> = {
1800 let mut guard = self.service_handles.write().await;
1801 std::mem::take(&mut *guard)
1802 };
1803 for handle in service_handles {
1804 let _ = handle.await;
1805 }
1806
1807 info!("Awa worker runtime stopped");
1808 }
1809
1810 pub fn pool(&self) -> &PgPool {
1812 &self.pool
1813 }
1814
1815 pub async fn resolve_callback(
1833 &self,
1834 callback_id: Uuid,
1835 payload: Option<serde_json::Value>,
1836 default_action: awa_model::DefaultAction,
1837 run_lease: Option<i64>,
1838 ) -> Result<awa_model::ResolveOutcome, awa_model::AwaError> {
1839 let mut tx = self.pool.begin().await?;
1840 let outcome =
1841 admin::resolve_callback_in_tx(&mut tx, callback_id, payload, default_action, run_lease)
1842 .await?;
1843 let event = match &outcome {
1848 awa_model::ResolveOutcome::Completed { job, .. } => {
1849 self.dispatch_callback_followups_in_tx(
1850 &mut tx,
1851 job,
1852 crate::enqueue_specs::Outcome::Completed,
1853 None,
1854 )
1855 .await?;
1856 Some(UntypedJobEvent::Completed {
1857 job: job.clone(),
1858 duration: Duration::ZERO,
1859 })
1860 }
1861 awa_model::ResolveOutcome::Failed { job } => {
1862 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
1863 error: latest_error_message(job),
1864 attempt: job.attempt,
1865 };
1866 self.dispatch_callback_followups_in_tx(
1867 &mut tx,
1868 job,
1869 crate::enqueue_specs::Outcome::Exhausted,
1870 Some(outcome_ctx),
1871 )
1872 .await?;
1873 Some(UntypedJobEvent::Exhausted {
1874 job: job.clone(),
1875 error: latest_error_message(job),
1876 attempt: job.attempt,
1877 })
1878 }
1879 awa_model::ResolveOutcome::Ignored { .. } => None,
1880 };
1881 tx.commit().await?;
1882 if let Some(event) = event {
1883 self.dispatch_callback_event(event).await;
1884 }
1885 Ok(outcome)
1886 }
1887
1888 pub async fn complete_external(
1894 &self,
1895 callback_id: Uuid,
1896 payload: Option<serde_json::Value>,
1897 run_lease: Option<i64>,
1898 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1899 let mut tx = self.pool.begin().await?;
1900 let job = admin::complete_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
1901 self.dispatch_callback_followups_in_tx(
1902 &mut tx,
1903 &job,
1904 crate::enqueue_specs::Outcome::Completed,
1905 None,
1906 )
1907 .await?;
1908 tx.commit().await?;
1909 self.dispatch_callback_event(UntypedJobEvent::Completed {
1910 job: job.clone(),
1911 duration: Duration::ZERO,
1912 })
1913 .await;
1914 Ok(job)
1915 }
1916
1917 pub async fn fail_external(
1923 &self,
1924 callback_id: Uuid,
1925 error: &str,
1926 run_lease: Option<i64>,
1927 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1928 let mut tx = self.pool.begin().await?;
1929 let job = admin::fail_external_in_tx(&mut tx, callback_id, error, run_lease).await?;
1930 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
1931 error: error.to_string(),
1932 attempt: job.attempt,
1933 };
1934 self.dispatch_callback_followups_in_tx(
1935 &mut tx,
1936 &job,
1937 crate::enqueue_specs::Outcome::Exhausted,
1938 Some(outcome_ctx),
1939 )
1940 .await?;
1941 tx.commit().await?;
1942 self.dispatch_callback_event(UntypedJobEvent::Exhausted {
1943 job: job.clone(),
1944 error: error.to_string(),
1945 attempt: job.attempt,
1946 })
1947 .await;
1948 Ok(job)
1949 }
1950
1951 pub async fn retry_external(
1966 &self,
1967 callback_id: Uuid,
1968 run_lease: Option<i64>,
1969 ) -> Result<awa_model::JobRow, awa_model::AwaError> {
1970 let mut tx = self.pool.begin().await?;
1971 let parked_attempt: Option<i16> = sqlx::query_scalar(
1972 "SELECT attempt FROM awa.jobs \
1973 WHERE callback_id = $1 AND state = 'waiting_external' \
1974 AND ($2::bigint IS NULL OR run_lease = $2)",
1975 )
1976 .bind(callback_id)
1977 .bind(run_lease)
1978 .fetch_optional(&mut *tx)
1979 .await?;
1980
1981 let job = admin::retry_external_in_tx(&mut tx, callback_id, run_lease).await?;
1982 let attempt = parked_attempt.unwrap_or(job.attempt);
1983 let error_msg = latest_error_message(&job);
1984 let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
1985 error: error_msg.clone(),
1986 attempt,
1987 next_run_at: job.run_at,
1988 };
1989 self.dispatch_callback_followups_in_tx(
1990 &mut tx,
1991 &job,
1992 crate::enqueue_specs::Outcome::Retried,
1993 Some(outcome_ctx),
1994 )
1995 .await?;
1996 tx.commit().await?;
1997 self.dispatch_callback_event(UntypedJobEvent::Retried {
1998 job: job.clone(),
1999 error: error_msg,
2000 attempt,
2001 next_run_at: job.run_at,
2002 })
2003 .await;
2004 Ok(job)
2005 }
2006
2007 async fn dispatch_callback_event(&self, event: UntypedJobEvent) {
2008 let kind = event.job().kind.clone();
2009 crate::executor::dispatch_lifecycle_event(&self.lifecycle_handlers, &kind, event).await;
2010 }
2011
2012 async fn dispatch_callback_followups_in_tx(
2019 &self,
2020 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2021 job: &awa_model::JobRow,
2022 outcome: crate::enqueue_specs::Outcome,
2023 outcome_context: Option<crate::enqueue_specs::OutcomeContext>,
2024 ) -> Result<(), awa_model::AwaError> {
2025 let Some(specs) = self
2026 .enqueue_specs
2027 .get(&outcome)
2028 .and_then(|by_kind| by_kind.get(&job.kind))
2029 .cloned()
2030 else {
2031 return Ok(());
2032 };
2033 if specs.is_empty() {
2034 return Ok(());
2035 }
2036 crate::enqueue_specs::dispatch_specs_in_tx(tx, job, &specs, outcome_context.as_ref()).await
2037 }
2038
2039 pub async fn health_check(&self) -> HealthCheck {
2041 let postgres_connected = sqlx::query("SELECT 1").execute(&self.pool).await.is_ok();
2042 let poll_loop_alive = self
2043 .dispatcher_alive
2044 .values()
2045 .all(|alive| alive.load(Ordering::SeqCst));
2046 let heartbeat_alive = self.heartbeat_alive.load(Ordering::SeqCst);
2047 let maintenance_alive = self.maintenance_alive.load(Ordering::SeqCst);
2048 let shutting_down = self.dispatch_cancel.is_cancelled();
2049 let leader = self.leader.load(Ordering::SeqCst);
2050 let effective_storage = self.effective_storage.read().await.clone();
2051 let available_rows = if let Some(store) = effective_storage.queue_storage_store() {
2052 sqlx::query_as::<_, (String, i64)>(&format!(
2053 r#"
2054 SELECT
2055 enqueues.queue,
2056 COALESCE(
2057 sum(GREATEST(
2058 {}.sequence_next_value(enqueues.seq_name)
2059 - {}.sequence_next_value(claims.seq_name),
2060 0
2061 )),
2062 0
2063 )::bigint AS available
2064 FROM {}.queue_enqueue_heads AS enqueues
2065 JOIN {}.queue_claim_heads AS claims
2066 ON claims.queue = enqueues.queue
2067 AND claims.priority = enqueues.priority
2068 AND claims.enqueue_shard = enqueues.enqueue_shard
2069 GROUP BY enqueues.queue
2070 "#,
2071 store.schema(),
2072 store.schema(),
2073 store.schema(),
2074 store.schema()
2075 ))
2076 .fetch_all(&self.pool)
2077 .await
2078 .unwrap_or_default()
2079 } else {
2080 sqlx::query_as::<_, (String, i64)>(
2081 r#"
2082 SELECT queue, count(*)::bigint AS available
2083 FROM awa.jobs_hot
2084 WHERE state = 'available'
2085 GROUP BY queue
2086 "#,
2087 )
2088 .fetch_all(&self.pool)
2089 .await
2090 .unwrap_or_default()
2091 };
2092 let available_by_queue: HashMap<_, _> = available_rows.into_iter().collect();
2093 let queues = self
2094 .queues
2095 .iter()
2096 .map(|(queue, config)| {
2097 let in_flight = self
2098 .queue_in_flight
2099 .get(queue)
2100 .map(|counter| counter.load(Ordering::SeqCst))
2101 .unwrap_or(0);
2102 let available = available_by_queue.get(queue).copied().unwrap_or(0).max(0) as u64;
2103 let capacity = if let Some(overflow_pool) = &self.overflow_pool {
2104 QueueCapacity::Weighted {
2105 min_workers: config.min_workers,
2106 weight: config.weight,
2107 overflow_held: overflow_pool.held(queue),
2108 }
2109 } else {
2110 QueueCapacity::HardReserved {
2111 max_workers: config.max_workers,
2112 }
2113 };
2114 (
2115 queue.clone(),
2116 QueueHealth {
2117 in_flight,
2118 available,
2119 capacity,
2120 },
2121 )
2122 })
2123 .collect();
2124
2125 HealthCheck {
2126 healthy: postgres_connected
2127 && poll_loop_alive
2128 && heartbeat_alive
2129 && maintenance_alive
2130 && !shutting_down,
2131 postgres_connected,
2132 poll_loop_alive,
2133 heartbeat_alive,
2134 maintenance_alive,
2135 shutting_down,
2136 leader,
2137 queues,
2138 }
2139 }
2140}
2141
2142impl RuntimeReporterState {
2143 async fn storage_capability(&self) -> StorageCapability {
2144 if !self.queue_storage_capable {
2145 return StorageCapability::Canonical;
2146 }
2147
2148 let effective_storage = self.effective_storage.read().await.clone();
2149 if matches!(effective_storage, RuntimeStorage::QueueStorage(_)) {
2150 return StorageCapability::QueueStorage;
2151 }
2152
2153 match transition::status(&self.pool).await {
2154 Ok(status)
2155 if matches!(status.state.as_str(), "mixed_transition" | "active")
2156 && status.active_engine == "queue_storage" =>
2157 {
2158 StorageCapability::CanonicalDrainOnly
2159 }
2160 Ok(_) => StorageCapability::QueueStorage,
2161 Err(err) => {
2162 warn!(
2163 error = %err,
2164 "Failed to resolve storage transition status for runtime snapshot"
2165 );
2166 StorageCapability::QueueStorage
2167 }
2168 }
2169 }
2170
2171 fn queue_descriptor_hashes(&self) -> HashMap<String, String> {
2172 self.declared_queue_descriptors()
2173 .into_iter()
2174 .map(|named| (named.queue, named.descriptor.descriptor_hash()))
2175 .collect()
2176 }
2177
2178 fn job_kind_descriptor_hashes(&self) -> HashMap<String, String> {
2179 self.declared_job_kind_descriptors()
2180 .into_iter()
2181 .map(|named| (named.kind, named.descriptor.descriptor_hash()))
2182 .collect()
2183 }
2184
2185 fn declared_queue_descriptors(&self) -> Vec<NamedQueueDescriptor> {
2186 self.queues
2187 .iter()
2188 .map(|(queue, _)| NamedQueueDescriptor {
2189 queue: queue.clone(),
2190 descriptor: self
2191 .queue_descriptors
2192 .get(queue)
2193 .cloned()
2194 .unwrap_or_default(),
2195 })
2196 .collect()
2197 }
2198
2199 fn declared_job_kind_descriptors(&self) -> Vec<NamedJobKindDescriptor> {
2200 let mut kinds = self.worker_kinds.clone();
2201 for kind in self.job_kind_descriptors.keys() {
2202 if !kinds.iter().any(|existing| existing == kind) {
2203 kinds.push(kind.clone());
2204 }
2205 }
2206 kinds.sort();
2207 kinds.dedup();
2208
2209 kinds
2210 .into_iter()
2211 .map(|kind| NamedJobKindDescriptor {
2212 descriptor: self
2213 .job_kind_descriptors
2214 .get(&kind)
2215 .cloned()
2216 .unwrap_or_default(),
2217 kind,
2218 })
2219 .collect()
2220 }
2221
2222 fn queue_snapshot(&self, queue: &str, config: &QueueConfig) -> QueueRuntimeSnapshot {
2223 let in_flight = self
2224 .queue_in_flight
2225 .get(queue)
2226 .map(|counter| counter.load(Ordering::SeqCst))
2227 .unwrap_or(0);
2228
2229 let (mode, max_workers, min_workers, weight, overflow_held) =
2230 if let Some(overflow_pool) = &self.overflow_pool {
2231 (
2232 QueueRuntimeMode::Weighted,
2233 None,
2234 Some(config.min_workers),
2235 Some(config.weight),
2236 Some(overflow_pool.held(queue)),
2237 )
2238 } else {
2239 (
2240 QueueRuntimeMode::HardReserved,
2241 Some(config.max_workers),
2242 None,
2243 None,
2244 None,
2245 )
2246 };
2247
2248 QueueRuntimeSnapshot {
2249 queue: queue.to_string(),
2250 in_flight,
2251 overflow_held,
2252 config: QueueRuntimeConfigSnapshot {
2253 mode,
2254 max_workers,
2255 min_workers,
2256 weight,
2257 global_max_workers: self.global_max_workers,
2258 poll_interval_ms: config.poll_interval.as_millis() as u64,
2259 deadline_duration_secs: config.deadline_duration.as_secs(),
2260 priority_aging_interval_secs: config.priority_aging_interval.as_secs(),
2261 claimers: Some(config.claimers),
2262 claim_batch_size: Some(config.claim_batch_size),
2263 dlq_enabled: Some(self.dlq_policy.enabled_for(queue)),
2264 rate_limit: config.rate_limit.as_ref().map(|rl| RateLimitSnapshot {
2265 max_rate: rl.max_rate,
2266 burst: rl.burst,
2267 }),
2268 },
2269 }
2270 }
2271
2272 async fn snapshot_input(&self) -> RuntimeSnapshotInput {
2273 let postgres_connected = sqlx::query("SELECT 1").execute(&self.pool).await.is_ok();
2274 let poll_loop_alive = self
2275 .dispatcher_alive
2276 .values()
2277 .all(|alive| alive.load(Ordering::SeqCst));
2278 let heartbeat_alive = self.heartbeat_alive.load(Ordering::SeqCst);
2279 let maintenance_alive = self.maintenance_alive.load(Ordering::SeqCst);
2280 let shutting_down = self.dispatch_cancel.is_cancelled();
2281 let leader = self.leader.load(Ordering::SeqCst);
2282 let healthy = postgres_connected
2283 && poll_loop_alive
2284 && heartbeat_alive
2285 && maintenance_alive
2286 && !shutting_down;
2287 let storage_capability = self.storage_capability().await;
2288 let queues = self
2289 .queues
2290 .iter()
2291 .map(|(queue, config)| self.queue_snapshot(queue, config))
2292 .collect();
2293
2294 RuntimeSnapshotInput {
2295 instance_id: self.instance_id,
2296 hostname: self.hostname.clone(),
2297 pid: self.pid,
2298 version: self.version.to_string(),
2299 storage_capability,
2300 transition_role: TransitionRole::from(self.transition_role),
2301 started_at: self.started_at,
2302 snapshot_interval_ms: self.snapshot_interval.as_millis() as i64,
2303 healthy,
2304 postgres_connected,
2305 poll_loop_alive,
2306 heartbeat_alive,
2307 maintenance_alive,
2308 shutting_down,
2309 leader,
2310 global_max_workers: self.global_max_workers,
2311 queues,
2312 queue_descriptor_hashes: self.queue_descriptor_hashes(),
2313 job_kind_descriptor_hashes: self.job_kind_descriptor_hashes(),
2314 }
2315 }
2316
2317 async fn publish_snapshot(&self) {
2318 let queue_descriptors = self.declared_queue_descriptors();
2319 let kind_descriptors = self.declared_job_kind_descriptors();
2320
2321 if let Err(err) =
2322 admin::sync_queue_descriptors(&self.pool, &queue_descriptors, self.snapshot_interval)
2323 .await
2324 {
2325 warn!(error = %err, "Failed to sync queue descriptors");
2326 }
2327 if let Err(err) =
2328 admin::sync_job_kind_descriptors(&self.pool, &kind_descriptors, self.snapshot_interval)
2329 .await
2330 {
2331 warn!(error = %err, "Failed to sync job kind descriptors");
2332 }
2333
2334 for named in &queue_descriptors {
2339 self.metrics.record_queue_info(
2340 &named.queue,
2341 named.descriptor.display_name.as_deref(),
2342 named.descriptor.description.as_deref(),
2343 named.descriptor.owner.as_deref(),
2344 named.descriptor.docs_url.as_deref(),
2345 &named.descriptor.tags,
2346 );
2347 }
2348 for named in &kind_descriptors {
2349 self.metrics.record_job_kind_info(
2350 &named.kind,
2351 named.descriptor.display_name.as_deref(),
2352 named.descriptor.description.as_deref(),
2353 named.descriptor.owner.as_deref(),
2354 named.descriptor.docs_url.as_deref(),
2355 &named.descriptor.tags,
2356 );
2357 }
2358
2359 let snapshot = self.snapshot_input().await;
2360 if let Err(err) = admin::upsert_runtime_snapshot(&self.pool, &snapshot).await {
2361 warn!(error = %err, "Failed to publish runtime snapshot");
2362 }
2363
2364 if self.queue_storage_capable {
2365 match transition::status_report(&self.pool).await {
2366 Ok(report) => {
2367 self.metrics.record_storage_state(&report.status);
2368 self.metrics.record_storage_transition_ready(
2369 "enter_mixed_transition",
2370 report.can_enter_mixed_transition,
2371 );
2372 self.metrics
2373 .record_storage_transition_ready("finalize", report.can_finalize);
2374 self.metrics
2375 .record_storage_canonical_live_backlog(report.canonical_live_backlog);
2376
2377 for capability in ["canonical", "canonical_drain_only", "queue_storage"] {
2378 let count = report
2379 .live_runtime_capability_counts
2380 .get(capability)
2381 .copied()
2382 .unwrap_or(0) as i64;
2383 self.metrics
2384 .record_storage_live_runtime_capability(capability, count);
2385 }
2386
2387 for (capability, count) in report.live_runtime_capability_counts {
2388 if capability != "canonical"
2389 && capability != "canonical_drain_only"
2390 && capability != "queue_storage"
2391 {
2392 self.metrics
2393 .record_storage_live_runtime_capability(&capability, count as i64);
2394 }
2395 }
2396 }
2397 Err(err) => {
2398 warn!(error = %err, "Failed to publish storage transition metrics");
2399 }
2400 }
2401 }
2402 }
2403
2404 async fn run(self) {
2405 let mut interval = tokio::time::interval(self.snapshot_interval);
2406 interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
2407 interval.tick().await;
2408 loop {
2409 tokio::select! {
2410 _ = self.dispatch_cancel.cancelled() => {
2411 self.publish_snapshot().await;
2412 break;
2413 }
2414 _ = interval.tick() => {
2415 self.publish_snapshot().await;
2416 }
2417 }
2418 }
2419 }
2420}
2421
2422#[cfg(test)]
2423mod tests {
2424 use super::*;
2425 use awa_model::{migrations, storage, JobArgs, QueueStorage, QueueStorageConfig};
2426 use sqlx::postgres::PgPoolOptions;
2427 use sqlx::PgPool;
2428 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2429 use std::sync::{Arc, Mutex, OnceLock};
2430 use std::time::{Duration, Instant};
2431 use tokio::sync::{oneshot, Notify};
2432
2433 static TEST_MUTEX: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
2434
2435 fn test_mutex() -> &'static tokio::sync::Mutex<()> {
2436 TEST_MUTEX.get_or_init(|| tokio::sync::Mutex::new(()))
2437 }
2438
2439 fn lazy_pool() -> PgPool {
2440 PgPoolOptions::new()
2441 .connect_lazy("postgres://postgres:test@localhost/awa_test")
2442 .expect("lazy pool should build")
2443 }
2444
2445 #[test]
2446 fn queue_config_defaults_use_throughput_oriented_claim_batch() {
2447 let config = QueueConfig::default();
2448 assert_eq!(config.claimers, 1);
2449 assert_eq!(config.claim_batch_size, 512);
2450 }
2451
2452 #[tokio::test]
2453 async fn register_accepts_send_non_sync_handler_future() {
2454 #[derive(serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2455 struct NonSyncFutureJob;
2456
2457 let _client = Client::builder(lazy_pool())
2458 .queue("non_sync_future", QueueConfig::default())
2459 .register::<NonSyncFutureJob, _, _>(|_, _| async move {
2460 let cell = std::cell::Cell::new(0_u8);
2461 tokio::task::yield_now().await;
2462 cell.set(1);
2463 Ok(JobResult::Completed)
2464 })
2465 .build()
2466 .expect("send but non-sync handler future should compile");
2467 }
2468
2469 fn base_database_url() -> String {
2470 std::env::var("DATABASE_URL")
2471 .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string())
2472 }
2473
2474 fn replace_database_name(url: &str, database_name: &str) -> String {
2475 let (without_query, query_suffix) = match url.split_once('?') {
2476 Some((prefix, query)) => (prefix, Some(query)),
2477 None => (url, None),
2478 };
2479 let (base, _) = without_query
2480 .rsplit_once('/')
2481 .expect("database URL should include a database name");
2482 let mut out = format!("{base}/{database_name}");
2483 if let Some(query) = query_suffix {
2484 out.push('?');
2485 out.push_str(query);
2486 }
2487 out
2488 }
2489
2490 fn database_url() -> String {
2491 std::env::var("DATABASE_URL_WORKER_CLIENT").unwrap_or_else(|_| {
2492 replace_database_name(&base_database_url(), "awa_test_worker_client")
2493 })
2494 }
2495
2496 async fn ensure_database_exists(url: &str) {
2497 let database_name = url
2498 .split_once('?')
2499 .map(|(prefix, _)| prefix)
2500 .unwrap_or(url)
2501 .rsplit_once('/')
2502 .map(|(_, database_name)| database_name.to_string())
2503 .expect("database URL should include a database name");
2504 let admin_url = replace_database_name(url, "postgres");
2505 let admin_pool = PgPoolOptions::new()
2506 .max_connections(1)
2507 .connect(&admin_url)
2508 .await
2509 .expect("Failed to connect to admin database for client tests");
2510 let create_sql = format!("CREATE DATABASE {database_name}");
2511 match sqlx::query(&create_sql).execute(&admin_pool).await {
2512 Ok(_) => {}
2513 Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("42P04") => {}
2514 Err(sqlx::Error::Database(db_err)) if db_err.code().as_deref() == Some("23505") => {}
2515 Err(err) => panic!("Failed to create client test database {database_name}: {err}"),
2516 }
2517 }
2518
2519 async fn setup_pool(max_connections: u32) -> PgPool {
2520 let url = database_url();
2521 ensure_database_exists(&url).await;
2522 PgPoolOptions::new()
2523 .max_connections(max_connections)
2524 .acquire_timeout(Duration::from_secs(5))
2525 .connect(&url)
2526 .await
2527 .expect("Failed to connect to client test database")
2528 }
2529
2530 async fn reset_schema(pool: &PgPool) {
2531 sqlx::raw_sql("DROP SCHEMA IF EXISTS awa CASCADE")
2532 .execute(pool)
2533 .await
2534 .expect("Failed to drop awa schema");
2535 }
2536
2537 async fn apply_migrations_through(pool: &PgPool, version: i32) {
2538 for (_version, _desc, sql) in migrations::migration_sql_range(0, version) {
2539 sqlx::raw_sql(&sql).execute(pool).await.unwrap();
2540 }
2541 }
2542
2543 async fn drop_queue_storage_schema(pool: &PgPool, schema: &str) {
2544 let sql = format!("DROP SCHEMA IF EXISTS {schema} CASCADE");
2545 sqlx::query(&sql)
2546 .execute(pool)
2547 .await
2548 .expect("Failed to drop queue storage schema");
2549 }
2550
2551 async fn insert_available_job(pool: &PgPool, kind: &str, queue: &str) -> i64 {
2552 sqlx::query_scalar(
2553 r#"
2554 INSERT INTO awa.jobs (
2555 kind,
2556 queue,
2557 args,
2558 state,
2559 priority,
2560 max_attempts,
2561 run_at,
2562 metadata,
2563 tags
2564 )
2565 VALUES (
2566 $1,
2567 $2,
2568 '{}'::jsonb,
2569 'available'::awa.job_state,
2570 2,
2571 25,
2572 clock_timestamp(),
2573 '{}'::jsonb,
2574 '{}'::text[]
2575 )
2576 RETURNING id
2577 "#,
2578 )
2579 .bind(kind)
2580 .bind(queue)
2581 .fetch_one(pool)
2582 .await
2583 .expect("Failed to insert job")
2584 }
2585
2586 async fn insert_canonical_available_job(pool: &PgPool, kind: &str, queue: &str) -> i64 {
2587 sqlx::query_scalar(
2588 r#"
2589 INSERT INTO awa.jobs_hot (
2590 kind,
2591 queue,
2592 args,
2593 state,
2594 priority,
2595 max_attempts,
2596 run_at,
2597 metadata,
2598 tags
2599 )
2600 VALUES (
2601 $1,
2602 $2,
2603 '{}'::jsonb,
2604 'available'::awa.job_state,
2605 2,
2606 25,
2607 clock_timestamp(),
2608 '{}'::jsonb,
2609 '{}'::text[]
2610 )
2611 RETURNING id
2612 "#,
2613 )
2614 .bind(kind)
2615 .bind(queue)
2616 .fetch_one(pool)
2617 .await
2618 .expect("Failed to insert canonical job")
2619 }
2620
2621 async fn active_queue_storage_schema(pool: &PgPool) -> Option<String> {
2622 sqlx::query_scalar("SELECT awa.active_queue_storage_schema()")
2623 .fetch_one(pool)
2624 .await
2625 .expect("Failed to fetch active queue storage schema")
2626 }
2627
2628 async fn insert_fake_queue_storage_target(pool: &PgPool) {
2633 sqlx::query(
2634 r#"
2635 INSERT INTO awa.runtime_instances (
2636 instance_id, hostname, pid, version,
2637 started_at, last_seen_at, snapshot_interval_ms,
2638 healthy, postgres_connected, poll_loop_alive,
2639 heartbeat_alive, maintenance_alive, shutting_down,
2640 leader, global_max_workers, queues,
2641 storage_capability, transition_role
2642 )
2643 VALUES (
2644 $1, 'fake-target', 7777, '0.6.0-test',
2645 now() - interval '1 minute', now(), 1000,
2646 TRUE, TRUE, TRUE,
2647 TRUE, TRUE, FALSE,
2648 FALSE, NULL, '[]'::jsonb,
2649 'queue_storage', 'queue_storage_target'
2650 )
2651 "#,
2652 )
2653 .bind(Uuid::new_v4())
2654 .execute(pool)
2655 .await
2656 .expect("Failed to insert fake queue_storage_target runtime row");
2657 }
2658
2659 async fn wait_for_runtime_capability(
2660 pool: &PgPool,
2661 instance_id: Uuid,
2662 capability: StorageCapability,
2663 timeout: Duration,
2664 ) {
2665 let start = Instant::now();
2666 loop {
2667 let current: Option<String> = sqlx::query_scalar(
2668 "SELECT storage_capability FROM awa.runtime_instances WHERE instance_id = $1",
2669 )
2670 .bind(instance_id)
2671 .fetch_optional(pool)
2672 .await
2673 .expect("Failed to fetch runtime storage capability");
2674 if current.as_deref() == Some(capability.as_str()) {
2675 return;
2676 }
2677 assert!(
2678 start.elapsed() <= timeout,
2679 "Timed out waiting for runtime {instance_id} to report capability {}; last={current:?}",
2680 capability.as_str()
2681 );
2682 tokio::time::sleep(Duration::from_millis(25)).await;
2683 }
2684 }
2685
2686 async fn expire_runtime_instance(pool: &PgPool, instance_id: Uuid) {
2687 sqlx::query(
2688 "UPDATE awa.runtime_instances SET last_seen_at = now() - interval '1 hour' WHERE instance_id = $1",
2689 )
2690 .bind(instance_id)
2691 .execute(pool)
2692 .await
2693 .expect("Failed to expire runtime instance");
2694 }
2695
2696 async fn prepare_queue_storage_transition_with_config(
2704 pool: &PgPool,
2705 config: QueueStorageConfig,
2706 ) -> QueueStorage {
2707 let schema = config.schema.clone();
2708 let store = QueueStorage::new(config).expect("Failed to build queue storage store");
2709 drop_queue_storage_schema(pool, &schema).await;
2710 store
2711 .prepare_schema(pool)
2712 .await
2713 .expect("Failed to prepare queue storage schema");
2714 storage::prepare(
2715 pool,
2716 "queue_storage",
2717 serde_json::json!({ "schema": schema }),
2718 )
2719 .await
2720 .expect("Failed to prepare queue storage transition");
2721 store
2722 }
2723
2724 async fn wait_for_state(pool: &PgPool, job_id: i64, state: &str, timeout: Duration) {
2725 let start = Instant::now();
2726 loop {
2727 let current: Option<String> = sqlx::query_scalar(
2728 "SELECT state::text FROM awa.jobs_hot WHERE id = $1 UNION ALL SELECT state::text FROM awa.scheduled_jobs WHERE id = $1 LIMIT 1",
2729 )
2730 .bind(job_id)
2731 .fetch_optional(pool)
2732 .await
2733 .expect("Failed to fetch canonical job state");
2734 if current.as_deref() == Some(state) {
2735 return;
2736 }
2737 assert!(
2738 start.elapsed() <= timeout,
2739 "Timed out waiting for job {job_id} to reach state {state}; last_state={current:?}"
2740 );
2741 tokio::time::sleep(Duration::from_millis(25)).await;
2742 }
2743 }
2744
2745 async fn wait_for_queue_storage_done(
2746 pool: &PgPool,
2747 schema: &str,
2748 job_id: i64,
2749 timeout: Duration,
2750 ) {
2751 let sql = format!(
2752 "SELECT EXISTS(SELECT 1 FROM {schema}.terminal_jobs WHERE job_id = $1 AND state = 'completed')"
2753 );
2754 let start = Instant::now();
2755 loop {
2756 let done: bool = sqlx::query_scalar(&sql)
2757 .bind(job_id)
2758 .fetch_one(pool)
2759 .await
2760 .expect("Failed to query queue storage terminal rows");
2761 if done {
2762 return;
2763 }
2764 assert!(
2765 start.elapsed() <= timeout,
2766 "Timed out waiting for queue storage job {job_id} to complete"
2767 );
2768 tokio::time::sleep(Duration::from_millis(25)).await;
2769 }
2770 }
2771
2772 fn force_canonical(mut builder: ClientBuilder) -> ClientBuilder {
2773 builder.storage = RuntimeStorage::Canonical;
2774 builder.storage_error = None;
2775 builder
2776 }
2777
2778 #[tokio::test]
2779 async fn queue_storage_target_requires_prepared_schema() {
2780 let _guard = test_mutex().lock().await;
2781 let pool = setup_pool(4).await;
2782 let queue_storage_schema = "awa_cutover_target_requires_prepare";
2783 reset_schema(&pool).await;
2784 migrations::run(&pool)
2785 .await
2786 .expect("fresh 0.6 schema install should succeed");
2787 drop_queue_storage_schema(&pool, queue_storage_schema).await;
2788
2789 storage::prepare(
2790 &pool,
2791 "queue_storage",
2792 serde_json::json!({ "schema": queue_storage_schema }),
2793 )
2794 .await
2795 .expect("Failed to prepare queue storage transition without schema");
2796
2797 let client = Client::builder(pool.clone())
2798 .queue(
2799 "cutover",
2800 QueueConfig {
2801 max_workers: 1,
2802 poll_interval: Duration::from_millis(25),
2803 ..QueueConfig::default()
2804 },
2805 )
2806 .queue_storage(
2807 QueueStorageConfig {
2808 schema: queue_storage_schema.to_string(),
2809 queue_slot_count: 4,
2810 lease_slot_count: 2,
2811 ..Default::default()
2812 },
2813 Duration::from_millis(1_000),
2814 Duration::from_millis(50),
2815 )
2816 .transition_role(TransitionWorkerRole::QueueStorageTarget)
2817 .register::<CutoverShortJob, _, _>(move |_args, _ctx| async move {
2818 Ok(JobResult::Completed)
2819 })
2820 .build()
2821 .expect("Failed to build queue-storage target client");
2822
2823 let err = client
2824 .start()
2825 .await
2826 .expect_err("queue-storage target should refuse to start without prepared schema");
2827 match err {
2828 awa_model::AwaError::Validation(msg) => {
2829 assert!(
2830 msg.contains("not prepared"),
2831 "unexpected validation message: {msg}"
2832 );
2833 }
2834 other => panic!("expected Validation error, got {other:?}"),
2835 }
2836 }
2837
2838 #[tokio::test]
2839 async fn queue_descriptor_requires_declared_queue() {
2840 let result = Client::builder(lazy_pool())
2841 .queue("default", QueueConfig::default())
2842 .queue_descriptor("billing", QueueDescriptor::new().display_name("Billing"))
2843 .build();
2844
2845 assert!(matches!(
2846 result,
2847 Err(BuildError::QueueDescriptorWithoutQueue { queue }) if queue == "billing"
2848 ));
2849 }
2850
2851 #[tokio::test]
2852 async fn queue_descriptor_allows_declared_queue() {
2853 let result = Client::builder(lazy_pool())
2854 .queue("billing", QueueConfig::default())
2855 .queue_descriptor("billing", QueueDescriptor::new().display_name("Billing"))
2856 .build();
2857
2858 assert!(result.is_ok(), "descriptor for declared queue should build");
2859 }
2860
2861 #[tokio::test]
2862 async fn partitioned_queue_declares_each_physical_queue() {
2863 let partitioned_queue =
2864 PartitionedQueue::new("email", 3).expect("partitioned queue should build");
2865
2866 let client = Client::builder(lazy_pool())
2867 .partitioned_queue(
2868 &partitioned_queue,
2869 QueueConfig {
2870 max_workers: 7,
2871 ..QueueConfig::default()
2872 },
2873 )
2874 .build()
2875 .expect("partitioned queue queues should build");
2876
2877 let queues: Vec<_> = client
2878 .queues
2879 .iter()
2880 .map(|(queue, config)| (queue.as_str(), config.max_workers))
2881 .collect();
2882 assert_eq!(
2883 queues,
2884 vec![("email", 7), ("email__p1", 7), ("email__p2", 7)]
2885 );
2886 }
2887
2888 #[tokio::test]
2889 async fn duplicate_queue_declarations_are_rejected() {
2890 let partitioned_queue =
2891 PartitionedQueue::new("email", 2).expect("partitioned queue should build");
2892
2893 let result = Client::builder(lazy_pool())
2894 .queue("email", QueueConfig::default())
2895 .partitioned_queue(&partitioned_queue, QueueConfig::default())
2896 .build();
2897
2898 assert!(matches!(
2899 result,
2900 Err(BuildError::DuplicateQueue { queue }) if queue == "email"
2901 ));
2902 }
2903
2904 #[tokio::test]
2905 async fn job_kind_descriptor_allows_registered_kind() {
2906 #[derive(serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2907 struct TestJob;
2908
2909 let result = Client::builder(lazy_pool())
2910 .queue("default", QueueConfig::default())
2911 .register::<TestJob, _, _>(|_args, _ctx| async { Ok(JobResult::Completed) })
2912 .job_kind_descriptor::<TestJob>(JobKindDescriptor::new().display_name("Test job"))
2913 .build();
2914
2915 assert!(
2916 result.is_ok(),
2917 "descriptor for registered kind should build"
2918 );
2919 }
2920
2921 #[tokio::test]
2922 async fn dlq_cleanup_batch_size_must_be_positive() {
2923 let result = Client::builder(lazy_pool())
2924 .queue("default", QueueConfig::default())
2925 .dlq_cleanup_batch_size(0)
2926 .build();
2927
2928 assert!(matches!(result, Err(BuildError::InvalidDlqBatchSize)));
2929 }
2930
2931 #[tokio::test]
2932 async fn terminal_count_rollup_interval_must_be_positive() {
2933 let result = Client::builder(lazy_pool())
2934 .queue("default", QueueConfig::default())
2935 .terminal_count_rollup_interval(Duration::ZERO)
2936 .build();
2937
2938 assert!(matches!(
2939 result,
2940 Err(BuildError::InvalidTerminalCountRollupInterval)
2941 ));
2942 }
2943
2944 #[tokio::test]
2945 async fn health_check_reads_available_from_active_queue_storage() {
2946 let _guard = test_mutex().lock().await;
2947 let pool = setup_pool(4).await;
2948 reset_schema(&pool).await;
2949 migrations::run(&pool)
2950 .await
2951 .expect("migrations should succeed");
2952 let queue = "health_queue_storage";
2956 let client = Client::builder(pool.clone())
2957 .queue(queue, QueueConfig::default())
2958 .build()
2959 .expect("queue-storage health client should build");
2960
2961 let store = client
2962 .storage
2963 .queue_storage_store()
2964 .expect("client should default to queue storage");
2965 store
2966 .install(&pool)
2967 .await
2968 .expect("queue storage install should succeed");
2969
2970 insert_available_job(&pool, "cutover_short_job", queue).await;
2971
2972 let health = client.health_check().await;
2973 let queue_health = health
2974 .queues
2975 .get(queue)
2976 .expect("queue should appear in health");
2977 assert_eq!(queue_health.available, 1);
2978 }
2979
2980 #[derive(Clone, serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2981 struct CutoverLongJob {}
2982
2983 #[derive(Clone, serde::Serialize, serde::Deserialize, awa_macros::JobArgs)]
2984 struct CutoverShortJob {}
2985
2986 #[tokio::test]
2987 async fn fresh_auto_finalize_uses_migrated_default_substrate_without_prepare_schema() {
2988 let _guard = test_mutex().lock().await;
2989 let pool = setup_pool(4).await;
2990 reset_schema(&pool).await;
2991 migrations::run(&pool)
2992 .await
2993 .expect("fresh 0.6 schema install should succeed");
2994
2995 sqlx::raw_sql(
2996 r#"
2997 CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate(
2998 p_schema TEXT,
2999 p_queue_slot_count INT DEFAULT 16,
3000 p_lease_slot_count INT DEFAULT 8,
3001 p_claim_slot_count INT DEFAULT 8,
3002 p_lease_claim_receipts BOOLEAN DEFAULT TRUE
3003 )
3004 RETURNS VOID
3005 LANGUAGE plpgsql
3006 AS $$
3007 BEGIN
3008 RAISE EXCEPTION 'prepare_schema should not run when default queue-storage substrate is already ready'
3009 USING ERRCODE = '55000';
3010 END
3011 $$;
3012 "#,
3013 )
3014 .execute(&pool)
3015 .await
3016 .expect("failed to poison queue-storage helper");
3017
3018 let client = Client::builder(pool.clone())
3019 .queue(
3020 "fresh_auto_finalize",
3021 QueueConfig {
3022 max_workers: 1,
3023 poll_interval: Duration::from_millis(25),
3024 deadline_duration: Duration::ZERO,
3025 ..QueueConfig::default()
3026 },
3027 )
3028 .register::<CutoverShortJob, _, _>(move |_args, _ctx| async move {
3029 Ok(JobResult::Completed)
3030 })
3031 .promote_interval(Duration::from_millis(25))
3032 .leader_election_interval(Duration::from_millis(100))
3033 .leader_check_interval(Duration::from_millis(50))
3034 .runtime_snapshot_interval(Duration::from_millis(100))
3035 .build()
3036 .expect("Failed to build fresh auto-finalize client");
3037
3038 client
3039 .start()
3040 .await
3041 .expect("fresh runtime should auto-finalize without prepare_schema");
3042 assert_eq!(
3043 active_queue_storage_schema(&pool).await.as_deref(),
3044 Some("awa")
3045 );
3046 client.shutdown(Duration::from_secs(5)).await;
3047 }
3048
3049 #[tokio::test]
3050 async fn canonical_runtime_drains_in_flight_jobs_across_schema_upgrade_before_queue_storage_cutover(
3051 ) {
3052 let _guard = test_mutex().lock().await;
3053 let pool = setup_pool(8).await;
3054 let queue_storage_schema = "awa_cutover_runtime";
3055 reset_schema(&pool).await;
3056 drop_queue_storage_schema(&pool, queue_storage_schema).await;
3057 apply_migrations_through(&pool, 9).await;
3058
3059 let long_started_flag = Arc::new(AtomicBool::new(false));
3060 let (long_started_tx_inner, long_started_rx) = oneshot::channel::<()>();
3061 let long_started_tx = Arc::new(Mutex::new(Some(long_started_tx_inner)));
3062 let long_release = Arc::new(Notify::new());
3063 let canonical_short_seen = Arc::new(AtomicUsize::new(0));
3064 let queue_storage_short_seen = Arc::new(AtomicUsize::new(0));
3065
3066 let canonical_client = {
3067 let started = long_started_flag.clone();
3068 let started_tx = long_started_tx.clone();
3069 let release = long_release.clone();
3070 let canonical_short_seen = canonical_short_seen.clone();
3071 let builder = Client::builder(pool.clone())
3072 .queue(
3073 "cutover",
3074 QueueConfig {
3075 max_workers: 2,
3076 poll_interval: Duration::from_millis(25),
3077 ..QueueConfig::default()
3078 },
3079 )
3080 .register::<CutoverLongJob, _, _>(move |_args, _ctx| {
3081 let started = started.clone();
3082 let started_tx = started_tx.clone();
3083 let release = release.clone();
3084 async move {
3085 started.store(true, Ordering::SeqCst);
3086 if let Some(tx) =
3087 started_tx.lock().expect("long-start mutex poisoned").take()
3088 {
3089 let _ = tx.send(());
3090 }
3091 release.notified().await;
3092 Ok(JobResult::Completed)
3093 }
3094 })
3095 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3096 let canonical_short_seen = canonical_short_seen.clone();
3097 async move {
3098 canonical_short_seen.fetch_add(1, Ordering::SeqCst);
3099 Ok(JobResult::Completed)
3100 }
3101 })
3102 .promote_interval(Duration::from_millis(25))
3103 .leader_election_interval(Duration::from_millis(100))
3104 .leader_check_interval(Duration::from_millis(50))
3105 .heartbeat_rescue_interval(Duration::from_millis(100))
3106 .deadline_rescue_interval(Duration::from_millis(100))
3107 .callback_rescue_interval(Duration::from_millis(100));
3108 force_canonical(builder)
3109 .build()
3110 .expect("Failed to build canonical client")
3111 };
3112
3113 canonical_client
3114 .start()
3115 .await
3116 .expect("Failed to start canonical client");
3117
3118 let long_id =
3119 insert_available_job(&pool, <CutoverLongJob as JobArgs>::kind(), "cutover").await;
3120 tokio::time::timeout(Duration::from_secs(5), long_started_rx)
3121 .await
3122 .expect("Timed out waiting for long canonical job to start")
3123 .expect("Long job start signal dropped");
3124 assert!(
3125 long_started_flag.load(Ordering::SeqCst),
3126 "long-running canonical job should be in flight before migration"
3127 );
3128
3129 migrations::run(&pool)
3130 .await
3131 .expect("Schema upgrade from 0.5.x to 0.6 should succeed during canonical runtime");
3132 assert_eq!(
3133 active_queue_storage_schema(&pool).await,
3134 None,
3135 "schema upgrade alone must not activate queue storage"
3136 );
3137
3138 let canonical_short_id =
3139 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3140 let canonical_short_start = Instant::now();
3141 while canonical_short_seen.load(Ordering::SeqCst) == 0 {
3142 assert!(
3143 canonical_short_start.elapsed() <= Duration::from_secs(5),
3144 "canonical worker stopped processing new jobs after schema upgrade"
3145 );
3146 tokio::time::sleep(Duration::from_millis(25)).await;
3147 }
3148 wait_for_state(
3149 &pool,
3150 canonical_short_id,
3151 "completed",
3152 Duration::from_secs(5),
3153 )
3154 .await;
3155
3156 long_release.notify_waiters();
3157 wait_for_state(&pool, long_id, "completed", Duration::from_secs(5)).await;
3158 canonical_client.shutdown(Duration::from_secs(5)).await;
3159 expire_runtime_instance(&pool, canonical_client.runtime_instance_id).await;
3160
3161 let store_config = QueueStorageConfig {
3171 schema: queue_storage_schema.to_string(),
3172 queue_slot_count: 4,
3173 lease_slot_count: 2,
3174 lease_claim_receipts: false,
3175 ..Default::default()
3176 };
3177 let _store =
3178 prepare_queue_storage_transition_with_config(&pool, store_config.clone()).await;
3179 assert_eq!(
3180 active_queue_storage_schema(&pool).await,
3181 None,
3182 "prepare alone must not activate queue storage routing"
3183 );
3184 let drain_only_client = {
3185 let queue_storage_short_seen = queue_storage_short_seen.clone();
3186 Client::builder(pool.clone())
3187 .queue(
3188 "cutover",
3189 QueueConfig {
3190 max_workers: 2,
3191 poll_interval: Duration::from_millis(25),
3192 ..QueueConfig::default()
3193 },
3194 )
3195 .queue_storage(
3196 store_config.clone(),
3197 Duration::from_millis(1_000),
3198 Duration::from_millis(50),
3199 )
3200 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3201 let queue_storage_short_seen = queue_storage_short_seen.clone();
3202 async move {
3203 queue_storage_short_seen.fetch_add(1, Ordering::SeqCst);
3204 Ok(JobResult::Completed)
3205 }
3206 })
3207 .promote_interval(Duration::from_millis(25))
3208 .leader_election_interval(Duration::from_millis(100))
3209 .leader_check_interval(Duration::from_millis(50))
3210 .heartbeat_rescue_interval(Duration::from_millis(100))
3211 .deadline_rescue_interval(Duration::from_millis(100))
3212 .callback_rescue_interval(Duration::from_millis(100))
3213 .runtime_snapshot_interval(Duration::from_millis(100))
3214 .build()
3215 .expect("Failed to build queue storage client")
3216 };
3217
3218 drain_only_client
3219 .start()
3220 .await
3221 .expect("Failed to start queue storage client");
3222 wait_for_runtime_capability(
3223 &pool,
3224 drain_only_client.runtime_instance_id,
3225 StorageCapability::QueueStorage,
3226 Duration::from_secs(5),
3227 )
3228 .await;
3229 assert_eq!(
3230 active_queue_storage_schema(&pool).await,
3231 None,
3232 "prepared queue storage runtime must stay canonical until mixed transition"
3233 );
3234
3235 let prepared_short_id =
3236 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3237 let queue_storage_start = Instant::now();
3238 while queue_storage_short_seen.load(Ordering::SeqCst) == 0 {
3239 assert!(
3240 queue_storage_start.elapsed() <= Duration::from_secs(5),
3241 "queue-storage-capable runtime failed to process canonical work before mixed transition"
3242 );
3243 tokio::time::sleep(Duration::from_millis(25)).await;
3244 }
3245 wait_for_state(
3246 &pool,
3247 prepared_short_id,
3248 "completed",
3249 Duration::from_secs(5),
3250 )
3251 .await;
3252
3253 insert_fake_queue_storage_target(&pool).await;
3260
3261 storage::enter_mixed_transition(&pool)
3262 .await
3263 .expect("enter_mixed_transition should succeed once only 0.6 workers remain");
3264 assert_eq!(
3265 active_queue_storage_schema(&pool).await,
3266 Some(queue_storage_schema.to_string()),
3267 "mixed transition should activate queue storage routing"
3268 );
3269 wait_for_runtime_capability(
3270 &pool,
3271 drain_only_client.runtime_instance_id,
3272 StorageCapability::CanonicalDrainOnly,
3273 Duration::from_secs(5),
3274 )
3275 .await;
3276
3277 let canonical_drain_id =
3278 insert_canonical_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover")
3279 .await;
3280 wait_for_state(
3281 &pool,
3282 canonical_drain_id,
3283 "completed",
3284 Duration::from_secs(5),
3285 )
3286 .await;
3287
3288 drain_only_client.shutdown(Duration::from_secs(5)).await;
3289
3290 let queue_storage_client = {
3291 let queue_storage_short_seen = queue_storage_short_seen.clone();
3292 Client::builder(pool.clone())
3293 .queue(
3294 "cutover",
3295 QueueConfig {
3296 max_workers: 2,
3297 poll_interval: Duration::from_millis(25),
3298 ..QueueConfig::default()
3299 },
3300 )
3301 .queue_storage(
3302 store_config.clone(),
3303 Duration::from_millis(1_000),
3304 Duration::from_millis(50),
3305 )
3306 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3307 let queue_storage_short_seen = queue_storage_short_seen.clone();
3308 async move {
3309 queue_storage_short_seen.fetch_add(1, Ordering::SeqCst);
3310 Ok(JobResult::Completed)
3311 }
3312 })
3313 .promote_interval(Duration::from_millis(25))
3314 .leader_election_interval(Duration::from_millis(100))
3315 .leader_check_interval(Duration::from_millis(50))
3316 .heartbeat_rescue_interval(Duration::from_millis(100))
3317 .deadline_rescue_interval(Duration::from_millis(100))
3318 .callback_rescue_interval(Duration::from_millis(100))
3319 .runtime_snapshot_interval(Duration::from_millis(100))
3320 .build()
3321 .expect("Failed to build post-transition queue storage client")
3322 };
3323
3324 queue_storage_client
3325 .start()
3326 .await
3327 .expect("Failed to start post-transition queue storage client");
3328 wait_for_runtime_capability(
3329 &pool,
3330 queue_storage_client.runtime_instance_id,
3331 StorageCapability::QueueStorage,
3332 Duration::from_secs(5),
3333 )
3334 .await;
3335
3336 let before_queue_storage = queue_storage_short_seen.load(Ordering::SeqCst);
3337 let queue_storage_job_id =
3338 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3339 let queue_storage_start = Instant::now();
3340 while queue_storage_short_seen.load(Ordering::SeqCst) == before_queue_storage {
3341 assert!(
3342 queue_storage_start.elapsed() <= Duration::from_secs(5),
3343 "queue storage runtime failed to process new work after cutover"
3344 );
3345 tokio::time::sleep(Duration::from_millis(25)).await;
3346 }
3347 wait_for_queue_storage_done(
3348 &pool,
3349 queue_storage_schema,
3350 queue_storage_job_id,
3351 Duration::from_secs(5),
3352 )
3353 .await;
3354
3355 queue_storage_client.shutdown(Duration::from_secs(5)).await;
3356 }
3357
3358 #[tokio::test]
3359 async fn queue_storage_target_started_before_mixed_transition_processes_new_work_immediately() {
3360 let _guard = test_mutex().lock().await;
3361 let pool = setup_pool(8).await;
3362 let queue_storage_schema = "awa_cutover_target_runtime";
3363 reset_schema(&pool).await;
3364 migrations::run(&pool)
3365 .await
3366 .expect("fresh 0.6 schema install should succeed");
3367 drop_queue_storage_schema(&pool, queue_storage_schema).await;
3368
3369 let canonical_seen = Arc::new(AtomicUsize::new(0));
3370 let queue_storage_seen = Arc::new(AtomicUsize::new(0));
3371 let store_config = QueueStorageConfig {
3379 schema: queue_storage_schema.to_string(),
3380 queue_slot_count: 4,
3381 lease_slot_count: 2,
3382 lease_claim_receipts: false,
3383 ..Default::default()
3384 };
3385
3386 prepare_queue_storage_transition_with_config(&pool, store_config.clone()).await;
3387 assert_eq!(
3388 active_queue_storage_schema(&pool).await,
3389 None,
3390 "prepare should not activate queue storage routing"
3391 );
3392
3393 let auto_client = {
3394 let canonical_seen = canonical_seen.clone();
3395 Client::builder(pool.clone())
3396 .queue(
3397 "cutover",
3398 QueueConfig {
3399 max_workers: 2,
3400 poll_interval: Duration::from_millis(25),
3401 ..QueueConfig::default()
3402 },
3403 )
3404 .queue_storage(
3405 store_config.clone(),
3406 Duration::from_millis(1_000),
3407 Duration::from_millis(50),
3408 )
3409 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3410 let canonical_seen = canonical_seen.clone();
3411 async move {
3412 canonical_seen.fetch_add(1, Ordering::SeqCst);
3413 Ok(JobResult::Completed)
3414 }
3415 })
3416 .promote_interval(Duration::from_millis(25))
3417 .leader_election_interval(Duration::from_millis(100))
3418 .leader_check_interval(Duration::from_millis(50))
3419 .heartbeat_rescue_interval(Duration::from_millis(100))
3420 .deadline_rescue_interval(Duration::from_millis(100))
3421 .callback_rescue_interval(Duration::from_millis(100))
3422 .runtime_snapshot_interval(Duration::from_millis(100))
3423 .build()
3424 .expect("Failed to build auto cutover client")
3425 };
3426 auto_client
3427 .start()
3428 .await
3429 .expect("Failed to start auto cutover client");
3430
3431 let target_client = {
3432 let queue_storage_seen = queue_storage_seen.clone();
3433 Client::builder(pool.clone())
3434 .queue(
3435 "cutover",
3436 QueueConfig {
3437 max_workers: 2,
3438 poll_interval: Duration::from_millis(25),
3439 ..QueueConfig::default()
3440 },
3441 )
3442 .queue_storage(
3443 store_config.clone(),
3444 Duration::from_millis(1_000),
3445 Duration::from_millis(50),
3446 )
3447 .transition_role(TransitionWorkerRole::QueueStorageTarget)
3448 .register::<CutoverShortJob, _, _>(move |_args, _ctx| {
3449 let queue_storage_seen = queue_storage_seen.clone();
3450 async move {
3451 queue_storage_seen.fetch_add(1, Ordering::SeqCst);
3452 Ok(JobResult::Completed)
3453 }
3454 })
3455 .promote_interval(Duration::from_millis(25))
3456 .leader_election_interval(Duration::from_millis(100))
3457 .leader_check_interval(Duration::from_millis(50))
3458 .heartbeat_rescue_interval(Duration::from_millis(100))
3459 .deadline_rescue_interval(Duration::from_millis(100))
3460 .callback_rescue_interval(Duration::from_millis(100))
3461 .runtime_snapshot_interval(Duration::from_millis(100))
3462 .build()
3463 .expect("Failed to build queue-storage target client")
3464 };
3465 target_client
3466 .start()
3467 .await
3468 .expect("Failed to start queue-storage target client");
3469
3470 wait_for_runtime_capability(
3471 &pool,
3472 auto_client.runtime_instance_id,
3473 StorageCapability::QueueStorage,
3474 Duration::from_secs(5),
3475 )
3476 .await;
3477 wait_for_runtime_capability(
3478 &pool,
3479 target_client.runtime_instance_id,
3480 StorageCapability::QueueStorage,
3481 Duration::from_secs(5),
3482 )
3483 .await;
3484
3485 let canonical_job_id =
3486 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3487 let canonical_start = Instant::now();
3488 while canonical_seen.load(Ordering::SeqCst) == 0 {
3489 assert!(
3490 canonical_start.elapsed() <= Duration::from_secs(5),
3491 "auto client failed to process canonical work before mixed transition"
3492 );
3493 tokio::time::sleep(Duration::from_millis(25)).await;
3494 }
3495 wait_for_state(&pool, canonical_job_id, "completed", Duration::from_secs(5)).await;
3496 assert_eq!(
3497 queue_storage_seen.load(Ordering::SeqCst),
3498 0,
3499 "queue-storage target should stay idle before routing flips"
3500 );
3501
3502 storage::enter_mixed_transition(&pool)
3503 .await
3504 .expect("enter_mixed_transition should succeed with prepared 0.6 fleet");
3505 wait_for_runtime_capability(
3506 &pool,
3507 auto_client.runtime_instance_id,
3508 StorageCapability::CanonicalDrainOnly,
3509 Duration::from_secs(5),
3510 )
3511 .await;
3512 wait_for_runtime_capability(
3513 &pool,
3514 target_client.runtime_instance_id,
3515 StorageCapability::QueueStorage,
3516 Duration::from_secs(5),
3517 )
3518 .await;
3519 assert_eq!(
3520 active_queue_storage_schema(&pool).await,
3521 Some(queue_storage_schema.to_string()),
3522 "mixed transition should activate queue storage routing"
3523 );
3524
3525 let before_queue_storage = queue_storage_seen.load(Ordering::SeqCst);
3526 let queue_storage_job_id =
3527 insert_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover").await;
3528 let queue_storage_start = Instant::now();
3529 while queue_storage_seen.load(Ordering::SeqCst) == before_queue_storage {
3530 assert!(
3531 queue_storage_start.elapsed() <= Duration::from_secs(5),
3532 "queue-storage target failed to process new work after routing flip"
3533 );
3534 tokio::time::sleep(Duration::from_millis(25)).await;
3535 }
3536 wait_for_queue_storage_done(
3537 &pool,
3538 queue_storage_schema,
3539 queue_storage_job_id,
3540 Duration::from_secs(5),
3541 )
3542 .await;
3543
3544 let canonical_drain_id =
3545 insert_canonical_available_job(&pool, <CutoverShortJob as JobArgs>::kind(), "cutover")
3546 .await;
3547 wait_for_state(
3548 &pool,
3549 canonical_drain_id,
3550 "completed",
3551 Duration::from_secs(5),
3552 )
3553 .await;
3554
3555 target_client.shutdown(Duration::from_secs(5)).await;
3556 auto_client.shutdown(Duration::from_secs(5)).await;
3557 }
3558}