Skip to main content

awa_worker/
enqueue_specs.rs

1//! Transactional follow-up enqueue specs (ADR-029).
2//!
3//! A spec is a per-(outcome, kind) registration that fires when its
4//! triggering state transition commits. Worker-driven outcomes and
5//! callback resolution via the worker `Client::*_external` APIs
6//! dispatch the follow-up `INSERT` in the same transaction as the
7//! state UPDATE (atomic with the trigger — a spec failure rolls the
8//! trigger back so the caller / external sender can retry). Maintenance
9//! rescue dispatches in a separate transaction after the rescue
10//! commits (best-effort — a failed `INSERT` is logged and the rescue
11//! stands). Either way, once the follow-up `INSERT` commits the row is
12//! a regular Awa job: at-least-once, retried, DLQ-aware, visible to
13//! admin tooling.
14//!
15//! Specs are type-erased here so the executor can dispatch them without
16//! knowing the trigger or follow-up types statically. The user-facing
17//! `ClientBuilder::on_*_enqueue` methods wrap their typed closures into
18//! impls of [`EnqueueFollowUp`] and accumulate them in a two-level
19//! `outcome -> kind -> specs` registry.
20
21use awa_model::{insert_with, AwaError, InsertOpts, JobArgs, JobRow};
22use serde::de::DeserializeOwned;
23use sqlx::PgConnection;
24use std::future::Future;
25use std::marker::PhantomData;
26use std::pin::Pin;
27use std::sync::Arc;
28
29/// Description of a follow-up job to enqueue when an `on_*_enqueue` spec
30/// fires. Carries the follow-up's `JobArgs` and the [`InsertOpts`] applied to
31/// the `INSERT`. For the common "default opts" case, users can return the
32/// `JobArgs` value directly from their closure — `EnqueueRequest::from(args)`
33/// is invoked automatically.
34#[derive(Debug, Clone)]
35pub struct EnqueueRequest<F> {
36    pub(crate) args: F,
37    pub(crate) opts: InsertOpts,
38}
39
40impl<F: JobArgs> EnqueueRequest<F> {
41    /// Build a request with default [`InsertOpts`].
42    pub fn new(args: F) -> Self {
43        Self {
44            args,
45            opts: InsertOpts::default(),
46        }
47    }
48
49    /// Override the follow-up's queue.
50    pub fn queue(mut self, queue: impl Into<String>) -> Self {
51        self.opts.queue = queue.into();
52        self
53    }
54
55    /// Override the follow-up's priority.
56    pub fn priority(mut self, priority: i16) -> Self {
57        self.opts.priority = priority;
58        self
59    }
60
61    /// Override the follow-up's `max_attempts`.
62    pub fn max_attempts(mut self, max_attempts: i16) -> Self {
63        self.opts.max_attempts = max_attempts;
64        self
65    }
66
67    /// Replace the follow-up's [`InsertOpts`] wholesale — useful for fields
68    /// without dedicated builder methods (`metadata`, `tags`, `unique`,
69    /// `run_at`, `deadline_duration`, `ordering_key`).
70    pub fn with_opts(mut self, opts: InsertOpts) -> Self {
71        self.opts = opts;
72        self
73    }
74}
75
76impl<F: JobArgs> From<F> for EnqueueRequest<F> {
77    fn from(args: F) -> Self {
78        Self::new(args)
79    }
80}
81
82/// The outcome whose state-commit triggers a registered spec.
83///
84/// One spec is tied to exactly one outcome; the registry is keyed on this so
85/// the executor can look up specs for the specific branch it just took
86/// without filtering.
87///
88/// `Started` is intentionally excluded (see ADR-029): claim-time follow-up
89/// enqueue would join the dispatcher's hot path and the durable-side-effect
90/// use case for "job started" is uncommon. Observation belongs to hooks.
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
92pub enum Outcome {
93    Completed,
94    Retried,
95    Exhausted,
96    Cancelled,
97    WaitingForCallback,
98    /// Maintenance rescued the job (expired callback, stale heartbeat, or
99    /// exceeded deadline). See [`crate::events::RescueReason`].
100    Rescued,
101}
102
103/// Per-outcome runtime context handed to non-`Completed` follow-up closures
104/// so they can specialise on outcome-specific fields (error / attempt /
105/// reason / next_run_at). Variants mirror the corresponding
106/// `UntypedJobEvent` variants in shape.
107#[derive(Debug, Clone)]
108pub enum OutcomeContext {
109    Retried {
110        error: String,
111        attempt: i16,
112        next_run_at: chrono::DateTime<chrono::Utc>,
113    },
114    Exhausted {
115        error: String,
116        attempt: i16,
117    },
118    Cancelled {
119        reason: String,
120    },
121    WaitingForCallback,
122    Rescued {
123        reason: crate::events::RescueReason,
124    },
125}
126
127/// Type-erased follow-up-enqueue spec for one (outcome, kind) pair.
128///
129/// `Completed` specs receive `outcome_context: None` (no extra context beyond
130/// the JobRow). Non-Completed specs receive `Some(ctx)` with the matching
131/// variant — the registry guarantees the variant matches because the spec is
132/// registered under its outcome key.
133pub(crate) trait EnqueueFollowUp: Send + Sync {
134    fn run<'a>(
135        &'a self,
136        conn: &'a mut PgConnection,
137        job: &'a JobRow,
138        outcome_context: Option<&'a OutcomeContext>,
139    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>>;
140}
141
142pub(crate) type BoxedEnqueueSpec = Arc<dyn EnqueueFollowUp + 'static>;
143
144fn decode_trigger_args<T: DeserializeOwned>(job: &JobRow) -> Result<T, AwaError> {
145    serde_json::from_value(job.args.clone()).map_err(|err| {
146        AwaError::Validation(format!(
147            "follow-up enqueue: failed to decode trigger args for kind {}: {err}",
148            job.kind
149        ))
150    })
151}
152
153/// Invoke a user-supplied `make` closure, converting any panic into an
154/// `AwaError` so the spec dispatcher can roll back its transaction and
155/// log the failure the same way it would for a returned error. Without
156/// this, a panic in user code would unwind:
157/// - on worker-driven outcomes, the executor's completion task (the
158///   transition tx would still roll back via Drop, but the worker would
159///   log a JoinError rather than the typed spec failure);
160/// - on callback-resolution paths, the resolver task (after the
161///   resolution itself has already committed) — observers and in-process
162///   hooks downstream of the panic point never run;
163/// - on rescue paths, the maintenance task — likewise dropping the
164///   detached hook spawn for that rescue.
165fn catch_make_panic<R>(
166    kind: &str,
167    f: impl FnOnce() -> R + std::panic::UnwindSafe,
168) -> Result<R, AwaError> {
169    std::panic::catch_unwind(f).map_err(|panic| {
170        let detail = if let Some(msg) = panic.downcast_ref::<&'static str>() {
171            (*msg).to_string()
172        } else if let Some(msg) = panic.downcast_ref::<String>() {
173            msg.clone()
174        } else {
175            "panic payload not a string".to_string()
176        };
177        AwaError::Validation(format!(
178            "follow-up enqueue closure for kind {kind} panicked: {detail}"
179        ))
180    })
181}
182
183/// Spec for the `Completed` outcome. Captures a typed closure that maps the
184/// trigger's deserialised args plus its post-completion `JobRow` to an
185/// [`EnqueueRequest<F>`] describing the follow-up.
186pub(crate) struct CompletedFollowUp<T, F, MakeFn> {
187    pub(crate) make: MakeFn,
188    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
189}
190
191impl<T, F, MakeFn> EnqueueFollowUp for CompletedFollowUp<T, F, MakeFn>
192where
193    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
194    F: JobArgs + Send + Sync + 'static,
195    MakeFn: Fn(T, &JobRow) -> EnqueueRequest<F> + Send + Sync + 'static,
196{
197    fn run<'a>(
198        &'a self,
199        conn: &'a mut PgConnection,
200        job: &'a JobRow,
201        _outcome_context: Option<&'a OutcomeContext>,
202    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
203        Box::pin(async move {
204            let args: T = decode_trigger_args(job)?;
205            let request = catch_make_panic(
206                &job.kind,
207                std::panic::AssertUnwindSafe(|| (self.make)(args, job)),
208            )?;
209            insert_with(&mut *conn, &request.args, request.opts).await?;
210            Ok(())
211        })
212    }
213}
214
215/// Spec for the `Retried` outcome.
216pub(crate) struct RetriedFollowUp<T, F, MakeFn> {
217    pub(crate) make: MakeFn,
218    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
219}
220
221impl<T, F, MakeFn> EnqueueFollowUp for RetriedFollowUp<T, F, MakeFn>
222where
223    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
224    F: JobArgs + Send + Sync + 'static,
225    MakeFn: Fn(T, &JobRow, &str, i16, chrono::DateTime<chrono::Utc>) -> EnqueueRequest<F>
226        + Send
227        + Sync
228        + 'static,
229{
230    fn run<'a>(
231        &'a self,
232        conn: &'a mut PgConnection,
233        job: &'a JobRow,
234        outcome_context: Option<&'a OutcomeContext>,
235    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
236        Box::pin(async move {
237            let Some(OutcomeContext::Retried {
238                error,
239                attempt,
240                next_run_at,
241            }) = outcome_context
242            else {
243                return Err(AwaError::Validation(
244                    "RetriedFollowUp dispatched without a Retried OutcomeContext".into(),
245                ));
246            };
247            let args: T = decode_trigger_args(job)?;
248            let request = catch_make_panic(
249                &job.kind,
250                std::panic::AssertUnwindSafe(|| {
251                    (self.make)(args, job, error, *attempt, *next_run_at)
252                }),
253            )?;
254            insert_with(&mut *conn, &request.args, request.opts).await?;
255            Ok(())
256        })
257    }
258}
259
260/// Spec for the `Exhausted` outcome (retries-exhausted or terminal-error).
261pub(crate) struct ExhaustedFollowUp<T, F, MakeFn> {
262    pub(crate) make: MakeFn,
263    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
264}
265
266impl<T, F, MakeFn> EnqueueFollowUp for ExhaustedFollowUp<T, F, MakeFn>
267where
268    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
269    F: JobArgs + Send + Sync + 'static,
270    MakeFn: Fn(T, &JobRow, &str, i16) -> EnqueueRequest<F> + Send + Sync + 'static,
271{
272    fn run<'a>(
273        &'a self,
274        conn: &'a mut PgConnection,
275        job: &'a JobRow,
276        outcome_context: Option<&'a OutcomeContext>,
277    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
278        Box::pin(async move {
279            let Some(OutcomeContext::Exhausted { error, attempt }) = outcome_context else {
280                return Err(AwaError::Validation(
281                    "ExhaustedFollowUp dispatched without an Exhausted OutcomeContext".into(),
282                ));
283            };
284            let args: T = decode_trigger_args(job)?;
285            let request = catch_make_panic(
286                &job.kind,
287                std::panic::AssertUnwindSafe(|| (self.make)(args, job, error, *attempt)),
288            )?;
289            insert_with(&mut *conn, &request.args, request.opts).await?;
290            Ok(())
291        })
292    }
293}
294
295/// Spec for the `Cancelled` outcome.
296pub(crate) struct CancelledFollowUp<T, F, MakeFn> {
297    pub(crate) make: MakeFn,
298    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
299}
300
301impl<T, F, MakeFn> EnqueueFollowUp for CancelledFollowUp<T, F, MakeFn>
302where
303    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
304    F: JobArgs + Send + Sync + 'static,
305    MakeFn: Fn(T, &JobRow, &str) -> EnqueueRequest<F> + Send + Sync + 'static,
306{
307    fn run<'a>(
308        &'a self,
309        conn: &'a mut PgConnection,
310        job: &'a JobRow,
311        outcome_context: Option<&'a OutcomeContext>,
312    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
313        Box::pin(async move {
314            let Some(OutcomeContext::Cancelled { reason }) = outcome_context else {
315                return Err(AwaError::Validation(
316                    "CancelledFollowUp dispatched without a Cancelled OutcomeContext".into(),
317                ));
318            };
319            let args: T = decode_trigger_args(job)?;
320            let request = catch_make_panic(
321                &job.kind,
322                std::panic::AssertUnwindSafe(|| (self.make)(args, job, reason)),
323            )?;
324            insert_with(&mut *conn, &request.args, request.opts).await?;
325            Ok(())
326        })
327    }
328}
329
330/// Spec for the `WaitingForCallback` outcome.
331pub(crate) struct WaitingForCallbackFollowUp<T, F, MakeFn> {
332    pub(crate) make: MakeFn,
333    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
334}
335
336impl<T, F, MakeFn> EnqueueFollowUp for WaitingForCallbackFollowUp<T, F, MakeFn>
337where
338    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
339    F: JobArgs + Send + Sync + 'static,
340    MakeFn: Fn(T, &JobRow) -> EnqueueRequest<F> + Send + Sync + 'static,
341{
342    fn run<'a>(
343        &'a self,
344        conn: &'a mut PgConnection,
345        job: &'a JobRow,
346        _outcome_context: Option<&'a OutcomeContext>,
347    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
348        Box::pin(async move {
349            let args: T = decode_trigger_args(job)?;
350            let request = catch_make_panic(
351                &job.kind,
352                std::panic::AssertUnwindSafe(|| (self.make)(args, job)),
353            )?;
354            insert_with(&mut *conn, &request.args, request.opts).await?;
355            Ok(())
356        })
357    }
358}
359
360/// Spec for the `Rescued` outcome. `make` receives the trigger's args,
361/// post-rescue `JobRow`, and the [`RescueReason`](crate::events::RescueReason).
362pub(crate) struct RescuedFollowUp<T, F, MakeFn> {
363    pub(crate) make: MakeFn,
364    pub(crate) _phantom: PhantomData<fn() -> (T, F)>,
365}
366
367impl<T, F, MakeFn> EnqueueFollowUp for RescuedFollowUp<T, F, MakeFn>
368where
369    T: JobArgs + DeserializeOwned + Send + Sync + 'static,
370    F: JobArgs + Send + Sync + 'static,
371    MakeFn:
372        Fn(T, &JobRow, crate::events::RescueReason) -> EnqueueRequest<F> + Send + Sync + 'static,
373{
374    fn run<'a>(
375        &'a self,
376        conn: &'a mut PgConnection,
377        job: &'a JobRow,
378        outcome_context: Option<&'a OutcomeContext>,
379    ) -> Pin<Box<dyn Future<Output = Result<(), AwaError>> + Send + 'a>> {
380        Box::pin(async move {
381            let Some(OutcomeContext::Rescued { reason }) = outcome_context else {
382                return Err(AwaError::Validation(
383                    "RescuedFollowUp dispatched without a Rescued OutcomeContext".into(),
384                ));
385            };
386            let args: T = decode_trigger_args(job)?;
387            let request = catch_make_panic(
388                &job.kind,
389                std::panic::AssertUnwindSafe(|| (self.make)(args, job, *reason)),
390            )?;
391            insert_with(&mut *conn, &request.args, request.opts).await?;
392            Ok(())
393        })
394    }
395}
396
397/// Helper used by the executor (and other emission sites) to drive a list of
398/// specs against a connection inside an already-open transaction.
399pub(crate) async fn dispatch_specs_in_tx(
400    conn: &mut PgConnection,
401    job: &JobRow,
402    specs: &[BoxedEnqueueSpec],
403    outcome_context: Option<&OutcomeContext>,
404) -> Result<(), AwaError> {
405    for spec in specs {
406        spec.run(conn, job, outcome_context).await?;
407    }
408    Ok(())
409}