Skip to main content

rust_web_server/jobs/
mod.rs

1//! Background job queue for one-shot, fire-and-forget work triggered from
2//! request handlers — e.g. "send a confirmation email after signup" — without
3//! blocking the response.
4//!
5//! Two flavors:
6//! - [`JobQueue`] — an in-memory worker pool. Simple and dependency-free, but
7//!   queued and retrying jobs are lost if the process crashes or restarts.
8//! - [`PersistentJobQueue`] (requires a `model-sqlite` / `model-postgres` /
9//!   `model-mysql` feature) — persists jobs via the model layer so pending
10//!   and retrying jobs survive a restart.
11//!
12//! # Example
13//!
14//! ```rust
15//! use rust_web_server::jobs::JobQueue;
16//!
17//! let queue = JobQueue::new(4); // 4 worker threads
18//! queue.submit(|| {
19//!     println!("sending welcome email...");
20//!     Ok(())
21//! });
22//! queue.join(); // wait for in-flight jobs to finish (tests/shutdown only)
23//! ```
24
25#[cfg(test)]
26mod tests;
27
28use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
29use std::sync::mpsc::{self, Sender};
30use std::sync::Arc;
31use std::thread;
32use std::time::Duration;
33
34// ── Job ──────────────────────────────────────────────────────────────────────
35
36/// A unit of background work.
37///
38/// Implement `run` directly for named, stateful jobs, or rely on the blanket
39/// implementation below to submit a plain closure.
40pub trait Job: Send {
41    /// Perform the work. Return `Err(message)` to trigger a retry (subject to
42    /// the queue's retry policy).
43    fn run(&self) -> Result<(), String>;
44
45    /// Human-readable name used in retry/failure log lines. Defaults to `"job"`.
46    fn name(&self) -> &str {
47        "job"
48    }
49}
50
51impl<F> Job for F
52where
53    F: Fn() -> Result<(), String> + Send,
54{
55    fn run(&self) -> Result<(), String> {
56        self()
57    }
58}
59
60struct Envelope {
61    job: Box<dyn Job>,
62    attempt: u32,
63}
64
65// ── JobQueue ─────────────────────────────────────────────────────────────────
66
67/// An in-memory background job queue backed by a fixed-size worker pool.
68///
69/// Jobs that return `Err` are retried with exponential backoff (default: 3
70/// retries, 500 ms initial backoff, 2x multiplier) — the same worker thread
71/// that picked up the job sleeps between attempts, so a job that retries
72/// repeatedly occupies one worker for the duration of its backoff; size the
73/// pool with that in mind. Queued and retrying jobs are lost if the process
74/// exits; see [`PersistentJobQueue`] for crash-safe persistence.
75pub struct JobQueue {
76    sender: Sender<Envelope>,
77    workers: Vec<thread::JoinHandle<()>>,
78    max_retries: Arc<AtomicU32>,
79    initial_backoff_ms: Arc<AtomicU64>,
80    backoff_multiplier: Arc<AtomicU32>,
81}
82
83impl JobQueue {
84    /// Create a queue backed by `workers` worker threads.
85    ///
86    /// Default retry policy: 3 retries, 500 ms initial backoff, 2x multiplier.
87    /// Adjust it with [`max_retries`](Self::max_retries) / [`backoff`](Self::backoff)
88    /// before submitting jobs.
89    pub fn new(workers: usize) -> Self {
90        assert!(workers > 0, "JobQueue requires at least one worker");
91
92        let (sender, receiver) = mpsc::channel::<Envelope>();
93        let receiver = Arc::new(std::sync::Mutex::new(receiver));
94
95        let max_retries = Arc::new(AtomicU32::new(3));
96        let initial_backoff_ms = Arc::new(AtomicU64::new(500));
97        let backoff_multiplier = Arc::new(AtomicU32::new(2));
98
99        let mut handles = Vec::with_capacity(workers);
100        for _ in 0..workers {
101            let receiver = Arc::clone(&receiver);
102            let max_retries = Arc::clone(&max_retries);
103            let initial_backoff_ms = Arc::clone(&initial_backoff_ms);
104            let backoff_multiplier = Arc::clone(&backoff_multiplier);
105            handles.push(thread::spawn(move || {
106                worker_loop(receiver, max_retries, initial_backoff_ms, backoff_multiplier);
107            }));
108        }
109
110        JobQueue {
111            sender,
112            workers: handles,
113            max_retries,
114            initial_backoff_ms,
115            backoff_multiplier,
116        }
117    }
118
119    /// Override the number of retries after the initial attempt (default: 3).
120    pub fn max_retries(self, n: u32) -> Self {
121        self.max_retries.store(n, Ordering::Relaxed);
122        self
123    }
124
125    /// Override the backoff policy: `initial` delay before the first retry,
126    /// doubled (or scaled by `multiplier`) on each subsequent attempt.
127    pub fn backoff(self, initial: Duration, multiplier: u32) -> Self {
128        self.initial_backoff_ms.store(initial.as_millis() as u64, Ordering::Relaxed);
129        self.backoff_multiplier.store(multiplier.max(1), Ordering::Relaxed);
130        self
131    }
132
133    /// Submit a job for background execution. Returns immediately.
134    pub fn submit(&self, job: impl Job + 'static) {
135        let _ = self.sender.send(Envelope { job: Box::new(job), attempt: 0 });
136    }
137
138    /// Stop accepting new jobs and wait for every worker to finish its
139    /// current job (including any retries already in progress). Consumes
140    /// the queue.
141    pub fn join(self) {
142        let JobQueue { sender, workers, .. } = self;
143        drop(sender);
144        for handle in workers {
145            let _ = handle.join();
146        }
147    }
148}
149
150fn worker_loop(
151    receiver: Arc<std::sync::Mutex<mpsc::Receiver<Envelope>>>,
152    max_retries: Arc<AtomicU32>,
153    initial_backoff_ms: Arc<AtomicU64>,
154    backoff_multiplier: Arc<AtomicU32>,
155) {
156    loop {
157        let received = {
158            let lock = receiver.lock().unwrap();
159            lock.recv()
160        };
161        let mut envelope = match received {
162            Ok(envelope) => envelope,
163            Err(_) => break, // sender dropped -> queue shut down
164        };
165
166        loop {
167            match envelope.job.run() {
168                Ok(()) => break,
169                Err(e) => {
170                    let max = max_retries.load(Ordering::Relaxed);
171                    if envelope.attempt >= max {
172                        eprintln!(
173                            "[jobs] '{}' failed permanently after {} attempt(s): {}",
174                            envelope.job.name(),
175                            envelope.attempt + 1,
176                            e
177                        );
178                        break;
179                    }
180
181                    let initial_ms = initial_backoff_ms.load(Ordering::Relaxed);
182                    let multiplier = backoff_multiplier.load(Ordering::Relaxed).max(1) as u64;
183                    let delay = Duration::from_millis(
184                        initial_ms.saturating_mul(multiplier.saturating_pow(envelope.attempt)),
185                    );
186                    eprintln!(
187                        "[jobs] '{}' failed (attempt {}/{}): {} — retrying in {:?}",
188                        envelope.job.name(),
189                        envelope.attempt + 1,
190                        max + 1,
191                        e,
192                        delay
193                    );
194                    thread::sleep(delay);
195                    envelope.attempt += 1;
196                }
197            }
198        }
199    }
200}
201
202// ── PersistentJobQueue ───────────────────────────────────────────────────────
203
204/// A job queue backed by the model layer (`model-sqlite` / `model-postgres` /
205/// `model-mysql`): jobs are written to a `rws_jobs` table before being
206/// acknowledged, so pending and retrying jobs survive a process crash or
207/// restart.
208///
209/// Unlike [`JobQueue`], a job here can't be an arbitrary closure — a closure
210/// can't be serialized to disk. Instead the queue stores `(job_type,
211/// payload)` as plain text and dispatches to a handler registered for that
212/// `job_type` at execution time. Register every handler you enqueue *before*
213/// calling [`start`](Self::start), including after a restart — unprocessed
214/// rows from a previous run (and any row left `running` when the process
215/// crashed) are picked up automatically.
216///
217/// ```sql
218/// CREATE TABLE IF NOT EXISTS rws_jobs (
219///     id          TEXT    PRIMARY KEY,
220///     job_type    TEXT    NOT NULL,
221///     payload     TEXT    NOT NULL DEFAULT '',
222///     status      TEXT    NOT NULL DEFAULT 'pending', -- pending | running | failed
223///     attempts    INTEGER NOT NULL DEFAULT 0,
224///     max_retries INTEGER NOT NULL DEFAULT 3,
225///     next_run_at INTEGER NOT NULL,
226///     created_at  INTEGER NOT NULL,
227///     last_error  TEXT
228/// )
229/// ```
230///
231/// Completed jobs are deleted; a job that exhausts its retries is left with
232/// `status = 'failed'` for inspection instead of being deleted.
233#[cfg(any(feature = "model-sqlite", feature = "model-postgres", feature = "model-mysql"))]
234pub struct PersistentJobQueue {
235    pool: crate::model::DbPool,
236    handlers: std::sync::RwLock<std::collections::HashMap<String, Arc<dyn Fn(&str) -> Result<(), String> + Send + Sync>>>,
237    default_max_retries: AtomicU32,
238    poll_interval_ms: AtomicU64,
239    initial_backoff_ms: AtomicU64,
240    backoff_multiplier: AtomicU32,
241    stop: Arc<std::sync::atomic::AtomicBool>,
242    id_counter: AtomicU64,
243}
244
245#[cfg(any(feature = "model-sqlite", feature = "model-postgres", feature = "model-mysql"))]
246impl PersistentJobQueue {
247    /// Creates the queue and its backing `rws_jobs` table if absent, then
248    /// resets any row left `running` (interrupted by a previous crash) back
249    /// to `pending` so it gets picked up again.
250    pub async fn new(pool: crate::model::DbPool) -> Result<Self, crate::model::DbError> {
251        let queue = PersistentJobQueue {
252            pool,
253            handlers: std::sync::RwLock::new(std::collections::HashMap::new()),
254            default_max_retries: AtomicU32::new(3),
255            poll_interval_ms: AtomicU64::new(500),
256            initial_backoff_ms: AtomicU64::new(500),
257            backoff_multiplier: AtomicU32::new(2),
258            stop: Arc::new(std::sync::atomic::AtomicBool::new(false)),
259            id_counter: AtomicU64::new(0),
260        };
261        queue.ensure_table().await?;
262        queue.recover_interrupted().await?;
263        Ok(queue)
264    }
265
266    async fn ensure_table(&self) -> Result<(), crate::model::DbError> {
267        self.pool
268            .execute(
269                "CREATE TABLE IF NOT EXISTS rws_jobs \
270                 (id TEXT PRIMARY KEY, job_type TEXT NOT NULL, payload TEXT NOT NULL DEFAULT '', \
271                  status TEXT NOT NULL DEFAULT 'pending', attempts INTEGER NOT NULL DEFAULT 0, \
272                  max_retries INTEGER NOT NULL DEFAULT 3, next_run_at INTEGER NOT NULL, \
273                  created_at INTEGER NOT NULL, last_error TEXT)",
274                &[],
275            )
276            .await?;
277        Ok(())
278    }
279
280    async fn recover_interrupted(&self) -> Result<(), crate::model::DbError> {
281        self.pool
282            .execute("UPDATE rws_jobs SET status = 'pending' WHERE status = 'running'", &[])
283            .await?;
284        Ok(())
285    }
286
287    /// Override the default retry count applied to jobs enqueued via
288    /// [`enqueue`](Self::enqueue) (default: 3). Use
289    /// [`enqueue_with_retries`](Self::enqueue_with_retries) for a per-job override.
290    pub fn max_retries(self, n: u32) -> Self {
291        self.default_max_retries.store(n, Ordering::Relaxed);
292        self
293    }
294
295    /// Override how often idle workers poll for due jobs (default: 500 ms).
296    pub fn poll_interval(self, interval: Duration) -> Self {
297        self.poll_interval_ms.store(interval.as_millis() as u64, Ordering::Relaxed);
298        self
299    }
300
301    /// Override the backoff policy applied between retries: `initial` delay
302    /// before the first retry, doubled (or scaled by `multiplier`) on each
303    /// subsequent attempt. Default: 500 ms initial, 2x multiplier.
304    pub fn backoff(self, initial: Duration, multiplier: u32) -> Self {
305        self.initial_backoff_ms.store(initial.as_millis() as u64, Ordering::Relaxed);
306        self.backoff_multiplier.store(multiplier.max(1), Ordering::Relaxed);
307        self
308    }
309
310    /// Register the function that executes jobs of `job_type`. Call before
311    /// [`start`](Self::start) — jobs enqueued for an unregistered `job_type`
312    /// fail immediately (and retry, same as any other handler error).
313    pub fn register(&self, job_type: &str, handler: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static) {
314        self.handlers.write().unwrap().insert(job_type.to_string(), Arc::new(handler));
315    }
316
317    fn generate_id(&self) -> String {
318        let nanos = std::time::SystemTime::now()
319            .duration_since(std::time::UNIX_EPOCH)
320            .map(|d| d.as_nanos() as u64)
321            .unwrap_or(0);
322        let count = self.id_counter.fetch_add(1, Ordering::Relaxed);
323        format!("{:x}-{:x}", nanos, count)
324    }
325
326    fn now_epoch() -> i64 {
327        std::time::SystemTime::now()
328            .duration_since(std::time::UNIX_EPOCH)
329            .map(|d| d.as_secs() as i64)
330            .unwrap_or(0)
331    }
332
333    /// Persist a new job using the queue's default retry count. Returns the
334    /// generated job id.
335    pub async fn enqueue(&self, job_type: &str, payload: &str) -> Result<String, crate::model::DbError> {
336        let max_retries = self.default_max_retries.load(Ordering::Relaxed);
337        self.enqueue_with_retries(job_type, payload, max_retries).await
338    }
339
340    /// Persist a new job with a per-job retry override.
341    pub async fn enqueue_with_retries(
342        &self,
343        job_type: &str,
344        payload: &str,
345        max_retries: u32,
346    ) -> Result<String, crate::model::DbError> {
347        use crate::model::Value;
348        let id = self.generate_id();
349        let now = Self::now_epoch();
350        self.pool
351            .execute(
352                "INSERT INTO rws_jobs (id, job_type, payload, status, attempts, max_retries, next_run_at, created_at) \
353                 VALUES (?, ?, ?, 'pending', 0, ?, ?, ?)",
354                &[
355                    Value::Text(id.clone()),
356                    Value::Text(job_type.to_string()),
357                    Value::Text(payload.to_string()),
358                    Value::Int(max_retries as i64),
359                    Value::Int(now),
360                    Value::Int(now),
361                ],
362            )
363            .await?;
364        Ok(id)
365    }
366
367    /// Atomically claims the oldest due job (`status = 'pending'` and
368    /// `next_run_at <= now`), marking it `running`. Returns `None` if no job
369    /// is due. The `UPDATE ... WHERE status = 'pending'` guards against two
370    /// workers claiming the same row.
371    async fn claim_next(&self) -> Result<Option<(String, String, String, i64, i64)>, crate::model::DbError> {
372        use crate::model::Value;
373        let now = Self::now_epoch();
374        let rows = self
375            .pool
376            .query_rows(
377                "SELECT id, job_type, payload, attempts, max_retries FROM rws_jobs \
378                 WHERE status = 'pending' AND next_run_at <= ? ORDER BY next_run_at ASC LIMIT 1",
379                &[Value::Int(now)],
380            )
381            .await?;
382        let row = match rows.into_iter().next() {
383            Some(r) => r,
384            None => return Ok(None),
385        };
386        let id: String = row.get("id")?;
387        let claimed = self
388            .pool
389            .execute(
390                "UPDATE rws_jobs SET status = 'running' WHERE id = ? AND status = 'pending'",
391                &[Value::Text(id.clone())],
392            )
393            .await?;
394        if claimed == 0 {
395            return Ok(None); // another worker won the race
396        }
397        Ok(Some((id, row.get("job_type")?, row.get("payload")?, row.get("attempts")?, row.get("max_retries")?)))
398    }
399
400    async fn mark_done(&self, id: &str) -> Result<(), crate::model::DbError> {
401        use crate::model::Value;
402        self.pool.execute("DELETE FROM rws_jobs WHERE id = ?", &[Value::Text(id.to_string())]).await?;
403        Ok(())
404    }
405
406    async fn mark_failed(&self, id: &str, error: &str) -> Result<(), crate::model::DbError> {
407        use crate::model::Value;
408        self.pool
409            .execute(
410                "UPDATE rws_jobs SET status = 'failed', last_error = ? WHERE id = ?",
411                &[Value::Text(error.to_string()), Value::Text(id.to_string())],
412            )
413            .await?;
414        Ok(())
415    }
416
417    async fn reschedule(&self, id: &str, attempts: i64, error: &str, delay: Duration) -> Result<(), crate::model::DbError> {
418        use crate::model::Value;
419        let next_run_at = Self::now_epoch() + delay.as_secs() as i64;
420        self.pool
421            .execute(
422                "UPDATE rws_jobs SET status = 'pending', attempts = ?, last_error = ?, next_run_at = ? WHERE id = ?",
423                &[Value::Int(attempts), Value::Text(error.to_string()), Value::Int(next_run_at), Value::Text(id.to_string())],
424            )
425            .await?;
426        Ok(())
427    }
428
429    /// Run one poll-claim-execute cycle. Returns `Ok(true)` if a job was
430    /// claimed (whether it succeeded, failed, or is retrying), `Ok(false)` if
431    /// none was due. Exposed for tests and for embedding in a caller-owned
432    /// loop; most applications should use [`start`](Self::start) instead.
433    pub async fn tick(&self) -> Result<bool, crate::model::DbError> {
434        let (id, job_type, payload, attempts, max_retries) = match self.claim_next().await? {
435            Some(claimed) => claimed,
436            None => return Ok(false),
437        };
438
439        let handler = self.handlers.read().unwrap().get(&job_type).cloned();
440        let result = match handler {
441            Some(h) => h(&payload),
442            None => Err(format!("no handler registered for job_type '{}'", job_type)),
443        };
444
445        match result {
446            Ok(()) => self.mark_done(&id).await?,
447            Err(e) => {
448                // `max_retries` counts retries *after* the first attempt, so
449                // `max_retries + 1` total attempts are allowed before giving
450                // up — matching `JobQueue`'s semantics.
451                let next_attempts = attempts + 1;
452                if next_attempts > max_retries {
453                    eprintln!(
454                        "[jobs] persistent job '{}' ({}) failed permanently after {} attempt(s): {}",
455                        job_type, id, next_attempts, e
456                    );
457                    self.mark_failed(&id, &e).await?;
458                } else {
459                    let initial_ms = self.initial_backoff_ms.load(Ordering::Relaxed);
460                    let multiplier = self.backoff_multiplier.load(Ordering::Relaxed).max(1) as u64;
461                    let delay = Duration::from_millis(
462                        initial_ms.saturating_mul(multiplier.saturating_pow(attempts as u32)),
463                    );
464                    eprintln!(
465                        "[jobs] persistent job '{}' ({}) failed (attempt {}/{}): {} — retrying in {:?}",
466                        job_type, id, next_attempts, max_retries + 1, e, delay
467                    );
468                    self.reschedule(&id, next_attempts, &e, delay).await?;
469                }
470            }
471        }
472        Ok(true)
473    }
474
475    /// Spawn `workers` background OS threads, each running its own
476    /// single-threaded Tokio runtime, polling for due jobs at
477    /// [`poll_interval`](Self::poll_interval) (default: 500 ms) whenever the
478    /// queue is empty. Returns the thread handles; call
479    /// [`stop`](Self::stop) then join them for a graceful shutdown.
480    pub fn start(self: Arc<Self>, workers: usize) -> Vec<thread::JoinHandle<()>> {
481        assert!(workers > 0, "PersistentJobQueue requires at least one worker");
482        (0..workers)
483            .map(|_| {
484                let queue = Arc::clone(&self);
485                thread::spawn(move || {
486                    let rt = tokio::runtime::Builder::new_current_thread()
487                        .enable_all()
488                        .build()
489                        .expect("failed to start tokio runtime for PersistentJobQueue worker");
490                    rt.block_on(async {
491                        loop {
492                            if queue.stop.load(Ordering::Relaxed) {
493                                break;
494                            }
495                            let poll_interval = Duration::from_millis(queue.poll_interval_ms.load(Ordering::Relaxed));
496                            match queue.tick().await {
497                                Ok(true) => {}
498                                Ok(false) => tokio::time::sleep(poll_interval).await,
499                                Err(e) => {
500                                    eprintln!("[jobs] persistent queue poll error: {}", e);
501                                    tokio::time::sleep(poll_interval).await;
502                                }
503                            }
504                        }
505                    });
506                })
507            })
508            .collect()
509    }
510
511    /// Signal all worker threads spawned by [`start`](Self::start) to stop
512    /// after their current tick. Does not itself join the returned thread
513    /// handles — join them separately to wait for a clean shutdown.
514    pub fn stop(&self) {
515        self.stop.store(true, Ordering::Relaxed);
516    }
517}