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

awa_worker/
executor.rs
