rainier_queue/
job.rs

1//! Jobs — the [`Job`] contract, the serialised [`QueuedJob`] payload, and the
2//! [`JobRegistry`] that turns one back into the other.
3//!
4//! A queued job crosses a process boundary: it is written by a web request and
5//! read, perhaps minutes later, by a worker that may be a different process on
6//! a different machine. So a job is defined by two things — a **payload** that
7//! serialises, and a **name** that survives serialisation.
8//!
9//! The name is explicit rather than derived from the Rust type path, because
10//! the type path is not a stable identifier: renaming a struct or moving it
11//! between modules would strand every job already in the queue.
12
13use 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
23/// What a job can reach while it runs.
24pub struct JobContext {
25    container: Arc<Container>,
26    attempt: u32,
27    max_attempts: u32,
28    id: String,
29    queue: String,
30}
31
32impl JobContext {
33    /// Build a context.
34    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    /// The service container, for resolving whatever the job needs.
45    ///
46    /// A job cannot capture its dependencies — it was serialised — so it
47    /// resolves them here instead.
48    pub fn container(&self) -> &Arc<Container> {
49        &self.container
50    }
51
52    /// Resolve a service.
53    pub fn resolve<T: Send + Sync + 'static>(&self) -> Result<Arc<T>> {
54        self.container.resolve::<T>()
55    }
56
57    /// Which attempt this is, starting at 1.
58    pub fn attempt(&self) -> u32 {
59        self.attempt
60    }
61
62    /// How many attempts the job gets in total.
63    pub fn max_attempts(&self) -> u32 {
64        self.max_attempts
65    }
66
67    /// Whether this is the last attempt — a job that wants to record a
68    /// permanent failure can check this before returning `Err`.
69    pub fn is_last_attempt(&self) -> bool {
70        self.attempt >= self.max_attempts
71    }
72
73    /// The queued job's id.
74    pub fn id(&self) -> &str {
75        &self.id
76    }
77
78    /// The queue it came from.
79    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/// Work to be done later.
95///
96/// ```
97/// use rainier_queue::{Job, JobContext};
98/// use serde::{Deserialize, Serialize};
99///
100/// #[derive(Serialize, Deserialize)]
101/// struct SendWelcomeEmail {
102///     user_id: u64,
103/// }
104///
105/// #[async_trait::async_trait]
106/// impl Job for SendWelcomeEmail {
107///     const NAME: &'static str = "mail.welcome";
108///
109///     async fn handle(&self, _context: &JobContext) -> rainier_support::Result<()> {
110///         println!("emailing user {}", self.user_id);
111///         Ok(())
112///     }
113/// }
114/// ```
115#[async_trait::async_trait]
116pub trait Job: Serialize + DeserializeOwned + Send + Sync + 'static {
117    /// The job's stable name on the wire.
118    ///
119    /// Must not change once jobs of this type exist in a queue, and must be
120    /// unique across the application.
121    const NAME: &'static str;
122
123    /// Do the work.
124    ///
125    /// Returning `Err` releases the job for another attempt, until
126    /// [`TRIES`](Job::TRIES) is exhausted and it is marked failed.
127    async fn handle(&self, context: &JobContext) -> Result<()>;
128
129    /// How many times to attempt the job before giving up.
130    const TRIES: u32 = 3;
131
132    /// The queue this job goes on by default.
133    const QUEUE: &'static str = "default";
134
135    /// How long to wait before retrying after `attempt` failures.
136    ///
137    /// Exponential by default — 1s, 2s, 4s, … — because the usual reason a job
138    /// fails is a dependency that is briefly unavailable, and retrying at full
139    /// speed makes that worse.
140    fn backoff(attempt: u32) -> Duration {
141        Duration::from_secs(1u64 << attempt.min(6))
142    }
143
144    /// Called after the final attempt fails. For recording the failure
145    /// somewhere the application will notice.
146    async fn failed(&self, context: &JobContext, error: &Error) {
147        let _ = (context, error);
148    }
149
150    /// What makes two of these job "the same one", or `None` for jobs that are
151    /// never duplicates.
152    ///
153    /// Returning `Some` makes dispatch take an
154    /// [atomic lock](rainier_cache::Lock) on this id, so a second dispatch
155    /// while the first is still pending is **dropped rather than queued**.
156    ///
157    /// ```ignore
158    /// impl Job for RebuildSearchIndex {
159    ///     const NAME: &'static str = "search.rebuild";
160    ///
161    ///     // One rebuild at a time, whatever asks for it.
162    ///     fn unique_id(&self) -> Option<String> {
163    ///         Some(String::new())
164    ///     }
165    /// }
166    ///
167    /// impl Job for SendInvoice {
168    ///     const NAME: &'static str = "billing.invoice";
169    ///
170    ///     // One per invoice, but two different invoices are two jobs.
171    ///     fn unique_id(&self) -> Option<String> {
172    ///         Some(self.invoice_id.to_string())
173    ///     }
174    /// }
175    /// ```
176    ///
177    /// The lock key is `NAME` **and** this id, so two job types returning the
178    /// same id do not collide.
179    ///
180    /// # This deduplicates dispatch, not execution
181    ///
182    /// The lock is released when the job finishes, so it stops a *queue* filling
183    /// with a hundred copies of the same work. It does not stop two workers
184    /// running one job twice — that is the queue's reservation, a different
185    /// mechanism — and it does not survive the lock expiring under a job that
186    /// overran [`UNIQUE_FOR`](Job::UNIQUE_FOR).
187    fn unique_id(&self) -> Option<String> {
188        None
189    }
190
191    /// How long the uniqueness lock is held if the job never finishes.
192    ///
193    /// The safety net for a worker that dies mid-job: without it, one crash
194    /// would block that job's id forever. Set it longer than the job takes.
195    const UNIQUE_FOR: Duration = Duration::from_secs(3600);
196}
197
198/// A job serialised for the queue.
199#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
200pub struct QueuedJob {
201    /// A unique id for this queued instance.
202    pub id: String,
203    /// The [`Job::NAME`] the registry resolves.
204    pub name: String,
205    /// The job's serialised body.
206    pub payload: serde_json::Value,
207    /// Which queue it is on.
208    pub queue: String,
209    /// How many times it has been attempted.
210    pub attempts: u32,
211    /// How many attempts it gets.
212    pub max_attempts: u32,
213    /// The earliest it may run — how delays are expressed.
214    pub available_at: DateTime<Utc>,
215    /// When it was enqueued.
216    pub created_at: DateTime<Utc>,
217    /// The uniqueness lock this job holds, if it declared a
218    /// [`unique_id`](Job::unique_id).
219    ///
220    /// Carried on the job rather than kept by the dispatcher, because the
221    /// process that releases it is a **worker**, in another process, minutes
222    /// later. `None` for the ordinary case.
223    #[serde(default, skip_serializing_if = "Option::is_none")]
224    pub unique_key: Option<String>,
225}
226
227impl QueuedJob {
228    /// Serialise `job` for the queue.
229    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    /// Put it on a different queue.
245    pub fn on_queue(mut self, queue: impl Into<String>) -> Self {
246        self.queue = queue.into();
247        self
248    }
249
250    /// Hold it back for `delay`.
251    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    /// Make it available no earlier than `at`.
258    pub fn available_at(mut self, at: DateTime<Utc>) -> Self {
259        self.available_at = at;
260        self
261    }
262
263    /// Whether it may run now.
264    pub fn is_available(&self) -> bool {
265        self.available_at <= Utc::now()
266    }
267
268    /// Whether every attempt has been used.
269    pub fn is_exhausted(&self) -> bool {
270        self.attempts >= self.max_attempts
271    }
272}
273
274/// A unique id for a queued job: a timestamp for rough ordering plus
275/// randomness for uniqueness.
276///
277/// Not a UUID, to avoid the dependency — the requirement is only that two ids
278/// minted in the same process at the same moment differ.
279fn 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
288/// Runs a [`QueuedJob`] whose concrete type has been erased.
289type Runner =
290    Arc<dyn Fn(serde_json::Value, Arc<JobContext>) -> BoxedFuture<Result<()>> + Send + Sync>;
291
292/// Maps a job's wire name back to code that can run it.
293///
294/// A worker reads `{"name": "mail.welcome", …}` and has no idea what type that
295/// is; the registry is what closes the gap. Every job type an application
296/// dispatches must be registered, or its jobs will fail as unknown.
297#[derive(Default, Clone)]
298pub struct JobRegistry {
299    runners: HashMap<String, Runner>,
300}
301
302impl JobRegistry {
303    /// An empty registry.
304    pub fn new() -> Self {
305        Self::default()
306    }
307
308    /// Register `J` so the worker can run it.
309    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    /// Builder form of [`register`](Self::register).
335    pub fn with<J: Job>(mut self) -> Self {
336        self.register::<J>();
337        self
338    }
339
340    /// Whether `name` is registered.
341    pub fn knows(&self, name: &str) -> bool {
342        self.runners.contains_key(name)
343    }
344
345    /// Every registered job name, sorted.
346    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    /// Run a queued job.
353    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        // Capped, so a long-failing job does not schedule itself years out.
468        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}
rainier_queue/job.rs

rainier_queue/
job.rs
