1use 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#[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 pub fn new(args: F) -> Self {
43 Self {
44 args,
45 opts: InsertOpts::default(),
46 }
47 }
48
49 pub fn queue(mut self, queue: impl Into<String>) -> Self {
51 self.opts.queue = queue.into();
52 self
53 }
54
55 pub fn priority(mut self, priority: i16) -> Self {
57 self.opts.priority = priority;
58 self
59 }
60
61 pub fn max_attempts(mut self, max_attempts: i16) -> Self {
63 self.opts.max_attempts = max_attempts;
64 self
65 }
66
67 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
92pub enum Outcome {
93 Completed,
94 Retried,
95 Exhausted,
96 Cancelled,
97 WaitingForCallback,
98 Rescued,
101}
102
103#[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
127pub(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
153fn 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
183pub(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
215pub(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
260pub(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
295pub(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
330pub(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
360pub(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
397pub(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}