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