1use std::collections::HashMap;
14use std::sync::Arc;
15use std::time::Duration;
16
17use chrono::{DateTime, Utc};
18use rainier_container::Container;
19use rainier_support::{BoxedFuture, Error, Result};
20use serde::de::DeserializeOwned;
21use serde::{Deserialize, Serialize};
22
23pub struct JobContext {
25 container: Arc<Container>,
26 attempt: u32,
27 max_attempts: u32,
28 id: String,
29 queue: String,
30}
31
32impl JobContext {
33 pub fn new(
35 container: Arc<Container>,
36 id: impl Into<String>,
37 queue: impl Into<String>,
38 attempt: u32,
39 max_attempts: u32,
40 ) -> Self {
41 Self { container, attempt, max_attempts, id: id.into(), queue: queue.into() }
42 }
43
44 pub fn container(&self) -> &Arc<Container> {
49 &self.container
50 }
51
52 pub fn resolve<T: Send + Sync + 'static>(&self) -> Result<Arc<T>> {
54 self.container.resolve::<T>()
55 }
56
57 pub fn attempt(&self) -> u32 {
59 self.attempt
60 }
61
62 pub fn max_attempts(&self) -> u32 {
64 self.max_attempts
65 }
66
67 pub fn is_last_attempt(&self) -> bool {
70 self.attempt >= self.max_attempts
71 }
72
73 pub fn id(&self) -> &str {
75 &self.id
76 }
77
78 pub fn queue(&self) -> &str {
80 &self.queue
81 }
82}
83
84impl std::fmt::Debug for JobContext {
85 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86 f.debug_struct("JobContext")
87 .field("id", &self.id)
88 .field("queue", &self.queue)
89 .field("attempt", &format_args!("{}/{}", self.attempt, self.max_attempts))
90 .finish()
91 }
92}
93
94#[async_trait::async_trait]
116pub trait Job: Serialize + DeserializeOwned + Send + Sync + 'static {
117 const NAME: &'static str;
122
123 async fn handle(&self, context: &JobContext) -> Result<()>;
128
129 const TRIES: u32 = 3;
131
132 const QUEUE: &'static str = "default";
134
135 fn backoff(attempt: u32) -> Duration {
141 Duration::from_secs(1u64 << attempt.min(6))
142 }
143
144 async fn failed(&self, context: &JobContext, error: &Error) {
147 let _ = (context, error);
148 }
149
150 fn unique_id(&self) -> Option<String> {
188 None
189 }
190
191 const UNIQUE_FOR: Duration = Duration::from_secs(3600);
196}
197
198#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
200pub struct QueuedJob {
201 pub id: String,
203 pub name: String,
205 pub payload: serde_json::Value,
207 pub queue: String,
209 pub attempts: u32,
211 pub max_attempts: u32,
213 pub available_at: DateTime<Utc>,
215 pub created_at: DateTime<Utc>,
217 #[serde(default, skip_serializing_if = "Option::is_none")]
224 pub unique_key: Option<String>,
225}
226
227impl QueuedJob {
228 pub fn from_job<J: Job>(job: &J) -> Result<Self> {
230 let now = Utc::now();
231 Ok(Self {
232 id: generate_id(),
233 name: J::NAME.to_string(),
234 payload: serde_json::to_value(job)?,
235 queue: J::QUEUE.to_string(),
236 attempts: 0,
237 max_attempts: J::TRIES,
238 unique_key: None,
239 available_at: now,
240 created_at: now,
241 })
242 }
243
244 pub fn on_queue(mut self, queue: impl Into<String>) -> Self {
246 self.queue = queue.into();
247 self
248 }
249
250 pub fn delayed_by(mut self, delay: Duration) -> Self {
252 let delay = chrono::Duration::from_std(delay).unwrap_or_else(|_| chrono::Duration::zero());
253 self.available_at = Utc::now() + delay;
254 self
255 }
256
257 pub fn available_at(mut self, at: DateTime<Utc>) -> Self {
259 self.available_at = at;
260 self
261 }
262
263 pub fn is_available(&self) -> bool {
265 self.available_at <= Utc::now()
266 }
267
268 pub fn is_exhausted(&self) -> bool {
270 self.attempts >= self.max_attempts
271 }
272}
273
274fn generate_id() -> String {
280 use std::sync::atomic::{AtomicU64, Ordering};
281 static COUNTER: AtomicU64 = AtomicU64::new(0);
282
283 let micros = Utc::now().timestamp_micros();
284 let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
285 format!("{micros:x}-{sequence:x}")
286}
287
288type Runner =
290 Arc<dyn Fn(serde_json::Value, Arc<JobContext>) -> BoxedFuture<Result<()>> + Send + Sync>;
291
292#[derive(Default, Clone)]
298pub struct JobRegistry {
299 runners: HashMap<String, Runner>,
300}
301
302impl JobRegistry {
303 pub fn new() -> Self {
305 Self::default()
306 }
307
308 pub fn register<J: Job>(&mut self) -> &mut Self {
310 let runner: Runner = Arc::new(|payload, context| {
311 Box::pin(async move {
312 let job: J = serde_json::from_value(payload).map_err(|e| {
313 Error::internal(format!(
314 "could not deserialise a `{}` job — its payload and its type have \
315 diverged: {e}",
316 J::NAME
317 ))
318 })?;
319
320 let outcome = job.handle(&context).await;
321 if let Err(error) = &outcome {
322 if context.is_last_attempt() {
323 job.failed(&context, error).await;
324 }
325 }
326 outcome
327 })
328 });
329
330 self.runners.insert(J::NAME.to_string(), runner);
331 self
332 }
333
334 pub fn with<J: Job>(mut self) -> Self {
336 self.register::<J>();
337 self
338 }
339
340 pub fn knows(&self, name: &str) -> bool {
342 self.runners.contains_key(name)
343 }
344
345 pub fn names(&self) -> Vec<&str> {
347 let mut names: Vec<&str> = self.runners.keys().map(String::as_str).collect();
348 names.sort();
349 names
350 }
351
352 pub async fn run(&self, job: &QueuedJob, context: Arc<JobContext>) -> Result<()> {
354 let runner = self.runners.get(&job.name).ok_or_else(|| {
355 Error::internal(format!(
356 "no job is registered as `{}` — register it on the JobRegistry, or the worker \
357 cannot run it. Known jobs: {:?}",
358 job.name,
359 self.names()
360 ))
361 })?;
362
363 runner(job.payload.clone(), context).await
364 }
365}
366
367impl std::fmt::Debug for JobRegistry {
368 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
369 f.debug_struct("JobRegistry").field("jobs", &self.names()).finish()
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376 use std::sync::atomic::{AtomicU32, Ordering};
377 use std::sync::Mutex;
378
379 static RAN: Mutex<Vec<String>> = Mutex::new(Vec::new());
380 static FAILED: AtomicU32 = AtomicU32::new(0);
381
382 #[derive(Serialize, Deserialize)]
383 struct SendEmail {
384 to: String,
385 }
386
387 #[async_trait::async_trait]
388 impl Job for SendEmail {
389 const NAME: &'static str = "mail.send";
390
391 async fn handle(&self, _: &JobContext) -> Result<()> {
392 RAN.lock().unwrap().push(self.to.clone());
393 Ok(())
394 }
395 }
396
397 #[derive(Serialize, Deserialize)]
398 struct AlwaysFails;
399
400 #[async_trait::async_trait]
401 impl Job for AlwaysFails {
402 const NAME: &'static str = "test.fails";
403 const TRIES: u32 = 2;
404 const QUEUE: &'static str = "slow";
405
406 async fn handle(&self, _: &JobContext) -> Result<()> {
407 Err(Error::internal("nope"))
408 }
409
410 async fn failed(&self, _: &JobContext, _: &Error) {
411 FAILED.fetch_add(1, Ordering::SeqCst);
412 }
413 }
414
415 fn context(attempt: u32, max_attempts: u32) -> Arc<JobContext> {
416 Arc::new(JobContext::new(
417 Arc::new(Container::new()),
418 "job-1",
419 "default",
420 attempt,
421 max_attempts,
422 ))
423 }
424
425 #[test]
426 fn a_job_serialises_with_its_name_and_defaults() {
427 let queued = QueuedJob::from_job(&SendEmail { to: "ada@example.com".into() }).unwrap();
428
429 assert_eq!(queued.name, "mail.send");
430 assert_eq!(queued.queue, "default");
431 assert_eq!(queued.max_attempts, 3);
432 assert_eq!(queued.attempts, 0);
433 assert_eq!(queued.payload["to"], "ada@example.com");
434 assert!(queued.is_available());
435 assert!(!queued.is_exhausted());
436 }
437
438 #[test]
439 fn a_job_can_declare_its_own_queue_and_tries() {
440 let queued = QueuedJob::from_job(&AlwaysFails).unwrap();
441 assert_eq!(queued.queue, "slow");
442 assert_eq!(queued.max_attempts, 2);
443 }
444
445 #[test]
446 fn a_delay_pushes_availability_into_the_future() {
447 let queued = QueuedJob::from_job(&SendEmail { to: "a@b.c".into() })
448 .unwrap()
449 .delayed_by(Duration::from_secs(60));
450
451 assert!(!queued.is_available());
452 assert!(queued.available_at > Utc::now());
453 }
454
455 #[test]
456 fn queued_ids_are_unique() {
457 let a = QueuedJob::from_job(&SendEmail { to: "a".into() }).unwrap();
458 let b = QueuedJob::from_job(&SendEmail { to: "b".into() }).unwrap();
459 assert_ne!(a.id, b.id);
460 }
461
462 #[test]
463 fn backoff_grows_and_then_levels_off() {
464 assert_eq!(SendEmail::backoff(0), Duration::from_secs(1));
465 assert_eq!(SendEmail::backoff(1), Duration::from_secs(2));
466 assert_eq!(SendEmail::backoff(3), Duration::from_secs(8));
467 assert_eq!(SendEmail::backoff(20), Duration::from_secs(64));
469 }
470
471 #[tokio::test]
472 async fn the_registry_runs_a_job_by_name() {
473 RAN.lock().unwrap().clear();
474 let registry = JobRegistry::new().with::<SendEmail>();
475
476 let queued = QueuedJob::from_job(&SendEmail { to: "ada@example.com".into() }).unwrap();
477 registry.run(&queued, context(1, 3)).await.unwrap();
478
479 assert_eq!(*RAN.lock().unwrap(), vec!["ada@example.com"]);
480 }
481
482 #[tokio::test]
483 async fn an_unregistered_job_names_itself_and_what_is_known() {
484 let registry = JobRegistry::new().with::<SendEmail>();
485 let queued = QueuedJob::from_job(&AlwaysFails).unwrap();
486
487 let err = registry.run(&queued, context(1, 2)).await.unwrap_err();
488 assert!(err.message().contains("test.fails"), "{}", err.message());
489 assert!(err.message().contains("mail.send"), "{}", err.message());
490 }
491
492 #[tokio::test]
493 async fn a_payload_that_no_longer_fits_its_type_is_reported_clearly() {
494 let registry = JobRegistry::new().with::<SendEmail>();
495 let mut queued = QueuedJob::from_job(&SendEmail { to: "a".into() }).unwrap();
496 queued.payload = serde_json::json!({ "recipient": "a" });
497
498 let err = registry.run(&queued, context(1, 3)).await.unwrap_err();
499 assert!(err.message().contains("diverged"), "{}", err.message());
500 }
501
502 #[tokio::test]
503 async fn the_failed_hook_only_fires_on_the_last_attempt() {
504 FAILED.store(0, Ordering::SeqCst);
505 let registry = JobRegistry::new().with::<AlwaysFails>();
506 let queued = QueuedJob::from_job(&AlwaysFails).unwrap();
507
508 assert!(registry.run(&queued, context(1, 2)).await.is_err());
509 assert_eq!(FAILED.load(Ordering::SeqCst), 0, "not on the first attempt");
510
511 assert!(registry.run(&queued, context(2, 2)).await.is_err());
512 assert_eq!(FAILED.load(Ordering::SeqCst), 1, "but yes on the last");
513 }
514
515 #[test]
516 fn the_registry_lists_what_it_knows() {
517 let registry = JobRegistry::new().with::<SendEmail>().with::<AlwaysFails>();
518 assert_eq!(registry.names(), vec!["mail.send", "test.fails"]);
519 assert!(registry.knows("mail.send"));
520 assert!(!registry.knows("nope"));
521 }
522
523 #[test]
524 fn the_context_reports_the_attempt_window() {
525 let context = context(3, 3);
526 assert_eq!(context.attempt(), 3);
527 assert!(context.is_last_attempt());
528 assert!(!self::context(1, 3).is_last_attempt());
529 }
530}