Skip to main content

awa_worker/
client.rs

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/// Errors returned when building a worker client.
34#[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/// Health check result.
63#[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/// Per-queue health.
76#[derive(Debug, Clone)]
77pub struct QueueHealth {
78    pub in_flight: u32,
79    pub available: u64,
80    /// Capacity interpretation depends on mode.
81    pub capacity: QueueCapacity,
82}
83
84/// Capacity information for a queue, mode-dependent.
85#[derive(Debug, Clone)]
86pub enum QueueCapacity {
87    /// Hard-reserved: fixed max.
88    HardReserved { max_workers: u32 },
89    /// Weighted: min guaranteed + current overflow.
90    Weighted {
91        min_workers: u32,
92        weight: u32,
93        overflow_held: u32,
94    },
95}
96
97/// Temporary execution role used during a storage transition.
98///
99/// This is not intended as a long-term “run either backend forever” feature.
100/// It exists so a `0.6` rollout can keep some runtimes draining canonical
101/// backlog while other runtimes are already prepared to execute queue-storage
102/// work as soon as routing flips.
103#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
104pub enum TransitionWorkerRole {
105    /// Follow `awa.storage_status()`:
106    /// canonical in `canonical` / `prepared`,
107    /// queue storage in `mixed_transition` / `active`.
108    #[default]
109    Auto,
110    /// Stay on canonical execution even after routing flips.
111    CanonicalDrain,
112    /// Run queue storage immediately, even before routing flips.
113    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
126/// Builder for the Awa worker client.
127pub 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        // #169: keep lease rotation on the same conservative cadence as
171        // queue / claim rotation. Faster lease rotation only advances a
172        // one-row metadata cursor more often; under a pinned MVCC horizon
173        // that metadata churn can dominate the otherwise append-only queue
174        // path without improving job safety.
175        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    /// Add a queue with its configuration.
227    pub fn queue(mut self, name: impl Into<String>, config: QueueConfig) -> Self {
228        self.queues.push((name.into(), config));
229        self
230    }
231
232    /// Add every physical queue in a logical partitioned queue.
233    ///
234    /// This is equivalent to calling [`queue`] once for each physical queue in
235    /// the partitioned queue. Producers should route inserts through the same
236    /// [`PartitionedQueue`] so workers and producers agree on the physical queue
237    /// names.
238    ///
239    /// The `config` is applied to each physical queue. In hard-reserved mode,
240    /// total logical capacity is therefore roughly
241    /// `partitioned_queue.partitions() * config.max_workers`; per-queue rate
242    /// limits also apply per physical
243    /// queue. Divide those knobs yourself, or use weighted mode with a
244    /// `global_max_workers` cap, when you need a logical total.
245    ///
246    /// [`queue`]: ClientBuilder::queue
247    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    /// Attach descriptive metadata (display name, description, owner,
259    /// docs URL, tags, extra JSON) to a queue so it appears labelled in
260    /// the admin API and UI. The queue must also be declared via
261    /// [`queue`]; otherwise [`build`] fails with
262    /// [`BuildError::QueueDescriptorWithoutQueue`].
263    ///
264    /// [`queue`]: ClientBuilder::queue
265    /// [`build`]: ClientBuilder::build
266    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    /// Attach descriptive metadata to a typed job kind. The kind string is
276    /// taken from [`JobArgs::kind`] on `T`.
277    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    /// Attach descriptive metadata to a job kind by string name. Useful
284    /// when the kind is known dynamically (e.g. from language bridges).
285    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    /// Register a typed worker.
295    ///
296    /// The worker handles jobs of type `T` where `T: JobArgs + DeserializeOwned`.
297    /// The handler function receives the deserialized args and job context.
298    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    /// Register a typed lifecycle event handler for a job kind.
315    ///
316    /// `Started` dispatch is scheduled after the claim commits and before the
317    /// worker handler is invoked. Outcome handlers run only after the
318    /// corresponding DB state transition commits. Hooks are best-effort
319    /// in-process notifications, not a durable workflow mechanism; because
320    /// they run in detached tasks, a very short job can complete before a
321    /// `Started` handler finishes. Capture any shared dependencies you need in
322    /// the closure environment.
323    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    /// Register an untyped lifecycle event handler for a specific job kind.
358    ///
359    /// Use this with `register_worker(...)` or for cross-cutting logic that
360    /// doesn't need typed args. Timing matches `on_event(...)`.
361    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
382    /// completes successfully.
383    ///
384    /// `make` receives the trigger's deserialised args plus its post-completion
385    /// [`awa_model::JobRow`] and returns the follow-up's `JobArgs` value. The
386    /// follow-up is enqueued with default [`awa_model::InsertOpts`] — use
387    /// [`Self::on_completed_enqueue_with`] to override queue, priority, or
388    /// other insert options.
389    ///
390    /// # Atomicity
391    ///
392    /// Whether the trigger completes from the worker handler
393    /// (`Ok(Completed)`) or via callback resolution
394    /// ([`Client::complete_external`], [`Client::resolve_callback`] with
395    /// a `Complete` action), the follow-up `INSERT`s in the *same
396    /// database transaction* as the completion's state change. The
397    /// trigger and the follow-up commit or roll back together (ADR-029,
398    /// ADR-013 run-lease guard). A spec INSERT failure or a panic in
399    /// this closure rolls the completion back as well, so a failed
400    /// callback can be redelivered by the external sender rather than
401    /// leaving the job half-applied.
402    ///
403    /// This is the durable counterpart to [`Self::on_event`]: hooks are
404    /// best-effort, in-process; this enqueue is at-least-once and rides
405    /// Awa's existing retry/DLQ machinery.
406    ///
407    /// Multiple registrations for the same kind stack and all run in the
408    /// same dispatch.
409    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    /// Like [`Self::on_completed_enqueue`] but lets the closure return an
421    /// [`EnqueueRequest`](crate::EnqueueRequest) with explicit queue,
422    /// priority, or other [`awa_model::InsertOpts`] overrides.
423    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
448    /// is retried (whether via `Ok(RetryAfter)` or a retryable `Err`).
449    /// `make` receives the trigger's args, its post-retry `JobRow`, the error
450    /// string, the previous attempt number, and the next-run-at timestamp.
451    ///
452    /// Atomicity matches [`Self::on_completed_enqueue`]: worker-driven
453    /// retries dispatch in the transition's transaction, and a retry
454    /// driven by [`Client::retry_external`] dispatches in the same
455    /// transaction as the resolution. A spec failure rolls the retry
456    /// transition back so the external sender can redeliver.
457    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    /// Like [`Self::on_retried_enqueue`] but the closure returns an
472    /// [`EnqueueRequest`](crate::EnqueueRequest) with `InsertOpts` overrides.
473    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
504    /// is exhausted (retries used up *or* `Err(JobError::Terminal)`).
505    /// `make` receives the trigger's args, its post-failure `JobRow`, the
506    /// error string, and the attempt number.
507    ///
508    /// Atomicity matches [`Self::on_completed_enqueue`]: worker-driven
509    /// exhaustion dispatches in the transition's transaction, and an
510    /// exhaustion driven by [`Client::fail_external`] or
511    /// [`Client::resolve_callback`] with a `Fail` action dispatches in
512    /// the same transaction as the resolution. A spec failure rolls the
513    /// failure transition back so the external sender can redeliver.
514    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    /// Like [`Self::on_exhausted_enqueue`] but the closure returns an
526    /// [`EnqueueRequest`](crate::EnqueueRequest) with `InsertOpts` overrides.
527    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
552    /// is cancelled via `Ok(JobResult::Cancel(reason))`. `make` receives the
553    /// trigger's args, its post-cancellation `JobRow`, and the reason string.
554    ///
555    /// See [`Self::on_completed_enqueue`] for the same-transaction semantics.
556    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    /// Like [`Self::on_cancelled_enqueue`] but the closure returns an
568    /// [`EnqueueRequest`](crate::EnqueueRequest) with `InsertOpts` overrides.
569    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
594    /// parks on an external callback (`Ok(JobResult::WaitForCallback)`).
595    /// `make` receives the trigger's args plus the parked `JobRow`
596    /// (`job.callback_id` and `job.callback_timeout_at` identify the
597    /// pending callback).
598    ///
599    /// See [`Self::on_completed_enqueue`] for the same-transaction semantics.
600    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    /// Like [`Self::on_waiting_for_callback_enqueue`] but the closure returns
612    /// an [`EnqueueRequest`](crate::EnqueueRequest) with `InsertOpts` overrides.
613    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    /// Register a durable follow-up Awa job to enqueue when a job of type `T`
639    /// is rescued by maintenance (expired callback, stale heartbeat, or
640    /// deadline exceeded). `make` receives the trigger's args, its post-rescue
641    /// `JobRow`, and the [`RescueReason`](crate::events::RescueReason).
642    ///
643    /// **Atomicity: best-effort, separate transaction.** Maintenance rescues
644    /// commit the rescue UPDATE before this dispatcher runs; the follow-up
645    /// dispatch opens its own transaction afterwards. If the spec INSERT
646    /// fails the rescue stands and the failure is logged. Don't predicate
647    /// a workflow on a rescue-driven follow-up landing — for compensation
648    /// after a stuck attempt, build the follow-up handler to be safely
649    /// re-runnable. The atomic variant is tracked as an open extension in
650    /// ADR-029.
651    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    /// Like [`Self::on_rescued_enqueue`] but the closure returns an
663    /// [`EnqueueRequest`](crate::EnqueueRequest) with `InsertOpts` overrides.
664    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    /// Register a raw worker implementation.
693    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    /// Register an HTTP worker that dispatches jobs to a remote endpoint.
700    ///
701    /// In async mode the worker POSTs the job and parks in `waiting_external`.
702    /// In sync mode the worker awaits the HTTP response directly.
703    ///
704    /// Requires the `http-worker` feature.
705    #[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    /// Add shared state accessible via `ctx.extract::<T>()`.
716    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    /// Set the heartbeat interval (default: 30s).
722    pub fn heartbeat_interval(mut self, interval: Duration) -> Self {
723        self.heartbeat_interval = interval;
724        self
725    }
726
727    /// Set the scheduled/retryable promotion interval (default: 250ms).
728    pub fn promote_interval(mut self, interval: Duration) -> Self {
729        self.promote_interval = interval;
730        self
731    }
732
733    /// Set the stale-heartbeat rescue interval (default: 30s).
734    pub fn heartbeat_rescue_interval(mut self, interval: Duration) -> Self {
735        self.heartbeat_rescue_interval = Some(interval);
736        self
737    }
738
739    /// Set how long a heartbeat must be stale before the job is rescued (default: 90s).
740    ///
741    /// Should be at least 3× the heartbeat interval to avoid false rescues.
742    pub fn heartbeat_staleness(mut self, staleness: Duration) -> Self {
743        self.heartbeat_staleness = Some(staleness);
744        self
745    }
746
747    /// Set the deadline rescue interval (default: 30s).
748    pub fn deadline_rescue_interval(mut self, interval: Duration) -> Self {
749        self.deadline_rescue_interval = Some(interval);
750        self
751    }
752
753    /// Set the callback-timeout rescue interval (default: 30s).
754    pub fn callback_rescue_interval(mut self, interval: Duration) -> Self {
755        self.callback_rescue_interval = Some(interval);
756        self
757    }
758
759    /// Set the leader election retry interval (default: 10s).
760    ///
761    /// Controls how often a non-leader instance retries acquiring the maintenance
762    /// advisory lock. Lower values are useful in tests.
763    pub fn leader_election_interval(mut self, interval: Duration) -> Self {
764        self.leader_election_interval = Some(interval);
765        self
766    }
767
768    /// Set the leader connection health-check interval (default: 30s).
769    pub fn leader_check_interval(mut self, interval: Duration) -> Self {
770        self.leader_check_interval = Some(interval);
771        self
772    }
773
774    /// Set a global maximum worker count across all queues (enables weighted mode).
775    ///
776    /// When set, each queue gets `min_workers` guaranteed permits plus a share
777    /// of the remaining overflow capacity based on `weight`.
778    pub fn global_max_workers(mut self, max: u32) -> Self {
779        self.global_max_workers = Some(max);
780        self
781    }
782
783    /// Set retention for completed jobs (default: 24h).
784    pub fn completed_retention(mut self, retention: Duration) -> Self {
785        self.completed_retention = Some(retention);
786        self
787    }
788
789    /// Set retention for failed/cancelled jobs (default: 72h).
790    pub fn failed_retention(mut self, retention: Duration) -> Self {
791        self.failed_retention = Some(retention);
792        self
793    }
794
795    /// How long a descriptor catalog row can go un-refreshed before the
796    /// maintenance leader deletes it (default: 30 days). Pass
797    /// `Duration::ZERO` to disable — the catalog will then accumulate
798    /// rows indefinitely. See [`MaintenanceService::descriptor_retention`].
799    pub fn descriptor_retention(mut self, retention: Duration) -> Self {
800        self.descriptor_retention = Some(retention);
801        self
802    }
803
804    /// Set the maximum number of jobs to delete per cleanup pass (default: 1000).
805    pub fn cleanup_batch_size(mut self, batch_size: i64) -> Self {
806        self.cleanup_batch_size = Some(batch_size);
807        self
808    }
809
810    /// Set the cleanup interval (default: 60s).
811    pub fn cleanup_interval(mut self, interval: Duration) -> Self {
812        self.cleanup_interval = Some(interval);
813        self
814    }
815
816    /// Set a per-queue retention override.
817    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    /// Set how often runtime observability snapshots are published (default: 10s).
823    pub fn runtime_snapshot_interval(mut self, interval: Duration) -> Self {
824        self.runtime_snapshot_interval = interval;
825        self
826    }
827
828    /// Set the maintenance priority aging interval.
829    ///
830    /// This controls how often waiting available jobs are promoted toward
831    /// higher priority to prevent starvation. It is a global maintenance
832    /// setting for this worker runtime.
833    pub fn priority_aging_interval(mut self, interval: Duration) -> Self {
834        self.priority_aging_interval = Some(interval);
835        self
836    }
837
838    /// Set how often queue-storage terminal-count deltas are rolled into
839    /// folded live counters (default: 30s).
840    ///
841    /// Exact queue-depth reads include both folded counters and pending
842    /// deltas, so increasing this interval trades a larger append-only delta
843    /// ledger for fewer maintenance writes to the mutable counter table.
844    pub fn terminal_count_rollup_interval(mut self, interval: Duration) -> Self {
845        self.terminal_count_rollup_interval = Some(interval);
846        self
847    }
848
849    /// Set how often queue depth/lag metrics are published (default: 30s).
850    pub fn queue_stats_interval(mut self, interval: Duration) -> Self {
851        self.queue_stats_interval = Some(interval);
852        self
853    }
854
855    /// Enable or disable DLQ routing by default.
856    pub fn dlq_enabled_by_default(mut self, enabled: bool) -> Self {
857        self.dlq_enabled_by_default = enabled;
858        self
859    }
860
861    /// Override DLQ routing for a single queue.
862    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    /// Set retention for DLQ rows.
868    pub fn dlq_retention(mut self, retention: Duration) -> Self {
869        self.dlq_retention = Some(retention);
870        self
871    }
872
873    /// Set the maximum number of DLQ rows deleted per cleanup pass.
874    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    /// Override the segmented queue storage configuration for this runtime.
880    ///
881    /// Queue storage is the default worker engine. Canonical tables remain
882    /// migration compatibility, not a second supported worker runtime.
883    /// Use this to change the schema name or rotation sizing/timing.
884    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    /// Override the ADR-023 claim-ring rotation cadence.
904    ///
905    /// Defaults to `queue_rotate_interval` so claim partitions age out in
906    /// step with the ready / done partitions they reference. Only takes
907    /// effect when queue storage is active; no-op on the canonical engine.
908    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    /// Force the worker runtime onto canonical storage.
918    ///
919    /// This is primarily useful for migration/testing and benchmark
920    /// comparisons against the pre-0.6 engine. Production 0.6 runtimes should
921    /// normally use queue storage.
922    pub fn canonical_storage(mut self) -> Self {
923        self.storage = RuntimeStorage::Canonical;
924        self.storage_error = None;
925        self
926    }
927
928    /// Choose how this runtime participates in a storage transition.
929    pub fn transition_role(mut self, role: TransitionWorkerRole) -> Self {
930        self.transition_role = role;
931        self
932    }
933
934    /// Register a periodic (cron) job schedule.
935    ///
936    /// The schedule is synced to the database by the leader and evaluated
937    /// every second. When a fire is due, a job is atomically enqueued.
938    pub fn periodic(mut self, job: PeriodicJob) -> Self {
939        self.periodic_jobs.push(job);
940        self
941    }
942
943    /// Build the client.
944    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        // Validate rate limits and weights
971        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        // Validate batch size
989        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        // Validate weighted mode constraints
1007        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        // Warn if heartbeat_staleness is less than 3× heartbeat_interval
1027        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
1114/// A typed worker that deserializes args and calls a handler function.
1115struct 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
1145/// The Awa worker client — manages dispatchers, heartbeat, and maintenance.
1146pub 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    /// Cancellation token for dispatchers only — stops claiming new jobs.
1168    dispatch_cancel: CancellationToken,
1169    /// Cancellation token for heartbeat + maintenance — kept alive during drain.
1170    service_cancel: CancellationToken,
1171    /// Handles for dispatcher tasks.
1172    dispatcher_handles: RwLock<Vec<tokio::task::JoinHandle<()>>>,
1173    /// Handles for service tasks (heartbeat + maintenance).
1174    service_handles: RwLock<Vec<tokio::task::JoinHandle<()>>>,
1175    /// JoinSet tracking in-flight job tasks for graceful drain.
1176    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    /// Shared overflow pool for weighted mode (None in hard-reserved mode).
1184    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
1240/// Best-effort extraction of the most recent error message from a job's
1241/// `errors` history, for populating callback-driven `Exhausted`/`Retried`
1242/// events. Returns an empty string when no structured error is present.
1243fn 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    /// Create a new builder.
1255    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                // Fresh-install auto-finalize: state=canonical with no
1334                // operator commands run yet means a brand-new cluster
1335                // shouldn't have to step through prepare → enter-mixed-
1336                // transition → finalize manually. Install the queue-
1337                // storage schema (idempotent — concurrent workers are
1338                // safe), then ask the SQL gate to advance state directly
1339                // to `active` if the fresh-install conditions hold (no
1340                // canonical jobs, no live workers, prepared_engine
1341                // still NULL). Returns FALSE on any non-fresh DB; the
1342                // caller then falls back to the canonical-storage path
1343                // and the staged transition is unaffected.
1344                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                    // Another worker promoted concurrently while we
1360                    // were installing the schema — re-fetch and use
1361                    // the now-current state.
1362                    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                    // Function returned FALSE because conditions weren't
1369                    // met (canonical jobs exist, or another runtime is
1370                    // live but in canonical-only mode). Fall through to
1371                    // the canonical path; operators will need the
1372                    // staged transition.
1373                }
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    /// Start the worker runtime. Spawns dispatchers, heartbeat, and maintenance.
1499    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        // Completion batcher stays alive during drain so tasks can release
1529        // only after their completion has been acknowledged.
1530        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        // Create executor with metrics
1545        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        // Admin cancellation listener: fires the in-flight cancel flag
1560        // for any locally-running attempt when an admin issues
1561        // `cancel(job_id)` on the DB. Listen before dispatchers start
1562        // claiming so an early admin cancel cannot race listener setup.
1563        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        // Start heartbeat service (uses service_cancel — stays alive during drain)
1578        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        // Start maintenance service (uses service_cancel — stays alive during drain)
1592        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            // #456: a canonical-resolved worker keeps its configured
1606            // queue-storage runtime available so a pre-flip maintenance
1607            // leader can promote the queue-storage deferred backlog once
1608            // routing flips.
1609            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        // Start dispatcher/claimer loops per queue (uses dispatch_cancel — stops claiming first).
1675        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                    // queue_claimer_leases owner ids are independent lease tokens.
1717                    // Extra dispatcher loops in the same runtime need distinct tokens
1718                    // so they can hold separate bounded-claimer slots.
1719                    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    /// Graceful shutdown with drain timeout.
1756    ///
1757    /// Phased lifecycle:
1758    /// 1. Stop dispatchers (no new jobs claimed)
1759    /// 2. Signal in-flight jobs to cancel
1760    /// 3. Wait for dispatchers to exit
1761    /// 4. Drain in-flight jobs (heartbeat + maintenance still alive!)
1762    /// 5. Stop heartbeat + maintenance
1763    pub async fn shutdown(&self, timeout: Duration) {
1764        info!("Initiating graceful shutdown");
1765
1766        // Phase 1: Stop claiming new jobs
1767        self.dispatch_cancel.cancel();
1768
1769        self.publish_runtime_snapshot().await;
1770
1771        // Phase 2: Signal in-flight cancellation flags
1772        for flag in self.in_flight.flags() {
1773            flag.store(true, Ordering::SeqCst);
1774        }
1775
1776        // Phase 3: Wait for dispatchers to exit their poll loops
1777        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        // Phase 4: Drain in-flight jobs (heartbeat + maintenance still alive)
1786        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        // Phase 5: Stop background services (heartbeat + maintenance)
1798        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    /// Get the pool reference.
1811    pub fn pool(&self) -> &PgPool {
1812        &self.pool
1813    }
1814
1815    /// Resolve a pending external callback and dispatch the matching
1816    /// lifecycle event + ADR-029 follow-up specs.
1817    ///
1818    /// The callback transition (Complete / Fail) and any registered
1819    /// `on_completed_enqueue` / `on_exhausted_enqueue` follow-up `INSERT`s
1820    /// commit in a single transaction. A follow-up `INSERT` failure (or a
1821    /// panic in the user-supplied closure) rolls the callback transition
1822    /// back; the caller sees an `Err` and can surface a retryable failure
1823    /// to the external sender so the callback can be redelivered. `Ignored`
1824    /// produces no transition; `Resumed` does not currently fire a
1825    /// follow-up spec because resume re-enters execution and the executor
1826    /// emits the eventual outcome event itself.
1827    ///
1828    /// Prefer this over [`awa_model::admin::resolve_callback`] when you
1829    /// want `Completed`/`Exhausted` hooks or follow-up enqueues to fire
1830    /// for callback-driven outcomes. Hooks fire only in this process and
1831    /// only after the transaction commits.
1832    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        // ADR-029: dispatch follow-up specs in the same transaction. A
1844        // spec failure rolls the transition back together with the
1845        // follow-up, so callers can retry the external callback rather
1846        // than discovering a half-applied result.
1847        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    /// Complete a waiting job via its callback. The callback completion
1889    /// and any registered `on_completed_enqueue` follow-up `INSERT`s
1890    /// commit atomically. A spec INSERT failure rolls the completion
1891    /// back and returns `Err` so the external sender can retry. The
1892    /// in-process `Completed` hook fires after the transaction commits.
1893    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    /// Fail a waiting job via its callback. The callback failure and any
1918    /// registered `on_exhausted_enqueue` follow-up `INSERT`s commit
1919    /// atomically. A spec INSERT failure rolls the failure back and
1920    /// returns `Err` so the external sender can retry. The in-process
1921    /// `Exhausted` hook fires after the transaction commits.
1922    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    /// Requeue a waiting job via its callback. The callback retry and
1952    /// any registered `on_retried_enqueue` follow-up `INSERT`s commit
1953    /// atomically. A spec INSERT failure rolls the retry back and
1954    /// returns `Err` so the external sender can retry. The in-process
1955    /// `Retried` hook fires after the transaction commits.
1956    ///
1957    /// `admin::retry_external_in_tx` resets `attempt` to 0 as part of
1958    /// requeuing the job from scratch, so the returned `JobRow` no
1959    /// longer reflects the attempt that was being retried. The pre-retry
1960    /// attempt number is captured under the same row lock before the
1961    /// transition runs so the `Retried` event and spec context report
1962    /// the failed-attempt number — matching the inline `Retried`
1963    /// semantics where `attempt` is "the attempt that was retried", not
1964    /// the post-transition value.
1965    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    /// ADR-029 callback-resolution follow-up dispatch inside the
2013    /// caller-owned transaction. Errors (including caught panics from
2014    /// user-supplied `make` closures) propagate out so the caller can
2015    /// roll the callback transition back together with the failed
2016    /// follow-up; the external sender then sees a retryable failure
2017    /// rather than a half-applied result.
2018    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    /// Health check.
2040    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        // Emit OTel info gauges for every declared descriptor. One series per
2335        // descriptor, value=1, with all descriptor fields as attributes. Panels
2336        // lift descriptor fields into existing metrics via a Prometheus label
2337        // join: `awa_job_completed_total * on(awa_job_queue) group_left(awa_queue_display_name) awa_queue_info`.
2338        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    /// Insert a synthetic `transition_role=queue_storage_target` runtime
2629    /// row so the mixed-transition gate is satisfied without needing a
2630    /// second real client. Used by tests that only exercise the
2631    /// canonical-drain side of the transition.
2632    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    /// Drop and re-create the queue-storage schema using the supplied
2697    /// config, then advance the storage transition state machine to the
2698    /// `prepared` engine. The `lease_claim_receipts` flag must already
2699    /// match what the runtime under test will use, because the
2700    /// receipts-vs-legacy claim CTE is baked into `claim_ready_runtime`
2701    /// at `prepare_schema` time — a later store built with a different
2702    /// flag value would still hit the SQL function compiled here.
2703    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        // The explicit install remains an idempotence check now that
2953        // migrations materialize the default queue-storage substrate.
2954
2955        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        // Pin the legacy lease-materialization path. The receipt-plane
3162        // fast path (now the default since ADR-023 Phase 6) requires
3163        // deadline_duration=0 on QueueConfig and would error every claim
3164        // here; this test exercises the canonical-drain → cutover flow
3165        // with the standard 60s deadline, so lease_claim_receipts stays
3166        // off. Tests that specifically cover the receipt path opt back
3167        // in and zero the deadline. The flag must be set when
3168        // `prepare_schema` runs because the receipts-vs-legacy claim CTE
3169        // is baked into the SQL function definition at that point.
3170        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        // The mixed-transition gate now requires at least one runtime
3254        // running with `transition_role=queue_storage_target` (auto-role
3255        // runtimes downgrade to drain-only after the routing flip and
3256        // would leave the cluster with no queue-storage executor). This
3257        // test focuses on the canonical-drain side; insert a fake target
3258        // row so the gate passes without standing up a second real client.
3259        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        // See the matching note in
3372        // canonical_runtime_drains_in_flight_jobs_across_schema_upgrade_before_queue_storage_cutover:
3373        // pin the legacy materialization path so a 60s deadline_duration
3374        // on QueueConfig::default() doesn't collide with receipt-plane
3375        // mode. The flag must be set when `prepare_schema` runs because
3376        // the receipts-vs-legacy claim CTE is baked into the SQL
3377        // function definition at that point.
3378        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}