awa_worker/
executor.rs

1use crate::completion::CompletionBatcherHandle;
2use crate::context::{CallbackGuard, JobContext};
3use crate::events::{BoxedUntypedEventHandler, UntypedJobEvent};
4use crate::runtime::{InFlightMap, InFlightState, ProgressState};
5use crate::storage::{QueueStorageRuntime, RuntimeStorage};
6use awa_model::reschedule::{Reschedule, RescheduleOutcome};
7use awa_model::{AwaError, ClaimedEntry, ClaimedRuntimeJob, JobRow, JobState};
8use sqlx::PgPool;
9use std::any::Any;
10use std::collections::HashMap;
11use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
12use std::sync::Arc;
13use std::time::Duration;
14use tracing::{error, info, info_span, warn, Instrument};
15
16/// Result of executing a job handler.
17///
18/// # Picking the right variant for re-runs
19///
20/// Three primitives can put the job back on the queue. They differ in
21/// what they say about *why*:
22///
23/// | Primitive | Means | Increments `attempt` | Delay shape |
24/// |-----------|-------|----------------------|-------------|
25/// | [`JobError::Retryable`] | "this attempt failed; try again" | yes | DB-computed exponential backoff |
26/// | [`JobResult::RetryAfter`] | "this attempt failed; try again after delay X" | yes | caller-specified |
27/// | [`JobResult::Snooze`] | "this attempt didn't fail — it's just not time yet" | no | caller-specified |
28///
29/// Rule of thumb: if every "non-success" return is a "not yet" rather
30/// than a failure (polling, waiting for an upstream signal, rate
31/// limiting), use [`Snooze`] so `max_attempts` keeps its plain
32/// meaning of bounding genuine failures. If the handler observed a
33/// real failure and wants a specific retry delay rather than the
34/// default exponential backoff, use [`RetryAfter`]. See
35/// `awa/examples/poll_until_deadline.rs` for a deadline-bounded
36/// polling example.
37///
38/// [`Snooze`]: JobResult::Snooze
39/// [`RetryAfter`]: JobResult::RetryAfter
40#[derive(Debug)]
41pub enum JobResult {
42    /// Job completed successfully.
43    Completed,
44    /// Job should be retried after the given duration. Increments
45    /// `attempt`. Use when this attempt failed and you want a
46    /// caller-specified delay instead of the default exponential
47    /// backoff produced by [`JobError::Retryable`].
48    RetryAfter(std::time::Duration),
49    /// Job should be re-scheduled after the given duration without
50    /// counting as a failed attempt. Use for polling-style waits
51    /// where each "not yet" probe is normal — `max_attempts` should
52    /// only bound genuine handler failures, not the polling cadence.
53    Snooze(std::time::Duration),
54    /// Job should be cancelled. Records the reason in the job's
55    /// `errors` column and sets state to `cancelled` — no DLQ, no
56    /// failure event. Use for graceful give-up (e.g. handler-side
57    /// deadline expiry, user-requested abort).
58    Cancel(String),
59    /// Job is waiting for an external callback (webhook completion).
60    ///
61    /// Obtain the required guard from `ctx.register_callback()` or
62    /// `ctx.register_callback_with_config()`.
63    WaitForCallback(CallbackGuard),
64}
65
66/// Error type for job handlers — any error is retryable unless it's terminal.
67///
68/// [`JobError::Retryable`] triggers retry with database-computed exponential backoff.
69/// For explicit caller-controlled retry delay, return `Ok(`[`JobResult::RetryAfter`]`)` instead.
70#[derive(Debug, thiserror::Error)]
71pub enum JobError {
72    /// Retryable error — will be retried if attempts remain.
73    #[error("{0}")]
74    Retryable(#[source] Box<dyn std::error::Error + Send + Sync>),
75
76    /// Terminal error — immediately fails the job regardless of remaining attempts.
77    #[error("terminal: {0}")]
78    Terminal(String),
79}
80
81impl JobError {
82    /// Create a retryable error from any `std::error::Error`.
83    pub fn retryable(err: impl std::error::Error + Send + Sync + 'static) -> Self {
84        JobError::Retryable(Box::new(err))
85    }
86
87    /// Create a retryable error from a display message.
88    ///
89    /// Use this with `anyhow::Error` or other types that implement `Display`
90    /// but not `std::error::Error`:
91    /// ```ignore
92    /// Err(JobError::retryable_msg(format!("{err:#}")))
93    /// // or with anyhow:
94    /// Err(JobError::retryable_msg(err))
95    /// ```
96    pub fn retryable_msg(msg: impl std::fmt::Display) -> Self {
97        JobError::Retryable(Box::new(DisplayError(msg.to_string())))
98    }
99
100    /// Create a terminal error — immediately fails the job.
101    pub fn terminal(msg: impl Into<String>) -> Self {
102        JobError::Terminal(msg.into())
103    }
104}
105
106/// Per-queue DLQ policy resolved at `Client::start`.
107#[derive(Debug, Clone, Default)]
108pub struct DlqPolicy {
109    pub enabled_default: bool,
110    pub overrides: Arc<HashMap<String, bool>>,
111}
112
113impl DlqPolicy {
114    pub fn new(enabled_default: bool, overrides: HashMap<String, bool>) -> Self {
115        Self {
116            enabled_default,
117            overrides: Arc::new(overrides),
118        }
119    }
120
121    pub fn enabled_for(&self, queue: &str) -> bool {
122        self.overrides
123            .get(queue)
124            .copied()
125            .unwrap_or(self.enabled_default)
126    }
127}
128
129/// Wrapper to turn a Display string into a std::error::Error for retryable_msg.
130#[derive(Debug)]
131struct DisplayError(String);
132
133impl std::fmt::Display for DisplayError {
134    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135        f.write_str(&self.0)
136    }
137}
138
139impl std::error::Error for DisplayError {}
140
141/// With the `anyhow` feature, `?` works directly in handlers:
142/// ```ignore
143/// async fn perform(&self, ctx: &JobContext) -> Result<JobResult, JobError> {
144///     let data = fallible_thing().await?; // anyhow::Error → JobError::Retryable
145///     Ok(JobResult::Completed)
146/// }
147/// ```
148#[cfg(feature = "anyhow")]
149impl From<anyhow::Error> for JobError {
150    fn from(err: anyhow::Error) -> Self {
151        JobError::retryable_msg(format!("{err:#}"))
152    }
153}
154
155/// Worker trait — implement this for each job type.
156///
157/// # Handling permanent failure
158///
159/// When all retry attempts are exhausted, awa moves the job to `failed`.
160/// To run cleanup logic (update external state, send notifications), check
161/// the attempt count inside `perform`:
162///
163/// ```ignore
164/// async fn perform(&self, ctx: &JobContext) -> Result<JobResult, JobError> {
165///     match do_work(ctx).await {
166///         Ok(()) => Ok(JobResult::Completed),
167///         Err(err) if ctx.job.attempt >= ctx.job.max_attempts => {
168///             // Last attempt — run cleanup before awa marks as failed
169///             mark_permanently_failed(ctx.job.id).await;
170///             Err(JobError::retryable(err))
171///         }
172///         Err(err) => Err(JobError::retryable(err)),
173///     }
174/// }
175/// ```
176#[async_trait::async_trait]
177pub trait Worker: Send + Sync + 'static {
178    /// The kind string for this worker (must match the job's kind).
179    fn kind(&self) -> &'static str;
180
181    /// Execute the job. Access the job row via `ctx.job`.
182    async fn perform(&self, ctx: &JobContext) -> Result<JobResult, JobError>;
183}
184
185/// Type-erased worker wrapper for the registry.
186pub(crate) type BoxedWorker = Box<dyn Worker>;
187
188/// Result of a state-transition attempt in `complete_job`.
189#[allow(clippy::large_enum_variant)]
190enum CompletionOutcome {
191    /// The DB update was applied; optionally carries a lifecycle event to dispatch.
192    Applied {
193        event: Option<UntypedJobEvent>,
194        terminal: bool,
195    },
196    /// The job was already rescued/cancelled — stale completion, no event.
197    IgnoredStale,
198}
199
200#[derive(Debug, Clone)]
201pub(crate) struct DispatchedJob {
202    pub job: JobRow,
203    pub queue_storage_claim: Option<ClaimedEntry>,
204    pub queue_storage_unique_states: Option<String>,
205}
206
207/// Manages job execution — spawns worker futures and tracks in-flight jobs.
208pub struct JobExecutor {
209    pool: PgPool,
210    workers: Arc<HashMap<String, BoxedWorker>>,
211    lifecycle_handlers: Arc<HashMap<String, Vec<BoxedUntypedEventHandler>>>,
212    enqueue_specs: Arc<
213        HashMap<
214            crate::enqueue_specs::Outcome,
215            HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
216        >,
217    >,
218    in_flight: InFlightMap,
219    queue_in_flight: Arc<HashMap<String, Arc<AtomicU32>>>,
220    state: Arc<HashMap<std::any::TypeId, Box<dyn Any + Send + Sync>>>,
221    metrics: crate::metrics::AwaMetrics,
222    completion_batcher: CompletionBatcherHandle,
223    storage: RuntimeStorage,
224    dlq_policy: DlqPolicy,
225}
226
227impl JobExecutor {
228    #[allow(clippy::too_many_arguments)]
229    pub(crate) fn new(
230        pool: PgPool,
231        workers: Arc<HashMap<String, BoxedWorker>>,
232        lifecycle_handlers: Arc<HashMap<String, Vec<BoxedUntypedEventHandler>>>,
233        enqueue_specs: Arc<
234            HashMap<
235                crate::enqueue_specs::Outcome,
236                HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
237            >,
238        >,
239        in_flight: InFlightMap,
240        queue_in_flight: Arc<HashMap<String, Arc<AtomicU32>>>,
241        state: Arc<HashMap<std::any::TypeId, Box<dyn Any + Send + Sync>>>,
242        metrics: crate::metrics::AwaMetrics,
243        completion_batcher: CompletionBatcherHandle,
244        storage: RuntimeStorage,
245        dlq_policy: DlqPolicy,
246    ) -> Self {
247        Self {
248            pool,
249            workers,
250            lifecycle_handlers,
251            enqueue_specs,
252            in_flight,
253            queue_in_flight,
254            state,
255            metrics,
256            completion_batcher,
257            storage,
258            dlq_policy,
259        }
260    }
261
262    /// Build the future that executes a claimed job.
263    ///
264    /// The caller is responsible for spawning it onto the runtime.
265    pub(crate) fn execute_task(
266        &self,
267        dispatched: DispatchedJob,
268        cancel: Arc<AtomicBool>,
269    ) -> impl std::future::Future<Output = ()> + Send + 'static {
270        let job = dispatched.job;
271        let queue_storage_claim = dispatched.queue_storage_claim;
272        let queue_storage_unique_states = dispatched.queue_storage_unique_states;
273        let pool = self.pool.clone();
274        let workers = self.workers.clone();
275        let lifecycle_handlers = self.lifecycle_handlers.clone();
276        let enqueue_specs = self.enqueue_specs.clone();
277        let in_flight = self.in_flight.clone();
278        let queue_in_flight = self.queue_in_flight.clone();
279        let state = self.state.clone();
280        let metrics = self.metrics.clone();
281        let completion_batcher = self.completion_batcher.clone();
282        let storage = self.storage.clone();
283        let dlq_policy = self.dlq_policy.clone();
284        let job_id = job.id;
285        let job_run_lease = job.run_lease;
286        let job_kind = job.kind.clone();
287        let job_queue = job.queue.clone();
288
289        let span = info_span!(
290            "job.execute",
291            job.id = job_id,
292            job.kind = %job_kind,
293            job.queue = %job_queue,
294            job.attempt = job.attempt,
295            otel.name = %format!("job.execute {}", job_kind),
296            otel.status_code = tracing::field::Empty,
297        );
298
299        async move {
300            // Seed progress from the persisted checkpoint (for retries/snoozes)
301            let progress_state = Arc::new(std::sync::Mutex::new(ProgressState::new(
302                job.progress.clone(),
303            )));
304
305            // Register as in-flight with cancel + progress
306            let in_flight_state = InFlightState {
307                cancel: cancel.clone(),
308                progress: progress_state.clone(),
309            };
310            in_flight.insert((job_id, job_run_lease), in_flight_state);
311            if let Some(counter) = queue_in_flight.get(&job_queue) {
312                counter.fetch_add(1, Ordering::SeqCst);
313            }
314            metrics.record_in_flight_change(&job_queue, 1);
315            let has_lifecycle_handlers = lifecycle_handlers.contains_key(&job_kind);
316
317            let start = std::time::Instant::now();
318            let ctx = JobContext::new(
319                job.clone(),
320                cancel,
321                state,
322                pool.clone(),
323                storage.clone(),
324                progress_state.clone(),
325            );
326
327            let result = match workers.get(&job.kind) {
328                Some(worker) => {
329                    if has_lifecycle_handlers {
330                        let started_handlers = lifecycle_handlers.clone();
331                        let started_kind = job_kind.clone();
332                        let started_job = job.clone();
333                        tokio::spawn(async move {
334                            dispatch_lifecycle_event(
335                                &started_handlers,
336                                &started_kind,
337                                UntypedJobEvent::Started { job: started_job },
338                            )
339                            .await;
340                        });
341                    }
342                    worker.perform(&ctx).await
343                }
344                None => {
345                    error!(kind = %job.kind, job_id, "No worker registered for job kind");
346                    Err(JobError::Terminal(format!(
347                        "unknown job kind: {}",
348                        job.kind
349                    )))
350                }
351            };
352
353            let duration = start.elapsed();
354
355            // Snapshot progress for state transition
356            let progress_snapshot = {
357                let guard = progress_state.lock().expect("progress lock poisoned");
358                guard.clone_latest()
359            };
360
361            // Remove from in-flight immediately after the handler returns and
362            // the progress snapshot is captured. This keeps local worker
363            // capacity tied to active handler execution, not to the tail
364            // latency of durable completion bookkeeping.
365            in_flight.remove((job_id, job_run_lease));
366            if let Some(counter) = queue_in_flight.get(&job_queue) {
367                counter.fetch_sub(1, Ordering::SeqCst);
368            }
369            metrics.record_in_flight_change(&job_queue, -1);
370
371            let dlq_enabled = dlq_policy.enabled_for(&job_queue);
372            tokio::spawn(async move {
373                let outcome = complete_job(
374                    &pool,
375                    &job,
376                    queue_storage_claim.as_ref(),
377                    queue_storage_unique_states.as_deref(),
378                    &result,
379                    &completion_batcher,
380                    progress_snapshot,
381                    duration,
382                    has_lifecycle_handlers,
383                    &enqueue_specs,
384                    &storage,
385                    dlq_enabled,
386                    &metrics,
387                )
388                .await;
389
390                match &outcome {
391                    Ok(CompletionOutcome::Applied { terminal, .. }) => {
392                        // State transition succeeded — record metrics. `terminal`
393                        // is the source of truth for retry-vs-failure because
394                        // JobError::Retryable can resolve to either path.
395                        match &result {
396                            Ok(JobResult::Completed) => {
397                                metrics.record_job_completed(&job_kind, &job_queue, duration);
398                            }
399                            Ok(JobResult::RetryAfter(_)) => {
400                                metrics.record_job_retried(&job_kind, &job_queue);
401                            }
402                            Ok(JobResult::Cancel(_)) => {
403                                metrics.jobs_cancelled.add(
404                                    1,
405                                    &[
406                                        opentelemetry::KeyValue::new(
407                                            "awa.job.kind",
408                                            job_kind.clone(),
409                                        ),
410                                        opentelemetry::KeyValue::new(
411                                            "awa.job.queue",
412                                            job_queue.clone(),
413                                        ),
414                                    ],
415                                );
416                            }
417                            Ok(JobResult::Snooze(_)) => {}
418                            Ok(JobResult::WaitForCallback(_)) => {
419                                if *terminal {
420                                    metrics.record_job_failed(&job_kind, &job_queue, true);
421                                } else {
422                                    metrics.jobs_waiting_external.add(
423                                        1,
424                                        &[
425                                            opentelemetry::KeyValue::new(
426                                                "awa.job.kind",
427                                                job_kind.clone(),
428                                            ),
429                                            opentelemetry::KeyValue::new(
430                                                "awa.job.queue",
431                                                job_queue.clone(),
432                                            ),
433                                        ],
434                                    );
435                                }
436                            }
437                            Err(JobError::Terminal(_)) => {
438                                metrics.record_job_failed(&job_kind, &job_queue, true);
439                            }
440                            Err(JobError::Retryable(_)) => {
441                                if *terminal {
442                                    metrics.record_job_failed(&job_kind, &job_queue, true);
443                                } else {
444                                    metrics.record_job_retried(&job_kind, &job_queue);
445                                }
446                            }
447                        }
448                    }
449                    Ok(CompletionOutcome::IgnoredStale) => {}
450                    Err(err) => {
451                        error!(job_id, error = %err, "Failed to complete job");
452                    }
453                }
454
455                if let Ok(CompletionOutcome::Applied {
456                    event: Some(event), ..
457                }) = outcome
458                {
459                    dispatch_lifecycle_event(&lifecycle_handlers, &job_kind, event).await;
460                }
461            });
462        }
463        .instrument(span)
464    }
465}
466
467/// Update job state in the database based on handler result.
468///
469/// Returns a `CompletionOutcome` indicating whether the state transition was
470/// applied (with an optional lifecycle event) or ignored as stale.
471#[allow(clippy::too_many_arguments)]
472async fn complete_job(
473    pool: &PgPool,
474    job: &JobRow,
475    queue_storage_claim: Option<&ClaimedEntry>,
476    queue_storage_unique_states: Option<&str>,
477    result: &Result<JobResult, JobError>,
478    completion_batcher: &CompletionBatcherHandle,
479    progress_snapshot: Option<serde_json::Value>,
480    duration: Duration,
481    needs_event: bool,
482    enqueue_specs: &Arc<
483        HashMap<
484            crate::enqueue_specs::Outcome,
485            HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
486        >,
487    >,
488    storage: &RuntimeStorage,
489    dlq_enabled: bool,
490    metrics: &crate::metrics::AwaMetrics,
491) -> Result<CompletionOutcome, AwaError> {
492    match storage {
493        RuntimeStorage::Canonical => {
494            complete_job_canonical(
495                pool,
496                job,
497                result,
498                completion_batcher,
499                progress_snapshot,
500                duration,
501                needs_event,
502                enqueue_specs,
503                dlq_enabled,
504                metrics,
505            )
506            .await
507        }
508        RuntimeStorage::QueueStorage(runtime) => {
509            complete_job_queue_storage(
510                runtime,
511                pool,
512                job,
513                queue_storage_claim,
514                queue_storage_unique_states,
515                result,
516                completion_batcher,
517                progress_snapshot,
518                duration,
519                needs_event,
520                enqueue_specs,
521                dlq_enabled,
522                metrics,
523            )
524            .await
525        }
526    }
527}
528
529#[allow(clippy::too_many_arguments)]
530async fn complete_job_canonical(
531    pool: &PgPool,
532    job: &JobRow,
533    result: &Result<JobResult, JobError>,
534    completion_batcher: &CompletionBatcherHandle,
535    progress_snapshot: Option<serde_json::Value>,
536    duration: Duration,
537    needs_event: bool,
538    enqueue_specs: &Arc<
539        HashMap<
540            crate::enqueue_specs::Outcome,
541            HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
542        >,
543    >,
544    _dlq_enabled: bool,
545    _metrics: &crate::metrics::AwaMetrics,
546) -> Result<CompletionOutcome, AwaError> {
547    match result {
548        Ok(JobResult::Completed) => {
549            tracing::Span::current().record("otel.status_code", "OK");
550            info!(job_id = job.id, kind = %job.kind, attempt = job.attempt, "Job completed");
551
552            // ADR-029: when this kind has follow-up specs registered, drive
553            // completion through a dedicated transaction so the UPDATE and
554            // the follow-up INSERTs commit atomically. The batched path
555            // can't carry per-job follow-ups, so we bypass it here.
556            let kind_specs = enqueue_specs
557                .get(&crate::enqueue_specs::Outcome::Completed)
558                .and_then(|by_kind| by_kind.get(&job.kind))
559                .cloned();
560            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
561                let outcome = complete_canonical_with_followups(pool, job, &specs).await?;
562                return match outcome {
563                    None => {
564                        warn!(
565                            job_id = job.id,
566                            "Job already rescued/cancelled, completion ignored"
567                        );
568                        Ok(CompletionOutcome::IgnoredStale)
569                    }
570                    Some(updated_job) => {
571                        let event = if needs_event {
572                            Some(UntypedJobEvent::Completed {
573                                job: updated_job,
574                                duration,
575                            })
576                        } else {
577                            None
578                        };
579                        Ok(CompletionOutcome::Applied {
580                            event,
581                            terminal: false,
582                        })
583                    }
584                };
585            }
586
587            let result = match completion_batcher.complete(job.id, job.run_lease).await {
588                Ok(updated) => updated,
589                Err(err) => {
590                    warn!(
591                        job_id = job.id,
592                        error = %err,
593                        "Completion batch flush failed, falling back to direct finalize"
594                    );
595                    direct_complete_job(pool, job).await?
596                }
597            };
598            if !result {
599                warn!(
600                    job_id = job.id,
601                    "Job already rescued/cancelled, completion ignored"
602                );
603                return Ok(CompletionOutcome::IgnoredStale);
604            }
605            if needs_event {
606                let updated_job: JobRow =
607                    sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1")
608                        .bind(job.id)
609                        .fetch_one(pool)
610                        .await?;
611                Ok(CompletionOutcome::Applied {
612                    event: Some(UntypedJobEvent::Completed {
613                        job: updated_job,
614                        duration,
615                    }),
616                    terminal: false,
617                })
618            } else {
619                Ok(CompletionOutcome::Applied {
620                    event: None,
621                    terminal: false,
622                })
623            }
624        }
625
626        Ok(JobResult::RetryAfter(retry_duration)) => {
627            let seconds = retry_duration.as_secs() as f64;
628            info!(
629                job_id = job.id,
630                kind = %job.kind,
631                retry_after_secs = seconds,
632                "Job requested retry after duration"
633            );
634
635            // ADR-029: caller-requested retry. The Retried OutcomeContext
636            // carries an empty error string (parity with the event emitted
637            // below) and the attempt / next_run_at read from the post-UPDATE
638            // row.
639            let kind_specs = enqueue_specs
640                .get(&crate::enqueue_specs::Outcome::Retried)
641                .and_then(|by_kind| by_kind.get(&job.kind))
642                .cloned();
643            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
644                let result = retry_after_canonical_with_followups(
645                    pool,
646                    job,
647                    seconds,
648                    progress_snapshot.as_ref(),
649                    &specs,
650                )
651                .await?;
652                return match result {
653                    None => {
654                        warn!(
655                            job_id = job.id,
656                            "Job already rescued/cancelled, retry ignored"
657                        );
658                        Ok(CompletionOutcome::IgnoredStale)
659                    }
660                    Some(updated_job) => {
661                        let event = if needs_event {
662                            Some(UntypedJobEvent::Retried {
663                                job: updated_job.clone(),
664                                error: String::new(),
665                                attempt: updated_job.attempt,
666                                next_run_at: updated_job.run_at,
667                            })
668                        } else {
669                            None
670                        };
671                        Ok(CompletionOutcome::Applied {
672                            event,
673                            terminal: false,
674                        })
675                    }
676                };
677            }
678
679            let outcome = awa_model::reschedule::reschedule_canonical_attempt(
680                pool,
681                job.id,
682                job.run_lease,
683                Reschedule::RetryAfter {
684                    delay_secs: seconds,
685                },
686                None,
687                progress_snapshot.as_ref(),
688            )
689            .await?;
690            let (RescheduleOutcome::Rescheduled {
691                job_id,
692                run_at,
693                attempt,
694            }
695            | RescheduleOutcome::Migrated {
696                job_id,
697                run_at,
698                attempt,
699            }) = outcome
700            else {
701                warn!(
702                    job_id = job.id,
703                    "Job already rescued/cancelled, retry ignored"
704                );
705                return Ok(CompletionOutcome::IgnoredStale);
706            };
707            if needs_event {
708                let updated_job = rescheduled_event_row(
709                    job,
710                    JobState::Retryable,
711                    job_id,
712                    run_at,
713                    attempt,
714                    None,
715                    progress_snapshot,
716                );
717                Ok(CompletionOutcome::Applied {
718                    event: Some(UntypedJobEvent::Retried {
719                        job: updated_job,
720                        error: String::new(),
721                        attempt,
722                        next_run_at: run_at,
723                    }),
724                    terminal: false,
725                })
726            } else {
727                Ok(CompletionOutcome::Applied {
728                    event: None,
729                    terminal: false,
730                })
731            }
732        }
733
734        Ok(JobResult::Snooze(snooze_duration)) => {
735            let seconds = snooze_duration.as_secs() as f64;
736            info!(
737                job_id = job.id,
738                kind = %job.kind,
739                snooze_secs = seconds,
740                "Job snoozed (attempt not incremented)"
741            );
742            let outcome = awa_model::reschedule::reschedule_canonical_attempt(
743                pool,
744                job.id,
745                job.run_lease,
746                Reschedule::Snooze {
747                    delay_secs: seconds,
748                },
749                None,
750                progress_snapshot.as_ref(),
751            )
752            .await?;
753            if matches!(
754                outcome,
755                RescheduleOutcome::Stale | RescheduleOutcome::CancelledDuplicate { .. }
756            ) {
757                warn!(
758                    job_id = job.id,
759                    "Job already rescued/cancelled, snooze ignored"
760                );
761                return Ok(CompletionOutcome::IgnoredStale);
762            }
763            Ok(CompletionOutcome::Applied {
764                event: None,
765                terminal: false,
766            })
767        }
768
769        Ok(JobResult::Cancel(reason)) => {
770            tracing::Span::current().record("otel.status_code", "OK");
771            info!(
772                job_id = job.id,
773                kind = %job.kind,
774                reason = %reason,
775                "Job cancelled by handler"
776            );
777
778            // ADR-029: when this kind has Cancelled specs registered, drive
779            // the cancellation through a transaction so the UPDATE and the
780            // follow-up `INSERT`s commit atomically.
781            let kind_specs = enqueue_specs
782                .get(&crate::enqueue_specs::Outcome::Cancelled)
783                .and_then(|by_kind| by_kind.get(&job.kind))
784                .cloned();
785            let outcome_ctx = crate::enqueue_specs::OutcomeContext::Cancelled {
786                reason: reason.clone(),
787            };
788            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
789                let result = cancel_canonical_with_followups(
790                    pool,
791                    job,
792                    reason,
793                    progress_snapshot.as_ref(),
794                    &specs,
795                    &outcome_ctx,
796                )
797                .await?;
798                return match result {
799                    None => {
800                        warn!(
801                            job_id = job.id,
802                            "Job already rescued/cancelled, cancel ignored"
803                        );
804                        Ok(CompletionOutcome::IgnoredStale)
805                    }
806                    Some(updated_job) => {
807                        let event = if needs_event {
808                            Some(UntypedJobEvent::Cancelled {
809                                job: updated_job,
810                                reason: reason.clone(),
811                            })
812                        } else {
813                            None
814                        };
815                        Ok(CompletionOutcome::Applied {
816                            event,
817                            terminal: false,
818                        })
819                    }
820                };
821            }
822
823            // `awa.jobs_hot` directly, not the `awa.jobs` view: in canonical-drain
824            // mode the view is backed by the active queue-storage schema, so a view
825            // UPDATE silently misses the canonical row and the job wedges in
826            // `running` (same rationale as `complete_canonical_with_followups`).
827            // This applies to every canonical fast-path write and re-read here.
828            let result = sqlx::query(
829                r#"
830                UPDATE awa.jobs_hot
831                SET state = 'cancelled',
832                    finalized_at = now(),
833                    errors = errors || $2::jsonb,
834                    progress = $4
835                WHERE id = $1 AND state = 'running' AND run_lease = $3
836                "#,
837            )
838            .bind(job.id)
839            .bind(serde_json::json!({
840                "error": format!("cancelled: {}", reason),
841                "attempt": job.attempt,
842                "at": chrono::Utc::now().to_rfc3339()
843            }))
844            .bind(job.run_lease)
845            .bind(&progress_snapshot)
846            .execute(pool)
847            .await?;
848            if result.rows_affected() == 0 {
849                warn!(
850                    job_id = job.id,
851                    "Job already rescued/cancelled, cancel ignored"
852                );
853                return Ok(CompletionOutcome::IgnoredStale);
854            }
855            if needs_event {
856                let updated_job: JobRow =
857                    sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1")
858                        .bind(job.id)
859                        .fetch_one(pool)
860                        .await?;
861                Ok(CompletionOutcome::Applied {
862                    event: Some(UntypedJobEvent::Cancelled {
863                        job: updated_job,
864                        reason: reason.clone(),
865                    }),
866                    terminal: false,
867                })
868            } else {
869                Ok(CompletionOutcome::Applied {
870                    event: None,
871                    terminal: false,
872                })
873            }
874        }
875
876        Ok(JobResult::WaitForCallback(_guard)) => {
877            info!(
878                job_id = job.id,
879                kind = %job.kind,
880                "Job waiting for external callback"
881            );
882
883            // ADR-029: WaitingForCallback follow-up enqueue. The triggering
884            // UPDATE and the follow-up INSERTs commit in the same transaction.
885            // Race / missing-callback paths below are unchanged and dispatch
886            // no follow-ups because the row didn't actually park.
887            let kind_specs = enqueue_specs
888                .get(&crate::enqueue_specs::Outcome::WaitingForCallback)
889                .and_then(|by_kind| by_kind.get(&job.kind))
890                .cloned();
891            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
892                let parked =
893                    park_canonical_with_followups(pool, job, progress_snapshot.as_ref(), &specs)
894                        .await?;
895                if let Some(parked_job) = parked {
896                    let event = if needs_event {
897                        Some(UntypedJobEvent::WaitingForCallback { job: parked_job })
898                    } else {
899                        None
900                    };
901                    return Ok(CompletionOutcome::Applied {
902                        event,
903                        terminal: false,
904                    });
905                }
906                // Fall through to the existing rows_affected == 0 handling.
907            }
908
909            let result = sqlx::query(
910                r#"
911                UPDATE awa.jobs_hot
912                SET state = 'waiting_external',
913                    heartbeat_at = NULL,
914                    deadline_at = NULL,
915                    progress = $3
916                WHERE id = $1 AND state = 'running' AND run_lease = $2 AND callback_id IS NOT NULL
917                "#,
918            )
919            .bind(job.id)
920            .bind(job.run_lease)
921            .bind(&progress_snapshot)
922            .execute(pool)
923            .await?;
924            if result.rows_affected() == 0 {
925                // A rescued row may have moved to `scheduled_jobs`, so the
926                // race check reads both canonical tables.
927                let current: Option<(JobState, Option<uuid::Uuid>)> = sqlx::query_as(
928                    "SELECT state, callback_id FROM awa.jobs_hot WHERE id = $1 \
929                     UNION ALL \
930                     SELECT state, callback_id FROM awa.scheduled_jobs WHERE id = $1",
931                )
932                .bind(job.id)
933                .fetch_optional(pool)
934                .await?;
935                match current {
936                    Some((state, _)) if state.is_terminal() => {
937                        info!(
938                            job_id = job.id,
939                            state = %state,
940                            "Job already completed by racing callback"
941                        );
942                        return Ok(CompletionOutcome::Applied {
943                            event: None,
944                            terminal: false,
945                        });
946                    }
947                    Some((_, None)) => {
948                        error!(
949                            job_id = job.id,
950                            "WaitForCallback returned without calling register_callback"
951                        );
952                        let result = sqlx::query(
953                            r#"
954                            UPDATE awa.jobs_hot
955                            SET state = 'failed',
956                                finalized_at = now(),
957                                errors = errors || $2::jsonb
958                            WHERE id = $1 AND state = 'running' AND run_lease = $3
959                            "#,
960                        )
961                        .bind(job.id)
962                        .bind(serde_json::json!({
963                            "error": "WaitForCallback returned without calling register_callback",
964                            "attempt": job.attempt,
965                            "at": chrono::Utc::now().to_rfc3339(),
966                            "terminal": true
967                        }))
968                        .bind(job.run_lease)
969                        .execute(pool)
970                        .await?;
971                        if result.rows_affected() == 0 {
972                            return Ok(CompletionOutcome::IgnoredStale);
973                        }
974                        return Ok(CompletionOutcome::Applied {
975                            event: None,
976                            terminal: true,
977                        });
978                    }
979                    _ => {
980                        warn!(
981                            job_id = job.id,
982                            "Job already rescued/cancelled, wait-for-callback ignored"
983                        );
984                        return Ok(CompletionOutcome::IgnoredStale);
985                    }
986                }
987            }
988            let event = if needs_event {
989                let parked_job: JobRow = sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1")
990                    .bind(job.id)
991                    .fetch_one(pool)
992                    .await?;
993                Some(UntypedJobEvent::WaitingForCallback { job: parked_job })
994            } else {
995                None
996            };
997            Ok(CompletionOutcome::Applied {
998                event,
999                terminal: false,
1000            })
1001        }
1002
1003        Err(JobError::Terminal(msg)) => {
1004            tracing::Span::current().record("otel.status_code", "ERROR");
1005            error!(
1006                job_id = job.id,
1007                kind = %job.kind,
1008                error = %msg,
1009                "Job failed terminally"
1010            );
1011
1012            // ADR-029: terminal error counts as Exhausted (the error is
1013            // fatal; no further attempts will run).
1014            let kind_specs = enqueue_specs
1015                .get(&crate::enqueue_specs::Outcome::Exhausted)
1016                .and_then(|by_kind| by_kind.get(&job.kind))
1017                .cloned();
1018            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1019                let result = terminal_canonical_with_followups(
1020                    pool,
1021                    job,
1022                    msg,
1023                    progress_snapshot.as_ref(),
1024                    &specs,
1025                )
1026                .await?;
1027                return match result {
1028                    None => {
1029                        warn!(
1030                            job_id = job.id,
1031                            "Job already rescued/cancelled, terminal failure ignored"
1032                        );
1033                        Ok(CompletionOutcome::IgnoredStale)
1034                    }
1035                    Some(updated_job) => {
1036                        let event = if needs_event {
1037                            Some(UntypedJobEvent::Exhausted {
1038                                job: updated_job,
1039                                error: msg.clone(),
1040                                attempt: job.attempt,
1041                            })
1042                        } else {
1043                            None
1044                        };
1045                        Ok(CompletionOutcome::Applied {
1046                            event,
1047                            terminal: true,
1048                        })
1049                    }
1050                };
1051            }
1052
1053            let result = sqlx::query(
1054                r#"
1055                UPDATE awa.jobs_hot
1056                SET state = 'failed',
1057                    finalized_at = now(),
1058                    errors = errors || $2::jsonb,
1059                    progress = $4
1060                WHERE id = $1 AND state = 'running' AND run_lease = $3
1061                "#,
1062            )
1063            .bind(job.id)
1064            .bind(serde_json::json!({
1065                "error": msg.to_string(),
1066                "attempt": job.attempt,
1067                "at": chrono::Utc::now().to_rfc3339(),
1068                "terminal": true
1069            }))
1070            .bind(job.run_lease)
1071            .bind(&progress_snapshot)
1072            .execute(pool)
1073            .await?;
1074            if result.rows_affected() == 0 {
1075                warn!(
1076                    job_id = job.id,
1077                    "Job already rescued/cancelled, terminal failure ignored"
1078                );
1079                return Ok(CompletionOutcome::IgnoredStale);
1080            }
1081            if needs_event {
1082                let updated_job: JobRow =
1083                    sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1")
1084                        .bind(job.id)
1085                        .fetch_one(pool)
1086                        .await?;
1087                Ok(CompletionOutcome::Applied {
1088                    event: Some(UntypedJobEvent::Exhausted {
1089                        job: updated_job,
1090                        error: msg.clone(),
1091                        attempt: job.attempt,
1092                    }),
1093                    terminal: true,
1094                })
1095            } else {
1096                Ok(CompletionOutcome::Applied {
1097                    event: None,
1098                    terminal: true,
1099                })
1100            }
1101        }
1102
1103        Err(JobError::Retryable(err)) => {
1104            let error_msg = err.to_string();
1105            if job.attempt >= job.max_attempts {
1106                tracing::Span::current().record("otel.status_code", "ERROR");
1107                error!(
1108                    job_id = job.id,
1109                    kind = %job.kind,
1110                    attempt = job.attempt,
1111                    max_attempts = job.max_attempts,
1112                    error = %error_msg,
1113                    "Job failed (max attempts exhausted)"
1114                );
1115
1116                // ADR-029: retries exhausted -> Exhausted outcome.
1117                let kind_specs = enqueue_specs
1118                    .get(&crate::enqueue_specs::Outcome::Exhausted)
1119                    .and_then(|by_kind| by_kind.get(&job.kind))
1120                    .cloned();
1121                if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1122                    let result = exhaust_canonical_with_followups(
1123                        pool,
1124                        job,
1125                        &error_msg,
1126                        progress_snapshot.as_ref(),
1127                        &specs,
1128                    )
1129                    .await?;
1130                    return match result {
1131                        None => {
1132                            warn!(
1133                                job_id = job.id,
1134                                "Job already rescued/cancelled, failure ignored"
1135                            );
1136                            Ok(CompletionOutcome::IgnoredStale)
1137                        }
1138                        Some(updated_job) => {
1139                            let event = if needs_event {
1140                                Some(UntypedJobEvent::Exhausted {
1141                                    job: updated_job,
1142                                    error: error_msg,
1143                                    attempt: job.attempt,
1144                                })
1145                            } else {
1146                                None
1147                            };
1148                            Ok(CompletionOutcome::Applied {
1149                                event,
1150                                terminal: true,
1151                            })
1152                        }
1153                    };
1154                }
1155
1156                let result = sqlx::query(
1157                    r#"
1158                    UPDATE awa.jobs_hot
1159                    SET state = 'failed',
1160                        finalized_at = now(),
1161                        errors = errors || $2::jsonb,
1162                        progress = $4
1163                    WHERE id = $1 AND state = 'running' AND run_lease = $3
1164                    "#,
1165                )
1166                .bind(job.id)
1167                .bind(serde_json::json!({
1168                    "error": error_msg,
1169                    "attempt": job.attempt,
1170                    "at": chrono::Utc::now().to_rfc3339()
1171                }))
1172                .bind(job.run_lease)
1173                .bind(&progress_snapshot)
1174                .execute(pool)
1175                .await?;
1176                if result.rows_affected() == 0 {
1177                    warn!(
1178                        job_id = job.id,
1179                        "Job already rescued/cancelled, failure ignored"
1180                    );
1181                    return Ok(CompletionOutcome::IgnoredStale);
1182                }
1183                if needs_event {
1184                    let updated_job: JobRow =
1185                        sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1")
1186                            .bind(job.id)
1187                            .fetch_one(pool)
1188                            .await?;
1189                    Ok(CompletionOutcome::Applied {
1190                        event: Some(UntypedJobEvent::Exhausted {
1191                            job: updated_job,
1192                            error: error_msg,
1193                            attempt: job.attempt,
1194                        }),
1195                        terminal: true,
1196                    })
1197                } else {
1198                    Ok(CompletionOutcome::Applied {
1199                        event: None,
1200                        terminal: true,
1201                    })
1202                }
1203            } else {
1204                warn!(
1205                    job_id = job.id,
1206                    kind = %job.kind,
1207                    attempt = job.attempt,
1208                    error = %error_msg,
1209                    "Job failed (will retry)"
1210                );
1211
1212                // ADR-029: retryable error with backoff -> Retried outcome.
1213                let kind_specs = enqueue_specs
1214                    .get(&crate::enqueue_specs::Outcome::Retried)
1215                    .and_then(|by_kind| by_kind.get(&job.kind))
1216                    .cloned();
1217                if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1218                    let result = retry_backoff_canonical_with_followups(
1219                        pool,
1220                        job,
1221                        &error_msg,
1222                        progress_snapshot.as_ref(),
1223                        &specs,
1224                    )
1225                    .await?;
1226                    return match result {
1227                        None => {
1228                            warn!(
1229                                job_id = job.id,
1230                                "Job already rescued/cancelled, retry ignored"
1231                            );
1232                            Ok(CompletionOutcome::IgnoredStale)
1233                        }
1234                        Some(updated_job) => {
1235                            let event = if needs_event {
1236                                Some(UntypedJobEvent::Retried {
1237                                    job: updated_job.clone(),
1238                                    error: error_msg,
1239                                    attempt: job.attempt,
1240                                    next_run_at: updated_job.run_at,
1241                                })
1242                            } else {
1243                                None
1244                            };
1245                            Ok(CompletionOutcome::Applied {
1246                                event,
1247                                terminal: false,
1248                            })
1249                        }
1250                    };
1251                }
1252
1253                let error_entry = serde_json::json!({
1254                    "error": error_msg,
1255                    "attempt": job.attempt,
1256                    "at": chrono::Utc::now().to_rfc3339()
1257                });
1258                let outcome = awa_model::reschedule::reschedule_canonical_attempt(
1259                    pool,
1260                    job.id,
1261                    job.run_lease,
1262                    Reschedule::RetryBackoff,
1263                    Some(&error_entry),
1264                    progress_snapshot.as_ref(),
1265                )
1266                .await?;
1267                let (RescheduleOutcome::Rescheduled {
1268                    job_id,
1269                    run_at,
1270                    attempt,
1271                }
1272                | RescheduleOutcome::Migrated {
1273                    job_id,
1274                    run_at,
1275                    attempt,
1276                }) = outcome
1277                else {
1278                    warn!(
1279                        job_id = job.id,
1280                        "Job already rescued/cancelled, retry ignored"
1281                    );
1282                    return Ok(CompletionOutcome::IgnoredStale);
1283                };
1284                if needs_event {
1285                    let updated_job = rescheduled_event_row(
1286                        job,
1287                        JobState::Retryable,
1288                        job_id,
1289                        run_at,
1290                        attempt,
1291                        Some(&error_entry),
1292                        progress_snapshot,
1293                    );
1294                    Ok(CompletionOutcome::Applied {
1295                        event: Some(UntypedJobEvent::Retried {
1296                            job: updated_job,
1297                            error: error_msg,
1298                            attempt: job.attempt,
1299                            next_run_at: run_at,
1300                        }),
1301                        terminal: false,
1302                    })
1303                } else {
1304                    Ok(CompletionOutcome::Applied {
1305                        event: None,
1306                        terminal: false,
1307                    })
1308                }
1309            }
1310        }
1311    }
1312}
1313
1314/// Synthesize the post-reschedule row for event emission. After a migrated
1315/// re-schedule the job lives in the queue-storage deferred backlog under a
1316/// new id, so re-reading `awa.jobs` by the old id would return nothing.
1317///
1318/// This must mirror every mutation the reschedule persisted, not just the
1319/// routing fields: hooks inspecting `JobEvent::Retried.job` read `errors` to
1320/// see the failure that triggered the retry, and `finalized_at` to see when
1321/// the attempt closed.
1322fn rescheduled_event_row(
1323    job: &JobRow,
1324    state: JobState,
1325    job_id: i64,
1326    run_at: chrono::DateTime<chrono::Utc>,
1327    attempt: i16,
1328    error: Option<&serde_json::Value>,
1329    progress: Option<serde_json::Value>,
1330) -> JobRow {
1331    let mut row = job.clone();
1332    row.id = job_id;
1333    row.state = state;
1334    row.run_at = run_at;
1335    row.attempt = attempt;
1336    row.heartbeat_at = None;
1337    row.deadline_at = None;
1338    row.progress = progress;
1339    if let Some(error) = error {
1340        row.errors.get_or_insert_with(Vec::new).push(error.clone());
1341    }
1342    if state == JobState::Retryable {
1343        row.finalized_at = Some(chrono::Utc::now());
1344    }
1345    row
1346}
1347
1348#[allow(clippy::too_many_arguments)]
1349async fn complete_job_queue_storage(
1350    runtime: &QueueStorageRuntime,
1351    pool: &PgPool,
1352    job: &JobRow,
1353    queue_storage_claim: Option<&ClaimedEntry>,
1354    queue_storage_unique_states: Option<&str>,
1355    result: &Result<JobResult, JobError>,
1356    completion_batcher: &CompletionBatcherHandle,
1357    progress_snapshot: Option<serde_json::Value>,
1358    duration: Duration,
1359    needs_event: bool,
1360    enqueue_specs: &Arc<
1361        HashMap<
1362            crate::enqueue_specs::Outcome,
1363            HashMap<String, Vec<crate::enqueue_specs::BoxedEnqueueSpec>>,
1364        >,
1365    >,
1366    dlq_enabled: bool,
1367    metrics: &crate::metrics::AwaMetrics,
1368) -> Result<CompletionOutcome, AwaError> {
1369    match result {
1370        Ok(JobResult::Completed) => {
1371            tracing::Span::current().record("otel.status_code", "OK");
1372            info!(job_id = job.id, kind = %job.kind, attempt = job.attempt, "Job completed");
1373
1374            // ADR-029: when this kind has follow-up specs registered, drive
1375            // completion through a dedicated transaction so the
1376            // receipt-plane / lease cleanup + done_entries append +
1377            // follow-up `INSERT`s commit atomically. The receipt-plane
1378            // fast-complete (ADR-023) can't carry per-job follow-ups, so we
1379            // bypass it for spec'd jobs and use the slow path's tx-aware
1380            // variant.
1381            let kind_specs = enqueue_specs
1382                .get(&crate::enqueue_specs::Outcome::Completed)
1383                .and_then(|by_kind| by_kind.get(&job.kind))
1384                .cloned();
1385            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1386                let outcome = complete_queue_storage_with_followups(
1387                    runtime,
1388                    pool,
1389                    job,
1390                    queue_storage_claim,
1391                    queue_storage_unique_states,
1392                    &specs,
1393                )
1394                .await?;
1395                return match outcome {
1396                    None => {
1397                        warn!(
1398                            job_id = job.id,
1399                            "Job already rescued/cancelled, completion ignored"
1400                        );
1401                        Ok(CompletionOutcome::IgnoredStale)
1402                    }
1403                    Some(updated_job) => {
1404                        let event = if needs_event {
1405                            Some(UntypedJobEvent::Completed {
1406                                job: updated_job,
1407                                duration,
1408                            })
1409                        } else {
1410                            None
1411                        };
1412                        Ok(CompletionOutcome::Applied {
1413                            event,
1414                            terminal: false,
1415                        })
1416                    }
1417                };
1418            }
1419
1420            let updated = match match queue_storage_claim {
1421                Some(claim) => {
1422                    completion_batcher
1423                        .complete_runtime_job(ClaimedRuntimeJob {
1424                            claim: claim.clone(),
1425                            job: job.clone(),
1426                            unique_states: queue_storage_unique_states
1427                                .map(std::string::ToString::to_string),
1428                        })
1429                        .await
1430                }
1431                None => completion_batcher.complete(job.id, job.run_lease).await,
1432            } {
1433                Ok(updated) => updated,
1434                Err(err) => {
1435                    warn!(
1436                        job_id = job.id,
1437                        error = %err,
1438                        "Completion batch flush failed, falling back to direct finalize"
1439                    );
1440                    direct_complete_job_queue_storage(
1441                        runtime,
1442                        pool,
1443                        job,
1444                        queue_storage_claim,
1445                        queue_storage_unique_states,
1446                    )
1447                    .await?
1448                }
1449            };
1450            if !updated {
1451                warn!(
1452                    job_id = job.id,
1453                    "Job already rescued/cancelled, completion ignored"
1454                );
1455                return Ok(CompletionOutcome::IgnoredStale);
1456            }
1457            if needs_event {
1458                let updated_job =
1459                    runtime
1460                        .store
1461                        .load_job(pool, job.id)
1462                        .await?
1463                        .unwrap_or_else(|| {
1464                            let mut completed_job = job.clone();
1465                            completed_job.state = JobState::Completed;
1466                            completed_job.finalized_at = Some(chrono::Utc::now());
1467                            completed_job.progress = None;
1468                            completed_job
1469                        });
1470                Ok(CompletionOutcome::Applied {
1471                    event: Some(UntypedJobEvent::Completed {
1472                        job: updated_job,
1473                        duration,
1474                    }),
1475                    terminal: false,
1476                })
1477            } else {
1478                Ok(CompletionOutcome::Applied {
1479                    event: None,
1480                    terminal: false,
1481                })
1482            }
1483        }
1484
1485        Ok(JobResult::RetryAfter(retry_duration)) => {
1486            info!(
1487                job_id = job.id,
1488                kind = %job.kind,
1489                retry_after_secs = retry_duration.as_secs_f64(),
1490                "Job requested retry after duration"
1491            );
1492
1493            // ADR-029: caller-requested retry on queue storage.
1494            let kind_specs = enqueue_specs
1495                .get(&crate::enqueue_specs::Outcome::Retried)
1496                .and_then(|by_kind| by_kind.get(&job.kind))
1497                .cloned();
1498            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1499                let result = retry_after_queue_storage_with_followups(
1500                    runtime,
1501                    pool,
1502                    job,
1503                    *retry_duration,
1504                    progress_snapshot.clone(),
1505                    &specs,
1506                )
1507                .await?;
1508                return match result {
1509                    None => {
1510                        warn!(
1511                            job_id = job.id,
1512                            "Job already rescued/cancelled, retry ignored"
1513                        );
1514                        Ok(CompletionOutcome::IgnoredStale)
1515                    }
1516                    Some(updated_job) => {
1517                        let event = if needs_event {
1518                            Some(UntypedJobEvent::Retried {
1519                                job: updated_job.clone(),
1520                                error: String::new(),
1521                                attempt: updated_job.attempt,
1522                                next_run_at: updated_job.run_at,
1523                            })
1524                        } else {
1525                            None
1526                        };
1527                        Ok(CompletionOutcome::Applied {
1528                            event,
1529                            terminal: false,
1530                        })
1531                    }
1532                };
1533            }
1534
1535            let Some(updated_job) = runtime
1536                .store
1537                .retry_after(
1538                    pool,
1539                    job.id,
1540                    job.run_lease,
1541                    *retry_duration,
1542                    progress_snapshot.clone(),
1543                )
1544                .await?
1545            else {
1546                warn!(
1547                    job_id = job.id,
1548                    "Job already rescued/cancelled, retry ignored"
1549                );
1550                return Ok(CompletionOutcome::IgnoredStale);
1551            };
1552            if needs_event {
1553                Ok(CompletionOutcome::Applied {
1554                    event: Some(UntypedJobEvent::Retried {
1555                        job: updated_job.clone(),
1556                        error: String::new(),
1557                        attempt: updated_job.attempt,
1558                        next_run_at: updated_job.run_at,
1559                    }),
1560                    terminal: false,
1561                })
1562            } else {
1563                Ok(CompletionOutcome::Applied {
1564                    event: None,
1565                    terminal: false,
1566                })
1567            }
1568        }
1569
1570        Ok(JobResult::Snooze(snooze_duration)) => {
1571            info!(
1572                job_id = job.id,
1573                kind = %job.kind,
1574                snooze_secs = snooze_duration.as_secs_f64(),
1575                "Job snoozed (attempt not incremented)"
1576            );
1577            let updated = runtime
1578                .store
1579                .snooze(
1580                    pool,
1581                    job.id,
1582                    job.run_lease,
1583                    *snooze_duration,
1584                    progress_snapshot.clone(),
1585                )
1586                .await?;
1587            if updated.is_none() {
1588                warn!(
1589                    job_id = job.id,
1590                    "Job already rescued/cancelled, snooze ignored"
1591                );
1592                return Ok(CompletionOutcome::IgnoredStale);
1593            }
1594            Ok(CompletionOutcome::Applied {
1595                event: None,
1596                terminal: false,
1597            })
1598        }
1599
1600        Ok(JobResult::Cancel(reason)) => {
1601            info!(
1602                job_id = job.id,
1603                kind = %job.kind,
1604                reason = %reason,
1605                "Job cancelled by handler"
1606            );
1607
1608            // ADR-029: queue-storage Cancelled follow-ups.
1609            let kind_specs = enqueue_specs
1610                .get(&crate::enqueue_specs::Outcome::Cancelled)
1611                .and_then(|by_kind| by_kind.get(&job.kind))
1612                .cloned();
1613            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1614                let result = cancel_queue_storage_with_followups(
1615                    runtime,
1616                    pool,
1617                    job,
1618                    reason,
1619                    progress_snapshot.clone(),
1620                    &specs,
1621                )
1622                .await?;
1623                return match result {
1624                    None => {
1625                        warn!(
1626                            job_id = job.id,
1627                            "Job already rescued/cancelled, cancel ignored"
1628                        );
1629                        Ok(CompletionOutcome::IgnoredStale)
1630                    }
1631                    Some(updated_job) => {
1632                        let event = if needs_event {
1633                            Some(UntypedJobEvent::Cancelled {
1634                                job: updated_job,
1635                                reason: reason.clone(),
1636                            })
1637                        } else {
1638                            None
1639                        };
1640                        Ok(CompletionOutcome::Applied {
1641                            event,
1642                            terminal: false,
1643                        })
1644                    }
1645                };
1646            }
1647
1648            let Some(updated_job) = runtime
1649                .store
1650                .cancel_running(
1651                    pool,
1652                    job.id,
1653                    job.run_lease,
1654                    reason,
1655                    progress_snapshot.clone(),
1656                )
1657                .await?
1658            else {
1659                warn!(
1660                    job_id = job.id,
1661                    "Job already rescued/cancelled, cancel ignored"
1662                );
1663                return Ok(CompletionOutcome::IgnoredStale);
1664            };
1665            if needs_event {
1666                Ok(CompletionOutcome::Applied {
1667                    event: Some(UntypedJobEvent::Cancelled {
1668                        job: updated_job,
1669                        reason: reason.clone(),
1670                    }),
1671                    terminal: false,
1672                })
1673            } else {
1674                Ok(CompletionOutcome::Applied {
1675                    event: None,
1676                    terminal: false,
1677                })
1678            }
1679        }
1680
1681        Ok(JobResult::WaitForCallback(guard)) => {
1682            info!(
1683                job_id = job.id,
1684                kind = %job.kind,
1685                "Job waiting for external callback"
1686            );
1687
1688            // ADR-029: queue-storage WaitingForCallback follow-ups. Race /
1689            // missing-callback paths below remain unchanged — they dispatch
1690            // no follow-ups because the row didn't actually park.
1691            let kind_specs = enqueue_specs
1692                .get(&crate::enqueue_specs::Outcome::WaitingForCallback)
1693                .and_then(|by_kind| by_kind.get(&job.kind))
1694                .cloned();
1695            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1696                let parked =
1697                    park_queue_storage_with_followups(runtime, pool, job, guard.id(), &specs)
1698                        .await?;
1699                if let Some(parked_job) = parked {
1700                    let event = if needs_event {
1701                        Some(UntypedJobEvent::WaitingForCallback { job: parked_job })
1702                    } else {
1703                        None
1704                    };
1705                    return Ok(CompletionOutcome::Applied {
1706                        event,
1707                        terminal: false,
1708                    });
1709                }
1710                // Fall through to existing rows_affected == 0 handling.
1711            }
1712
1713            let entered = runtime
1714                .store
1715                .enter_callback_wait(pool, job.id, job.run_lease, guard.id())
1716                .await?;
1717            if !entered {
1718                let current = runtime.store.load_job(pool, job.id).await?;
1719                match current {
1720                    Some(current) if current.state.is_terminal() => {
1721                        info!(
1722                            job_id = job.id,
1723                            state = %current.state,
1724                            "Job already completed by racing callback"
1725                        );
1726                        return Ok(CompletionOutcome::Applied {
1727                            event: None,
1728                            terminal: false,
1729                        });
1730                    }
1731                    Some(current)
1732                        if current.state == JobState::Running && current.callback_id.is_none() =>
1733                    {
1734                        error!(
1735                            job_id = job.id,
1736                            "WaitForCallback returned without calling register_callback"
1737                        );
1738                        let failed = if dlq_enabled {
1739                            let failed = runtime
1740                                .store
1741                                .fail_to_dlq(
1742                                    pool,
1743                                    job.id,
1744                                    job.run_lease,
1745                                    "wait_for_callback_contract_violation",
1746                                    "WaitForCallback returned without calling register_callback",
1747                                    progress_snapshot.clone(),
1748                                )
1749                                .await?;
1750                            if failed.is_some() {
1751                                metrics.record_dlq_moved(
1752                                    &job.kind,
1753                                    &job.queue,
1754                                    "wait_for_callback_contract_violation",
1755                                );
1756                            }
1757                            failed
1758                        } else {
1759                            runtime
1760                                .store
1761                                .fail_terminal(
1762                                    pool,
1763                                    job.id,
1764                                    job.run_lease,
1765                                    "WaitForCallback returned without calling register_callback",
1766                                    progress_snapshot.clone(),
1767                                )
1768                                .await?
1769                        };
1770                        if failed.is_none() {
1771                            return Ok(CompletionOutcome::IgnoredStale);
1772                        }
1773                        return Ok(CompletionOutcome::Applied {
1774                            event: None,
1775                            terminal: true,
1776                        });
1777                    }
1778                    _ => {
1779                        warn!(
1780                            job_id = job.id,
1781                            "Job already rescued/cancelled, wait-for-callback ignored"
1782                        );
1783                        return Ok(CompletionOutcome::IgnoredStale);
1784                    }
1785                }
1786            }
1787            let event = if needs_event {
1788                let parked_job = runtime
1789                    .store
1790                    .load_job(pool, job.id)
1791                    .await?
1792                    .unwrap_or_else(|| {
1793                        let mut parked = job.clone();
1794                        parked.state = JobState::WaitingExternal;
1795                        parked
1796                    });
1797                Some(UntypedJobEvent::WaitingForCallback { job: parked_job })
1798            } else {
1799                None
1800            };
1801            Ok(CompletionOutcome::Applied {
1802                event,
1803                terminal: false,
1804            })
1805        }
1806
1807        Err(JobError::Terminal(msg)) => {
1808            tracing::Span::current().record("otel.status_code", "ERROR");
1809            error!(
1810                job_id = job.id,
1811                kind = %job.kind,
1812                error = %msg,
1813                "Job failed terminally"
1814            );
1815
1816            // ADR-029: queue-storage Exhausted follow-ups (Terminal counts).
1817            let kind_specs = enqueue_specs
1818                .get(&crate::enqueue_specs::Outcome::Exhausted)
1819                .and_then(|by_kind| by_kind.get(&job.kind))
1820                .cloned();
1821            if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1822                let result = fail_queue_storage_with_followups(
1823                    runtime,
1824                    pool,
1825                    job,
1826                    "terminal_error",
1827                    msg,
1828                    progress_snapshot.clone(),
1829                    dlq_enabled,
1830                    metrics,
1831                    &specs,
1832                )
1833                .await?;
1834                return match result {
1835                    None => {
1836                        warn!(
1837                            job_id = job.id,
1838                            "Job already rescued/cancelled, terminal failure ignored"
1839                        );
1840                        Ok(CompletionOutcome::IgnoredStale)
1841                    }
1842                    Some(updated_job) => {
1843                        let event = if needs_event {
1844                            Some(UntypedJobEvent::Exhausted {
1845                                job: updated_job,
1846                                error: msg.clone(),
1847                                attempt: job.attempt,
1848                            })
1849                        } else {
1850                            None
1851                        };
1852                        Ok(CompletionOutcome::Applied {
1853                            event,
1854                            terminal: true,
1855                        })
1856                    }
1857                };
1858            }
1859
1860            let updated_job = if dlq_enabled {
1861                let moved = runtime
1862                    .store
1863                    .fail_to_dlq(
1864                        pool,
1865                        job.id,
1866                        job.run_lease,
1867                        "terminal_error",
1868                        msg,
1869                        progress_snapshot.clone(),
1870                    )
1871                    .await?;
1872                if moved.is_some() {
1873                    metrics.record_dlq_moved(&job.kind, &job.queue, "terminal_error");
1874                }
1875                moved
1876            } else {
1877                runtime
1878                    .store
1879                    .fail_terminal(pool, job.id, job.run_lease, msg, progress_snapshot.clone())
1880                    .await?
1881            };
1882            let Some(updated_job) = updated_job else {
1883                warn!(
1884                    job_id = job.id,
1885                    "Job already rescued/cancelled, terminal failure ignored"
1886                );
1887                return Ok(CompletionOutcome::IgnoredStale);
1888            };
1889            if needs_event {
1890                Ok(CompletionOutcome::Applied {
1891                    event: Some(UntypedJobEvent::Exhausted {
1892                        job: updated_job,
1893                        error: msg.clone(),
1894                        attempt: job.attempt,
1895                    }),
1896                    terminal: true,
1897                })
1898            } else {
1899                Ok(CompletionOutcome::Applied {
1900                    event: None,
1901                    terminal: true,
1902                })
1903            }
1904        }
1905
1906        Err(JobError::Retryable(err)) => {
1907            let error_msg = err.to_string();
1908            if job.attempt >= job.max_attempts {
1909                tracing::Span::current().record("otel.status_code", "ERROR");
1910                error!(
1911                        job_id = job.id,
1912                        kind = %job.kind,
1913                        attempt = job.attempt,
1914                        max_attempts = job.max_attempts,
1915                    error = %error_msg,
1916                    "Job failed (max attempts exhausted)"
1917                );
1918
1919                // ADR-029: retries exhausted -> Exhausted outcome.
1920                let kind_specs = enqueue_specs
1921                    .get(&crate::enqueue_specs::Outcome::Exhausted)
1922                    .and_then(|by_kind| by_kind.get(&job.kind))
1923                    .cloned();
1924                if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
1925                    let result = fail_queue_storage_with_followups(
1926                        runtime,
1927                        pool,
1928                        job,
1929                        "max_attempts_exhausted",
1930                        &error_msg,
1931                        progress_snapshot.clone(),
1932                        dlq_enabled,
1933                        metrics,
1934                        &specs,
1935                    )
1936                    .await?;
1937                    return match result {
1938                        None => {
1939                            warn!(
1940                                job_id = job.id,
1941                                "Job already rescued/cancelled, failure ignored"
1942                            );
1943                            Ok(CompletionOutcome::IgnoredStale)
1944                        }
1945                        Some(updated_job) => {
1946                            let event = if needs_event {
1947                                Some(UntypedJobEvent::Exhausted {
1948                                    job: updated_job,
1949                                    error: error_msg,
1950                                    attempt: job.attempt,
1951                                })
1952                            } else {
1953                                None
1954                            };
1955                            Ok(CompletionOutcome::Applied {
1956                                event,
1957                                terminal: true,
1958                            })
1959                        }
1960                    };
1961                }
1962
1963                let updated_job = if dlq_enabled {
1964                    let moved = runtime
1965                        .store
1966                        .fail_to_dlq(
1967                            pool,
1968                            job.id,
1969                            job.run_lease,
1970                            "max_attempts_exhausted",
1971                            &error_msg,
1972                            progress_snapshot.clone(),
1973                        )
1974                        .await?;
1975                    if moved.is_some() {
1976                        metrics.record_dlq_moved(&job.kind, &job.queue, "max_attempts_exhausted");
1977                    }
1978                    moved
1979                } else {
1980                    runtime
1981                        .store
1982                        .fail_terminal(
1983                            pool,
1984                            job.id,
1985                            job.run_lease,
1986                            &error_msg,
1987                            progress_snapshot.clone(),
1988                        )
1989                        .await?
1990                };
1991                let Some(updated_job) = updated_job else {
1992                    warn!(
1993                        job_id = job.id,
1994                        "Job already rescued/cancelled, failure ignored"
1995                    );
1996                    return Ok(CompletionOutcome::IgnoredStale);
1997                };
1998                if needs_event {
1999                    Ok(CompletionOutcome::Applied {
2000                        event: Some(UntypedJobEvent::Exhausted {
2001                            job: updated_job,
2002                            error: error_msg,
2003                            attempt: job.attempt,
2004                        }),
2005                        terminal: true,
2006                    })
2007                } else {
2008                    Ok(CompletionOutcome::Applied {
2009                        event: None,
2010                        terminal: true,
2011                    })
2012                }
2013            } else {
2014                warn!(
2015                    job_id = job.id,
2016                    kind = %job.kind,
2017                    attempt = job.attempt,
2018                    error = %error_msg,
2019                    "Job failed (will retry)"
2020                );
2021
2022                // ADR-029: retryable error with backoff -> Retried outcome.
2023                let kind_specs = enqueue_specs
2024                    .get(&crate::enqueue_specs::Outcome::Retried)
2025                    .and_then(|by_kind| by_kind.get(&job.kind))
2026                    .cloned();
2027                if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) {
2028                    let result = retry_backoff_queue_storage_with_followups(
2029                        runtime,
2030                        pool,
2031                        job,
2032                        &error_msg,
2033                        progress_snapshot.clone(),
2034                        &specs,
2035                    )
2036                    .await?;
2037                    return match result {
2038                        None => {
2039                            warn!(
2040                                job_id = job.id,
2041                                "Job already rescued/cancelled, retry ignored"
2042                            );
2043                            Ok(CompletionOutcome::IgnoredStale)
2044                        }
2045                        Some(updated_job) => {
2046                            let event = if needs_event {
2047                                Some(UntypedJobEvent::Retried {
2048                                    job: updated_job.clone(),
2049                                    error: error_msg,
2050                                    attempt: job.attempt,
2051                                    next_run_at: updated_job.run_at,
2052                                })
2053                            } else {
2054                                None
2055                            };
2056                            Ok(CompletionOutcome::Applied {
2057                                event,
2058                                terminal: false,
2059                            })
2060                        }
2061                    };
2062                }
2063
2064                let Some(updated_job) = runtime
2065                    .store
2066                    .fail_retryable(
2067                        pool,
2068                        job.id,
2069                        job.run_lease,
2070                        &error_msg,
2071                        progress_snapshot.clone(),
2072                    )
2073                    .await?
2074                else {
2075                    warn!(
2076                        job_id = job.id,
2077                        "Job already rescued/cancelled, retry ignored"
2078                    );
2079                    return Ok(CompletionOutcome::IgnoredStale);
2080                };
2081                if needs_event {
2082                    Ok(CompletionOutcome::Applied {
2083                        event: Some(UntypedJobEvent::Retried {
2084                            job: updated_job.clone(),
2085                            error: error_msg,
2086                            attempt: job.attempt,
2087                            next_run_at: updated_job.run_at,
2088                        }),
2089                        terminal: false,
2090                    })
2091                } else {
2092                    Ok(CompletionOutcome::Applied {
2093                        event: None,
2094                        terminal: false,
2095                    })
2096                }
2097            }
2098        }
2099    }
2100}
2101
2102/// Complete a queue-storage job and run its registered follow-up enqueue
2103/// specs atomically with the completion (ADR-029).
2104///
2105/// Bypasses the receipt-plane fast-complete path and goes straight to the
2106/// tx-aware slow path so the lease/receipt cleanup, `done_entries` append,
2107/// and follow-up `INSERT`s all commit together. `complete_runtime_batch_slow`
2108/// already handles both receipt-claimed and materialised leases, so the
2109/// trade-off is purely losing receipt-plane fast-complete throughput for
2110/// spec'd jobs — acceptable since spec'd jobs do extra DB work anyway.
2111///
2112/// Returns:
2113/// - `Ok(None)` if the completion was stale — the lease was already rescued
2114///   or cancelled. The transaction is rolled back; no follow-ups emitted.
2115/// - `Ok(Some(updated_job))` if the completion committed; follow-ups have
2116///   been INSERTed in the same transaction. The returned row is the
2117///   post-completion snapshot, mirroring the fallback constructed elsewhere
2118///   in this file when `runtime.store.load_job` doesn't return one.
2119#[allow(clippy::explicit_auto_deref)]
2120async fn complete_queue_storage_with_followups(
2121    runtime: &QueueStorageRuntime,
2122    pool: &PgPool,
2123    job: &JobRow,
2124    queue_storage_claim: Option<&ClaimedEntry>,
2125    queue_storage_unique_states: Option<&str>,
2126    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2127) -> Result<Option<JobRow>, AwaError> {
2128    let Some(claim) = queue_storage_claim else {
2129        // The slow path is keyed on `ClaimedEntry` (it needs `lease_slot`,
2130        // `lane_seq`, `claim_slot`, etc. to clean up receipt-plane and
2131        // materialised lease rows). The executor always passes a claim when
2132        // it dispatches a queue-storage job, so this branch is normally
2133        // unreachable; report as stale rather than silently dropping the
2134        // outcome.
2135        warn!(
2136            job_id = job.id,
2137            "queue-storage completion with follow-up specs but no claim — \
2138             treating as stale"
2139        );
2140        return Ok(None);
2141    };
2142
2143    let runtime_job = ClaimedRuntimeJob {
2144        claim: claim.clone(),
2145        job: job.clone(),
2146        unique_states: queue_storage_unique_states.map(std::string::ToString::to_string),
2147    };
2148
2149    let mut tx = pool.begin().await?;
2150
2151    let updated = runtime
2152        .store
2153        .complete_runtime_batch_slow_in_tx(&mut tx, std::slice::from_ref(&runtime_job))
2154        .await?;
2155
2156    if updated.is_empty() {
2157        tx.rollback().await?;
2158        return Ok(None);
2159    }
2160
2161    // Synthesise the post-completion snapshot — the store has just moved the
2162    // lease into `done_entries` so a live SELECT would race with rotation.
2163    // The happy-path fields are exactly those `complete_job_queue_storage`
2164    // reconstructs when `load_job` returns None.
2165    let mut updated_job = job.clone();
2166    updated_job.state = JobState::Completed;
2167    updated_job.finalized_at = Some(chrono::Utc::now());
2168    updated_job.progress = None;
2169
2170    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, None).await?;
2171
2172    tx.commit().await?;
2173    Ok(Some(updated_job))
2174}
2175
2176/// Cancel a running queue-storage job and dispatch any Cancelled
2177/// follow-ups in the same transaction (ADR-029).
2178async fn cancel_queue_storage_with_followups(
2179    runtime: &QueueStorageRuntime,
2180    pool: &PgPool,
2181    job: &JobRow,
2182    reason: &str,
2183    progress_snapshot: Option<serde_json::Value>,
2184    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2185) -> Result<Option<JobRow>, AwaError> {
2186    let mut tx = pool.begin().await?;
2187    let Some(updated_job) = runtime
2188        .store
2189        .cancel_running_in_tx(&mut tx, job.id, job.run_lease, reason, progress_snapshot)
2190        .await?
2191    else {
2192        tx.rollback().await?;
2193        return Ok(None);
2194    };
2195    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Cancelled {
2196        reason: reason.to_string(),
2197    };
2198    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2199        .await?;
2200    tx.commit().await?;
2201    Ok(Some(updated_job))
2202}
2203
2204/// RetryAfter on queue-storage + Retried follow-ups, in the same tx.
2205async fn retry_after_queue_storage_with_followups(
2206    runtime: &QueueStorageRuntime,
2207    pool: &PgPool,
2208    job: &JobRow,
2209    retry_duration: Duration,
2210    progress_snapshot: Option<serde_json::Value>,
2211    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2212) -> Result<Option<JobRow>, AwaError> {
2213    let mut tx = pool.begin().await?;
2214    let Some(updated_job) = runtime
2215        .store
2216        .retry_after_in_tx(
2217            &mut tx,
2218            job.id,
2219            job.run_lease,
2220            retry_duration,
2221            progress_snapshot,
2222        )
2223        .await?
2224    else {
2225        tx.rollback().await?;
2226        return Ok(None);
2227    };
2228    // Mirror the canonical `RetryAfter` event: empty error string.
2229    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
2230        error: String::new(),
2231        attempt: updated_job.attempt,
2232        next_run_at: updated_job.run_at,
2233    };
2234    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2235        .await?;
2236    tx.commit().await?;
2237    Ok(Some(updated_job))
2238}
2239
2240/// Retryable error on queue-storage that backs off (attempt < max_attempts)
2241/// plus Retried follow-ups, in the same tx.
2242async fn retry_backoff_queue_storage_with_followups(
2243    runtime: &QueueStorageRuntime,
2244    pool: &PgPool,
2245    job: &JobRow,
2246    error_msg: &str,
2247    progress_snapshot: Option<serde_json::Value>,
2248    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2249) -> Result<Option<JobRow>, AwaError> {
2250    let mut tx = pool.begin().await?;
2251    let Some(updated_job) = runtime
2252        .store
2253        .fail_retryable_in_tx(&mut tx, job.id, job.run_lease, error_msg, progress_snapshot)
2254        .await?
2255    else {
2256        tx.rollback().await?;
2257        return Ok(None);
2258    };
2259    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
2260        error: error_msg.to_string(),
2261        attempt: job.attempt,
2262        next_run_at: updated_job.run_at,
2263    };
2264    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2265        .await?;
2266    tx.commit().await?;
2267    Ok(Some(updated_job))
2268}
2269
2270/// Terminal failure (or retries exhausted) on queue-storage + Exhausted
2271/// follow-ups, in the same tx. Honours DLQ routing if enabled.
2272#[allow(clippy::too_many_arguments)]
2273async fn fail_queue_storage_with_followups(
2274    runtime: &QueueStorageRuntime,
2275    pool: &PgPool,
2276    job: &JobRow,
2277    dlq_reason: &str,
2278    error_msg: &str,
2279    progress_snapshot: Option<serde_json::Value>,
2280    dlq_enabled: bool,
2281    metrics: &crate::metrics::AwaMetrics,
2282    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2283) -> Result<Option<JobRow>, AwaError> {
2284    let mut tx = pool.begin().await?;
2285    // Track whether the row routed to DLQ so the metric records after a
2286    // successful commit only — a spec-INSERT failure rolls the whole tx
2287    // back, and a phantom DLQ count would lie about a transition that
2288    // didn't actually happen.
2289    let (updated_job, routed_to_dlq) = if dlq_enabled {
2290        let moved = runtime
2291            .store
2292            .fail_to_dlq_in_tx(
2293                &mut tx,
2294                job.id,
2295                job.run_lease,
2296                dlq_reason,
2297                error_msg,
2298                progress_snapshot,
2299            )
2300            .await?;
2301        let routed = moved.is_some();
2302        (moved, routed)
2303    } else {
2304        let moved = runtime
2305            .store
2306            .fail_terminal_in_tx(&mut tx, job.id, job.run_lease, error_msg, progress_snapshot)
2307            .await?;
2308        (moved, false)
2309    };
2310    let Some(updated_job) = updated_job else {
2311        tx.rollback().await?;
2312        return Ok(None);
2313    };
2314    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
2315        error: error_msg.to_string(),
2316        attempt: job.attempt,
2317    };
2318    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2319        .await?;
2320    tx.commit().await?;
2321    if routed_to_dlq {
2322        metrics.record_dlq_moved(&job.kind, &job.queue, dlq_reason);
2323    }
2324    Ok(Some(updated_job))
2325}
2326
2327/// Park a queue-storage job in `waiting_external` and dispatch any
2328/// WaitingForCallback follow-ups in the same tx (ADR-029).
2329///
2330/// Returns `Ok(Some(parked))` if the row transitioned. `Ok(None)` if the
2331/// guarded UPDATE didn't match — caller falls back to the existing race /
2332/// missing-callback handling.
2333async fn park_queue_storage_with_followups(
2334    runtime: &QueueStorageRuntime,
2335    pool: &PgPool,
2336    job: &JobRow,
2337    callback_id: uuid::Uuid,
2338    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2339) -> Result<Option<JobRow>, AwaError> {
2340    let mut tx = pool.begin().await?;
2341    let entered = runtime
2342        .store
2343        .enter_callback_wait_in_tx(&mut tx, job.id, job.run_lease, callback_id)
2344        .await?;
2345    if !entered {
2346        tx.rollback().await?;
2347        return Ok(None);
2348    }
2349    // Re-read the parked row inside the same transaction. The input
2350    // `job` is the snapshot the executor claimed earlier — it does not
2351    // carry `callback_id` / `callback_timeout_at` written by
2352    // `register_callback()` during the handler, and the parked-state
2353    // fields (state, heartbeat_at, deadline_at) only land after the
2354    // enter_callback_wait_in_tx UPDATE above. Fall back to a synthesised
2355    // snapshot only if the SELECT misses (shouldn't happen because the
2356    // UPDATE just matched, but defensive).
2357    let parked_job = match runtime
2358        .store
2359        .load_active_lease_in_tx(&mut tx, job.id, job.run_lease)
2360        .await?
2361    {
2362        Some(row) => row,
2363        None => {
2364            let mut parked = job.clone();
2365            parked.state = JobState::WaitingExternal;
2366            parked.heartbeat_at = None;
2367            parked.deadline_at = None;
2368            parked.callback_id = Some(callback_id);
2369            parked
2370        }
2371    };
2372    crate::enqueue_specs::dispatch_specs_in_tx(
2373        &mut tx,
2374        &parked_job,
2375        specs,
2376        Some(&crate::enqueue_specs::OutcomeContext::WaitingForCallback),
2377    )
2378    .await?;
2379    tx.commit().await?;
2380    Ok(Some(parked_job))
2381}
2382
2383async fn direct_complete_job_queue_storage(
2384    runtime: &QueueStorageRuntime,
2385    pool: &PgPool,
2386    job: &JobRow,
2387    queue_storage_claim: Option<&ClaimedEntry>,
2388    queue_storage_unique_states: Option<&str>,
2389) -> Result<bool, AwaError> {
2390    let updated = if let Some(claim) = queue_storage_claim {
2391        let runtime_job = ClaimedRuntimeJob {
2392            claim: claim.clone(),
2393            job: job.clone(),
2394            unique_states: queue_storage_unique_states.map(std::string::ToString::to_string),
2395        };
2396        runtime
2397            .store
2398            .complete_runtime_batch(pool, std::slice::from_ref(&runtime_job))
2399            .await?
2400    } else {
2401        runtime
2402            .store
2403            .complete_job_batch_by_id(pool, &[(job.id, job.run_lease)])
2404            .await?
2405    };
2406    Ok(!updated.is_empty())
2407}
2408
2409/// Dispatch a lifecycle event to all registered handlers for a job kind.
2410///
2411/// Handlers are called sequentially. Panics are caught and logged — a
2412/// misbehaving handler cannot crash the dispatch loop or lose events
2413/// for subsequent handlers.
2414pub(crate) async fn dispatch_lifecycle_event(
2415    handlers: &HashMap<String, Vec<BoxedUntypedEventHandler>>,
2416    kind: &str,
2417    event: UntypedJobEvent,
2418) {
2419    if let Some(handlers) = handlers.get(kind) {
2420        for handler in handlers {
2421            let handler = handler.clone();
2422            let event = event.clone();
2423            let result = tokio::spawn(async move {
2424                (handler)(event).await;
2425            })
2426            .await;
2427            if let Err(err) = result {
2428                tracing::warn!(
2429                    kind,
2430                    error = %err,
2431                    "Lifecycle event handler panicked"
2432                );
2433            }
2434        }
2435    }
2436}
2437
2438async fn direct_complete_job(pool: &PgPool, job: &JobRow) -> Result<bool, AwaError> {
2439    let result = sqlx::query(
2440        r#"
2441        UPDATE awa.jobs_hot
2442        SET state = 'completed',
2443            finalized_at = now(),
2444            progress = NULL
2445        WHERE id = $1 AND state = 'running' AND run_lease = $2
2446        "#,
2447    )
2448    .bind(job.id)
2449    .bind(job.run_lease)
2450    .execute(pool)
2451    .await?;
2452
2453    Ok(result.rows_affected() > 0)
2454}
2455
2456/// Complete a canonical-storage job and run its registered follow-up enqueue
2457/// specs atomically with the completion UPDATE (ADR-029).
2458///
2459/// Returns:
2460/// - `Ok(None)` if the completion was stale (`rows_affected == 0`) — the job
2461///   has already been rescued or cancelled, no follow-ups are emitted, no
2462///   event should fire.
2463/// - `Ok(Some(updated_job))` if the completion committed; follow-ups have
2464///   been INSERTed in the same transaction. The returned row is the
2465///   post-completion snapshot (state = `completed`, `finalized_at` set).
2466// The follow-up loop reborrows `&mut *tx` per spec invocation so the same
2467// transaction handle can be reused; clippy reads the `*tx` as a redundant
2468// deref but `fetch_one`'s Executor bound requires the inner connection.
2469#[allow(clippy::explicit_auto_deref)]
2470async fn complete_canonical_with_followups(
2471    pool: &PgPool,
2472    job: &JobRow,
2473    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2474) -> Result<Option<JobRow>, AwaError> {
2475    let mut tx = pool.begin().await?;
2476
2477    // `UPDATE ... RETURNING` against `awa.jobs_hot` directly (not the
2478    // `awa.jobs` compatibility view): in canonical-drain mode the view is
2479    // backed by the active queue-storage schema and would return RowNotFound
2480    // for a canonical row, rolling back the just-applied completion and
2481    // wedging the job. The hot table is the source of truth for canonical
2482    // attempts regardless of routing mode.
2483    let updated_job: Option<JobRow> = sqlx::query_as(
2484        r#"
2485        UPDATE awa.jobs_hot
2486        SET state = 'completed',
2487            finalized_at = now(),
2488            progress = NULL
2489        WHERE id = $1 AND state = 'running' AND run_lease = $2
2490        RETURNING *
2491        "#,
2492    )
2493    .bind(job.id)
2494    .bind(job.run_lease)
2495    .fetch_optional(&mut *tx)
2496    .await?;
2497
2498    let Some(updated_job) = updated_job else {
2499        // Stale: another writer already finalised this attempt. Drop the
2500        // transaction without emitting follow-ups.
2501        tx.rollback().await?;
2502        return Ok(None);
2503    };
2504
2505    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, None).await?;
2506
2507    tx.commit().await?;
2508    Ok(Some(updated_job))
2509}
2510
2511/// Cancelled-outcome equivalent: cancel the running canonical row and
2512/// dispatch any registered `Cancelled` follow-ups inside the same
2513/// transaction. Returns the post-update [`JobRow`] for event emission, or
2514/// `None` when the row was already finalised by another writer.
2515async fn cancel_canonical_with_followups(
2516    pool: &PgPool,
2517    job: &JobRow,
2518    reason: &str,
2519    progress_snapshot: Option<&serde_json::Value>,
2520    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2521    outcome_ctx: &crate::enqueue_specs::OutcomeContext,
2522) -> Result<Option<JobRow>, AwaError> {
2523    let mut tx = pool.begin().await?;
2524
2525    let updated_job: Option<JobRow> = sqlx::query_as(
2526        r#"
2527        UPDATE awa.jobs_hot
2528        SET state = 'cancelled',
2529            finalized_at = now(),
2530            errors = errors || $2::jsonb,
2531            progress = $4
2532        WHERE id = $1 AND state = 'running' AND run_lease = $3
2533        RETURNING *
2534        "#,
2535    )
2536    .bind(job.id)
2537    .bind(serde_json::json!({
2538        "error": format!("cancelled: {}", reason),
2539        "attempt": job.attempt,
2540        "at": chrono::Utc::now().to_rfc3339()
2541    }))
2542    .bind(job.run_lease)
2543    .bind(progress_snapshot)
2544    .fetch_optional(&mut *tx)
2545    .await?;
2546
2547    let Some(updated_job) = updated_job else {
2548        tx.rollback().await?;
2549        return Ok(None);
2550    };
2551
2552    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(outcome_ctx))
2553        .await?;
2554
2555    tx.commit().await?;
2556    Ok(Some(updated_job))
2557}
2558
2559/// Park the running canonical row into `waiting_external` and dispatch any
2560/// registered `WaitingForCallback` follow-ups inside the same transaction.
2561/// Returns the parked [`JobRow`] for event emission, or `None` when the
2562/// guarded UPDATE matched zero rows (caller falls back to the existing race
2563/// / missing-callback handling).
2564async fn park_canonical_with_followups(
2565    pool: &PgPool,
2566    job: &JobRow,
2567    progress_snapshot: Option<&serde_json::Value>,
2568    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2569) -> Result<Option<JobRow>, AwaError> {
2570    let mut tx = pool.begin().await?;
2571
2572    let parked_job: Option<JobRow> = sqlx::query_as(
2573        r#"
2574        UPDATE awa.jobs_hot
2575        SET state = 'waiting_external',
2576            heartbeat_at = NULL,
2577            deadline_at = NULL,
2578            progress = $3
2579        WHERE id = $1 AND state = 'running' AND run_lease = $2 AND callback_id IS NOT NULL
2580        RETURNING *
2581        "#,
2582    )
2583    .bind(job.id)
2584    .bind(job.run_lease)
2585    .bind(progress_snapshot)
2586    .fetch_optional(&mut *tx)
2587    .await?;
2588
2589    let Some(parked_job) = parked_job else {
2590        tx.rollback().await?;
2591        return Ok(None);
2592    };
2593
2594    crate::enqueue_specs::dispatch_specs_in_tx(
2595        &mut tx,
2596        &parked_job,
2597        specs,
2598        Some(&crate::enqueue_specs::OutcomeContext::WaitingForCallback),
2599    )
2600    .await?;
2601
2602    tx.commit().await?;
2603    Ok(Some(parked_job))
2604}
2605
2606/// Caller-requested retry (RetryAfter) + Retried follow-ups.
2607///
2608/// `retryable` is forbidden in `awa.jobs_hot` (CHECK constraint); the row
2609/// must move to `awa.scheduled_jobs`. We do the DELETE+INSERT explicitly
2610/// in a CTE so the move plus the follow-up enqueue commit atomically. This
2611/// mirrors what the `INSTEAD OF UPDATE` trigger on `awa.jobs` would do —
2612/// avoiding the view keeps us correct under canonical-drain mode, where the
2613/// view's trigger rejects UPDATEs while queue-storage is the active
2614/// engine.
2615async fn retry_after_canonical_with_followups(
2616    pool: &PgPool,
2617    job: &JobRow,
2618    seconds: f64,
2619    progress_snapshot: Option<&serde_json::Value>,
2620    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2621) -> Result<Option<JobRow>, AwaError> {
2622    let mut tx = pool.begin().await?;
2623
2624    let outcome = awa_model::reschedule::reschedule_canonical_attempt_tx(
2625        &mut tx,
2626        job.id,
2627        job.run_lease,
2628        Reschedule::RetryAfter {
2629            delay_secs: seconds,
2630        },
2631        None,
2632        progress_snapshot,
2633    )
2634    .await?;
2635    let (job_id, run_at, attempt) = match outcome {
2636        RescheduleOutcome::Rescheduled {
2637            job_id,
2638            run_at,
2639            attempt,
2640        }
2641        | RescheduleOutcome::Migrated {
2642            job_id,
2643            run_at,
2644            attempt,
2645        } => (job_id, run_at, attempt),
2646        RescheduleOutcome::CancelledDuplicate { .. } => {
2647            tx.commit().await?;
2648            return Ok(None);
2649        }
2650        RescheduleOutcome::Stale => {
2651            tx.rollback().await?;
2652            return Ok(None);
2653        }
2654    };
2655    let updated_job = rescheduled_event_row(
2656        job,
2657        JobState::Retryable,
2658        job_id,
2659        run_at,
2660        attempt,
2661        None,
2662        progress_snapshot.cloned(),
2663    );
2664
2665    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
2666        error: String::new(),
2667        attempt: updated_job.attempt,
2668        next_run_at: updated_job.run_at,
2669    };
2670    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2671        .await?;
2672
2673    tx.commit().await?;
2674    Ok(Some(updated_job))
2675}
2676
2677/// Terminal error -> failed + Exhausted follow-ups.
2678async fn terminal_canonical_with_followups(
2679    pool: &PgPool,
2680    job: &JobRow,
2681    msg: &str,
2682    progress_snapshot: Option<&serde_json::Value>,
2683    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2684) -> Result<Option<JobRow>, AwaError> {
2685    let mut tx = pool.begin().await?;
2686
2687    let updated_job: Option<JobRow> = sqlx::query_as(
2688        r#"
2689        UPDATE awa.jobs_hot
2690        SET state = 'failed',
2691            finalized_at = now(),
2692            errors = errors || $2::jsonb,
2693            progress = $4
2694        WHERE id = $1 AND state = 'running' AND run_lease = $3
2695        RETURNING *
2696        "#,
2697    )
2698    .bind(job.id)
2699    .bind(serde_json::json!({
2700        "error": msg,
2701        "attempt": job.attempt,
2702        "at": chrono::Utc::now().to_rfc3339(),
2703        "terminal": true
2704    }))
2705    .bind(job.run_lease)
2706    .bind(progress_snapshot)
2707    .fetch_optional(&mut *tx)
2708    .await?;
2709
2710    let Some(updated_job) = updated_job else {
2711        tx.rollback().await?;
2712        return Ok(None);
2713    };
2714
2715    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
2716        error: msg.to_string(),
2717        attempt: job.attempt,
2718    };
2719    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2720        .await?;
2721
2722    tx.commit().await?;
2723    Ok(Some(updated_job))
2724}
2725
2726/// Retries exhausted (attempt >= max_attempts) -> failed + Exhausted
2727/// follow-ups. Differs from terminal in the error envelope (no `terminal`
2728/// marker since exhaustion is not a fatal-by-handler-assertion).
2729async fn exhaust_canonical_with_followups(
2730    pool: &PgPool,
2731    job: &JobRow,
2732    error_msg: &str,
2733    progress_snapshot: Option<&serde_json::Value>,
2734    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2735) -> Result<Option<JobRow>, AwaError> {
2736    let mut tx = pool.begin().await?;
2737
2738    let updated_job: Option<JobRow> = sqlx::query_as(
2739        r#"
2740        UPDATE awa.jobs_hot
2741        SET state = 'failed',
2742            finalized_at = now(),
2743            errors = errors || $2::jsonb,
2744            progress = $4
2745        WHERE id = $1 AND state = 'running' AND run_lease = $3
2746        RETURNING *
2747        "#,
2748    )
2749    .bind(job.id)
2750    .bind(serde_json::json!({
2751        "error": error_msg,
2752        "attempt": job.attempt,
2753        "at": chrono::Utc::now().to_rfc3339()
2754    }))
2755    .bind(job.run_lease)
2756    .bind(progress_snapshot)
2757    .fetch_optional(&mut *tx)
2758    .await?;
2759
2760    let Some(updated_job) = updated_job else {
2761        tx.rollback().await?;
2762        return Ok(None);
2763    };
2764
2765    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Exhausted {
2766        error: error_msg.to_string(),
2767        attempt: job.attempt,
2768    };
2769    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2770        .await?;
2771
2772    tx.commit().await?;
2773    Ok(Some(updated_job))
2774}
2775
2776/// Retryable error within budget -> retryable + Retried follow-ups, using
2777/// the configured backoff function for run_at.
2778async fn retry_backoff_canonical_with_followups(
2779    pool: &PgPool,
2780    job: &JobRow,
2781    error_msg: &str,
2782    progress_snapshot: Option<&serde_json::Value>,
2783    specs: &[crate::enqueue_specs::BoxedEnqueueSpec],
2784) -> Result<Option<JobRow>, AwaError> {
2785    let mut tx = pool.begin().await?;
2786
2787    let error_entry = serde_json::json!({
2788        "error": error_msg,
2789        "attempt": job.attempt,
2790        "at": chrono::Utc::now().to_rfc3339()
2791    });
2792    let outcome = awa_model::reschedule::reschedule_canonical_attempt_tx(
2793        &mut tx,
2794        job.id,
2795        job.run_lease,
2796        Reschedule::RetryBackoff,
2797        Some(&error_entry),
2798        progress_snapshot,
2799    )
2800    .await?;
2801    let (job_id, run_at, attempt) = match outcome {
2802        RescheduleOutcome::Rescheduled {
2803            job_id,
2804            run_at,
2805            attempt,
2806        }
2807        | RescheduleOutcome::Migrated {
2808            job_id,
2809            run_at,
2810            attempt,
2811        } => (job_id, run_at, attempt),
2812        RescheduleOutcome::CancelledDuplicate { .. } => {
2813            tx.commit().await?;
2814            return Ok(None);
2815        }
2816        RescheduleOutcome::Stale => {
2817            tx.rollback().await?;
2818            return Ok(None);
2819        }
2820    };
2821    let updated_job = rescheduled_event_row(
2822        job,
2823        JobState::Retryable,
2824        job_id,
2825        run_at,
2826        attempt,
2827        Some(&error_entry),
2828        progress_snapshot.cloned(),
2829    );
2830
2831    let outcome_ctx = crate::enqueue_specs::OutcomeContext::Retried {
2832        error: error_msg.to_string(),
2833        attempt: job.attempt,
2834        next_run_at: updated_job.run_at,
2835    };
2836    crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, Some(&outcome_ctx))
2837        .await?;
2838
2839    tx.commit().await?;
2840    Ok(Some(updated_job))
2841}
awa_worker/executor.rs

awa_worker/
executor.rs
