#[cfg(feature = "jobs-postgres")]
pub mod pg;
#[cfg(feature = "jobs-postgres")]
pub type DatabaseJobQueue = pg::PgJobQueue;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use serde::{de::DeserializeOwned, Serialize};
use tokio::sync::{mpsc, watch, Mutex};
use tokio::task::{JoinHandle, JoinSet};
#[derive(Debug, thiserror::Error)]
pub enum JobError {
#[error("retryable: {0}")]
Retryable(String),
#[error("fatal: {0}")]
Fatal(String),
#[error("queue error: {0}")]
Queue(String),
}
#[async_trait::async_trait]
pub trait Job: Send + Sync + Sized + Serialize + DeserializeOwned + 'static {
const NAME: &'static str;
const MAX_ATTEMPTS: u32 = 5;
fn retry_backoff(failed_attempt: u32) -> Duration {
exponential_backoff(Duration::from_secs(1), failed_attempt)
}
async fn run(&self) -> Result<(), JobError>;
}
#[async_trait::async_trait]
pub trait JobQueue: Send + Sync + 'static {
async fn register<T: Job>(&self);
async fn register_with<T, F, Fut>(&self, run: F)
where
T: Job,
F: Fn(T) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), JobError>> + Send + 'static,
{
drop(run);
tracing::warn!(
target: "rustango::jobs",
job = T::NAME,
"this JobQueue does not override register_with; the handler is ignored \
and Job::run (with any process-wide fallback) runs instead"
);
self.register::<T>().await;
}
async fn dispatch<T: Job>(&self, payload: &T) -> Result<(), JobError>;
async fn start(&self);
async fn shutdown(&self);
async fn pending_count(&self) -> usize;
}
pub const DEFAULT_SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
pub(crate) struct Running {
workers: Vec<JoinHandle<()>>,
stop: watch::Sender<bool>,
}
#[derive(Clone)]
pub(crate) struct StopSignal(watch::Receiver<bool>);
impl StopSignal {
#[cfg_attr(not(feature = "jobs-postgres"), allow(dead_code))]
pub(crate) fn is_set(&self) -> bool {
*self.0.borrow() || self.0.has_changed().is_err()
}
pub(crate) async fn wait(&mut self) {
let _ = self.0.wait_for(|stop| *stop).await;
}
}
impl Running {
pub(crate) fn new() -> (Self, StopSignal) {
let (stop, rx) = watch::channel(false);
let run = Self {
workers: Vec::new(),
stop,
};
(run, StopSignal(rx))
}
pub(crate) fn push(&mut self, worker: JoinHandle<()>) {
self.workers.push(worker);
}
pub(crate) async fn stop(self, grace: Duration) -> Vec<usize> {
let _ = self.stop.send(true);
let deadline = tokio::time::Instant::now() + grace;
let mut aborted = Vec::new();
for (n, mut h) in self.workers.into_iter().enumerate() {
if tokio::time::timeout_at(deadline, &mut h).await.is_err() {
h.abort();
let _ = h.await;
aborted.push(n);
}
}
aborted
}
}
#[derive(Debug, Clone)]
struct JobEnvelope {
name: &'static str,
payload: serde_json::Value,
attempt: u32,
max_attempts: u32,
context: crate::task_context::TaskContext,
not_before: Option<tokio::time::Instant>,
}
type HandlerFn = Arc<
dyn Fn(serde_json::Value) -> Pin<Box<dyn Future<Output = Result<(), JobError>> + Send>>
+ Send
+ Sync,
>;
#[derive(Clone)]
struct Registered {
handler: HandlerFn,
#[cfg_attr(not(feature = "jobs-postgres"), allow(dead_code))]
name: &'static str,
retry_backoff: fn(u32) -> Duration,
}
impl Registered {
fn backoff(&self, failed_attempt: u32) -> Duration {
(self.retry_backoff)(failed_attempt).min(Duration::from_secs(365 * 86_400))
}
}
#[derive(Default)]
struct HandlerRegistry {
handlers: HashMap<&'static str, Registered>,
}
impl HandlerRegistry {
fn register<T: Job>(&mut self) {
self.register_with::<T, _, _>(|job: T| async move { job.run().await });
}
fn register_with<T, F, Fut>(&mut self, run: F)
where
T: Job,
F: Fn(T) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), JobError>> + Send + 'static,
{
let run = Arc::new(run);
let handler: HandlerFn = Arc::new(move |payload| {
let run = run.clone();
Box::pin(async move {
let job: T =
serde_json::from_value(payload).map_err(|e| JobError::Queue(e.to_string()))?;
crate::panic_guard::catch_unwind(run(job))
.await
.unwrap_or_else(|panic| {
let msg = crate::panic_guard::panic_message(&*panic);
Err(JobError::Retryable(format!("job panicked: {msg}")))
})
})
});
let entry = Registered {
handler,
name: T::NAME,
retry_backoff: T::retry_backoff,
};
self.handlers.insert(T::NAME, entry);
}
fn lookup(&self, name: &str) -> Option<Registered> {
self.handlers.get(name).cloned()
}
}
pub type DeadLetterFn =
Arc<dyn Fn(JobDeadLetter) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
#[derive(Debug, Clone)]
pub struct JobDeadLetter {
pub name: &'static str,
pub payload: serde_json::Value,
pub attempts: u32,
pub error: String,
}
pub(crate) async fn deliver_dead_letter(cb: DeadLetterFn, dl: JobDeadLetter) {
let name = dl.name;
if let Err(panic) = crate::panic_guard::catch_unwind(cb(dl)).await {
let msg = crate::panic_guard::panic_message(&*panic);
tracing::error!(job = name, panic = msg, "dead-letter callback panicked");
}
}
pub(crate) fn max_attempts<T: Job>() -> u32 {
T::MAX_ATTEMPTS.max(1)
}
#[must_use]
pub fn exponential_backoff(base: Duration, failed_attempt: u32) -> Duration {
base.saturating_mul(1u32 << failed_attempt.min(10))
}
pub struct InMemoryJobQueue {
tx: mpsc::UnboundedSender<JobEnvelope>,
rx: Arc<Mutex<mpsc::UnboundedReceiver<JobEnvelope>>>,
registry: Arc<Mutex<HandlerRegistry>>,
run: Mutex<Option<InMemoryRun>>,
worker_count: usize,
dead_letter: Arc<Mutex<Option<DeadLetterFn>>>,
pending: Arc<std::sync::atomic::AtomicUsize>,
retries: Arc<std::sync::Mutex<JoinSet<()>>>,
shutdown_grace: Duration,
}
type CurrentJob = Arc<std::sync::Mutex<Option<JobEnvelope>>>;
struct InMemoryRun {
run: Running,
current: Vec<CurrentJob>,
}
impl InMemoryJobQueue {
#[must_use]
pub fn with_workers(worker_count: usize) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
Self {
tx,
rx: Arc::new(Mutex::new(rx)),
registry: Arc::new(Mutex::new(HandlerRegistry::default())),
run: Mutex::new(None),
worker_count,
dead_letter: Arc::new(Mutex::new(None)),
pending: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
retries: Arc::default(),
shutdown_grace: DEFAULT_SHUTDOWN_GRACE,
}
}
#[must_use]
pub fn shutdown_grace(mut self, grace: Duration) -> Self {
self.shutdown_grace = grace;
self
}
#[must_use]
pub fn new() -> Self {
Self::with_workers(4)
}
pub async fn on_dead_letter<F, Fut>(&self, callback: F)
where
F: Fn(JobDeadLetter) -> Fut + Send + Sync + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let boxed: DeadLetterFn = Arc::new(move |dl| Box::pin(callback(dl)));
*self.dead_letter.lock().await = Some(boxed);
}
}
impl Default for InMemoryJobQueue {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "config")]
#[must_use]
pub fn inmemory_from_settings(s: &crate::config::JobsSettings) -> Arc<InMemoryJobQueue> {
let workers = s.concurrency.map_or(4, |c| c as usize);
if let Some(backend) = s.backend.as_deref() {
if backend != "memory" {
tracing::warn!(
target: "rustango::jobs",
backend = %backend,
"jobs.backend = `{backend}` but inmemory_from_settings only builds InMemoryJobQueue. \
Wire the desired backend directly via Arc::new(...). See the docstring."
);
}
}
Arc::new(InMemoryJobQueue::with_workers(workers))
}
#[async_trait::async_trait]
impl JobQueue for InMemoryJobQueue {
async fn register<T: Job>(&self) {
self.registry.lock().await.register::<T>();
}
async fn register_with<T, F, Fut>(&self, run: F)
where
T: Job,
F: Fn(T) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<(), JobError>> + Send + 'static,
{
self.registry.lock().await.register_with::<T, F, Fut>(run);
}
async fn dispatch<T: Job>(&self, payload: &T) -> Result<(), JobError> {
let value = serde_json::to_value(payload).map_err(|e| JobError::Queue(e.to_string()))?;
let envelope = JobEnvelope {
name: T::NAME,
payload: value,
attempt: 0,
max_attempts: max_attempts::<T>(),
context: crate::task_context::TaskContext::capture(),
not_before: None,
};
self.pending
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
self.tx.send(envelope).map_err(|e| {
self.pending
.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
JobError::Queue(e.to_string())
})?;
Ok(())
}
async fn start(&self) {
let mut slot = self.run.lock().await;
if slot.is_some() {
return; }
let (mut run, stop) = Running::new();
let mut current = Vec::with_capacity(self.worker_count);
for _ in 0..self.worker_count {
let job = CurrentJob::default();
current.push(job.clone());
let worker = InMemoryWorker {
rx: self.rx.clone(),
registry: self.registry.clone(),
dead_letter: self.dead_letter.clone(),
pending: self.pending.clone(),
tx: self.tx.clone(),
retries: self.retries.clone(),
stop: stop.clone(),
current: job,
};
run.push(tokio::spawn(worker_loop(worker)));
}
*slot = Some(InMemoryRun { run, current });
}
async fn shutdown(&self) {
let mut slot = self.run.lock().await;
let Some(InMemoryRun { run, current }) = slot.take() else {
return;
};
for n in run.stop(self.shutdown_grace).await {
if let Some(mut job) = current[n].lock().unwrap().take() {
tracing::warn!(job = job.name, "job aborted at shutdown; re-queued");
job.attempt += 1;
let _ = self.tx.send(job);
}
}
let mut timers = std::mem::take(&mut *self.retries.lock().unwrap());
while timers.join_next().await.is_some() {}
}
async fn pending_count(&self) -> usize {
self.pending.load(std::sync::atomic::Ordering::SeqCst)
}
}
struct InMemoryWorker {
rx: Arc<Mutex<mpsc::UnboundedReceiver<JobEnvelope>>>,
registry: Arc<Mutex<HandlerRegistry>>,
dead_letter: Arc<Mutex<Option<DeadLetterFn>>>,
pending: Arc<std::sync::atomic::AtomicUsize>,
tx: mpsc::UnboundedSender<JobEnvelope>,
retries: Arc<std::sync::Mutex<JoinSet<()>>>,
stop: StopSignal,
current: CurrentJob,
}
async fn worker_loop(w: InMemoryWorker) {
let InMemoryWorker {
rx,
registry,
dead_letter,
pending,
tx,
retries,
mut stop,
current,
} = w;
loop {
let envelope = tokio::select! {
biased;
() = stop.wait() => return,
e = async { rx.lock().await.recv().await } => match e {
Some(e) => e,
None => return, },
};
*current.lock().unwrap() = Some(envelope.clone());
if let Some(until) = envelope
.not_before
.filter(|t| *t > tokio::time::Instant::now())
{
current.lock().unwrap().take();
park(&retries, &tx, stop.clone(), envelope, until);
continue;
}
if envelope.attempt >= envelope.max_attempts {
current.lock().unwrap().take();
pending.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
let error = "attempts spent by runs aborted at shutdown".to_owned();
let attempts = envelope.attempt;
dead_letter_job(&dead_letter, envelope, attempts, error).await;
continue;
}
let handler = {
let reg = registry.lock().await;
reg.lookup(envelope.name)
};
let Some(entry) = handler else {
current.lock().unwrap().take();
tracing::error!(job = envelope.name, "no handler registered");
pending.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
continue;
};
let payload = envelope.payload.clone();
let result = envelope
.context
.clone()
.install((entry.handler)(payload))
.await;
current.lock().unwrap().take();
match result {
Ok(()) => {
pending.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
}
Err(JobError::Retryable(msg)) => {
let next_attempt = envelope.attempt + 1;
if next_attempt >= envelope.max_attempts {
pending.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
dead_letter_job(&dead_letter, envelope, next_attempt, msg).await;
} else {
let backoff = entry.backoff(envelope.attempt);
let mut retry = envelope;
retry.attempt = next_attempt;
let until = tokio::time::Instant::now() + backoff;
park(&retries, &tx, stop.clone(), retry, until);
}
}
Err(e @ (JobError::Fatal(_) | JobError::Queue(_))) => {
pending.fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
let attempts = envelope.attempt + 1;
dead_letter_job(&dead_letter, envelope, attempts, e.to_string()).await;
}
}
}
}
fn park(
retries: &std::sync::Mutex<JoinSet<()>>,
tx: &mpsc::UnboundedSender<JobEnvelope>,
mut stop: StopSignal,
mut job: JobEnvelope,
until: tokio::time::Instant,
) {
job.not_before = Some(until);
let tx = tx.clone();
let mut timers = retries.lock().unwrap();
while timers.try_join_next().is_some() {}
timers.spawn(async move {
tokio::select! {
() = tokio::time::sleep_until(until) => {}
() = stop.wait() => {}
}
let _ = tx.send(job);
});
}
async fn dead_letter_job(
dead_letter: &Mutex<Option<DeadLetterFn>>,
job: JobEnvelope,
attempts: u32,
error: String,
) {
let cb = dead_letter.lock().await.clone();
if let Some(cb) = cb {
let dl = JobDeadLetter {
name: job.name,
payload: job.payload,
attempts,
error,
};
job.context.install(deliver_dead_letter(cb, dl)).await;
} else {
tracing::error!(job = job.name, attempts, error = %error, "job dead-lettered");
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
fn retry_backoff_starts_at_one_second_and_doubles() {
let ms = |n| Increment::retry_backoff(n).as_millis();
assert_eq!(ms(0), 1_000, "the first retry waits 1s, not 2s");
assert_eq!(ms(1), 2_000);
assert_eq!(ms(2), 4_000);
assert_eq!(ms(3), 8_000);
}
#[test]
fn retry_backoff_caps_rather_than_overflowing() {
let ms = |n| Increment::retry_backoff(n).as_millis();
assert_eq!(ms(10), 1_024_000);
assert_eq!(
ms(u32::MAX),
1_024_000,
"a runaway attempt count must clamp, not shift past 63 and panic"
);
}
#[test]
fn a_runaway_backoff_hook_is_capped_at_a_year() {
#[derive(Serialize, Deserialize)]
struct Forever;
#[async_trait::async_trait]
impl Job for Forever {
const NAME: &'static str = "test:forever";
fn retry_backoff(_: u32) -> Duration {
Duration::MAX
}
async fn run(&self) -> Result<(), JobError> {
Ok(())
}
}
let mut reg = HandlerRegistry::default();
reg.register::<Forever>();
let wait = reg.lookup(Forever::NAME).unwrap().backoff(0);
assert_eq!(wait, Duration::from_secs(365 * 86_400));
let _ = tokio::time::Instant::now() + wait;
let _ = chrono::Utc::now() + chrono::Duration::from_std(wait).unwrap();
}
#[tokio::test]
async fn a_jobs_backoff_hook_sets_the_retry_delay() {
static RUNS: AtomicUsize = AtomicUsize::new(0);
#[derive(Serialize, Deserialize)]
struct Quick;
#[async_trait::async_trait]
impl Job for Quick {
const NAME: &'static str = "test:quick_backoff";
const MAX_ATTEMPTS: u32 = 3;
fn retry_backoff(_: u32) -> Duration {
Duration::from_millis(10)
}
async fn run(&self) -> Result<(), JobError> {
RUNS.fetch_add(1, Ordering::SeqCst);
Err(JobError::Retryable("again".into()))
}
}
let q = InMemoryJobQueue::with_workers(1);
q.register::<Quick>().await;
q.start().await;
q.dispatch(&Quick).await.unwrap();
tokio::time::sleep(Duration::from_millis(500)).await;
q.shutdown().await;
assert_eq!(RUNS.load(Ordering::SeqCst), 3);
}
#[derive(Serialize, Deserialize, Debug)]
struct Increment;
static COUNTER: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for Increment {
const NAME: &'static str = "test:increment";
async fn run(&self) -> Result<(), JobError> {
COUNTER.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[derive(Serialize, Deserialize, Debug)]
struct AlwaysFail {
fatal: bool,
}
#[async_trait::async_trait]
impl Job for AlwaysFail {
const NAME: &'static str = "test:always_fail";
const MAX_ATTEMPTS: u32 = 2;
async fn run(&self) -> Result<(), JobError> {
if self.fatal {
Err(JobError::Fatal("dead now".into()))
} else {
Err(JobError::Retryable("transient".into()))
}
}
}
#[derive(Serialize, Deserialize, Debug)]
struct EventuallyOk {
fail_n: u32,
success_marker_id: u64,
}
static SUCCESSES: std::sync::Mutex<Vec<u64>> = std::sync::Mutex::new(Vec::new());
static ATTEMPTS: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for EventuallyOk {
const NAME: &'static str = "test:eventually_ok";
const MAX_ATTEMPTS: u32 = 5;
async fn run(&self) -> Result<(), JobError> {
let n = ATTEMPTS.fetch_add(1, Ordering::SeqCst);
if (n as u32) < self.fail_n {
Err(JobError::Retryable(format!("attempt {n}")))
} else {
SUCCESSES.lock().unwrap().push(self.success_marker_id);
Ok(())
}
}
}
#[tokio::test]
async fn a_job_sees_the_context_of_whoever_enqueued_it() {
use crate::audit::{current_source, with_source, AuditSource};
static SEEN: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
#[derive(serde::Serialize, serde::Deserialize)]
struct RecordSource;
#[async_trait::async_trait]
impl Job for RecordSource {
const NAME: &'static str = "record_source";
async fn run(&self) -> Result<(), JobError> {
*SEEN.lock().unwrap() = Some(current_source().as_token());
Ok(())
}
}
let q = InMemoryJobQueue::with_workers(1);
q.register::<RecordSource>().await;
q.start().await;
with_source(AuditSource::User { id: "99".into() }, async {
q.dispatch(&RecordSource).await.unwrap();
})
.await;
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while SEEN.lock().unwrap().is_none() && std::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(10)).await;
}
q.shutdown().await;
assert_eq!(
SEEN.lock().unwrap().clone(),
Some("user:99".to_owned()),
"the job should see who enqueued it"
);
}
#[tokio::test]
async fn a_job_enqueued_outside_any_scope_is_system() {
use crate::audit::current_source;
static SEEN: std::sync::Mutex<Option<String>> = std::sync::Mutex::new(None);
#[derive(serde::Serialize, serde::Deserialize)]
struct RecordPlain;
#[async_trait::async_trait]
impl Job for RecordPlain {
const NAME: &'static str = "record_plain";
async fn run(&self) -> Result<(), JobError> {
*SEEN.lock().unwrap() = Some(current_source().as_token());
Ok(())
}
}
let q = InMemoryJobQueue::with_workers(1);
q.register::<RecordPlain>().await;
q.start().await;
q.dispatch(&RecordPlain).await.unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while SEEN.lock().unwrap().is_none() && std::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(10)).await;
}
q.shutdown().await;
assert_eq!(SEEN.lock().unwrap().clone(), Some("system".to_owned()));
}
#[tokio::test]
async fn dispatch_runs_handler() {
COUNTER.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(2);
q.register::<Increment>().await;
q.start().await;
q.dispatch(&Increment).await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
q.shutdown().await;
}
#[tokio::test]
async fn fatal_error_goes_to_dead_letter() {
let q = InMemoryJobQueue::with_workers(1);
q.register::<AlwaysFail>().await;
let captured: Arc<Mutex<Vec<JobDeadLetter>>> = Arc::new(Mutex::new(Vec::new()));
let cap = captured.clone();
q.on_dead_letter(move |dl| {
let cap = cap.clone();
async move {
cap.lock().await.push(dl);
}
})
.await;
q.start().await;
q.dispatch(&AlwaysFail { fatal: true }).await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(captured.lock().await.len(), 1);
assert!(captured.lock().await[0].error.contains("dead now"));
q.shutdown().await;
}
#[tokio::test]
async fn zero_max_attempts_runs_once() {
static RUNS: AtomicUsize = AtomicUsize::new(0);
#[derive(Serialize, Deserialize)]
struct NoAttempts;
#[async_trait::async_trait]
impl Job for NoAttempts {
const NAME: &'static str = "test:no_attempts";
const MAX_ATTEMPTS: u32 = 0;
async fn run(&self) -> Result<(), JobError> {
RUNS.fetch_add(1, Ordering::SeqCst);
Err(JobError::Retryable("again".into()))
}
}
let q = InMemoryJobQueue::with_workers(1);
q.register::<NoAttempts>().await;
let captured: Arc<std::sync::Mutex<Vec<JobDeadLetter>>> = Arc::default();
let cap = captured.clone();
q.on_dead_letter(move |dl| {
cap.lock().unwrap().push(dl);
async {}
})
.await;
q.start().await;
q.dispatch(&NoAttempts).await.unwrap();
wait_until("the dead letter", || !captured.lock().unwrap().is_empty()).await;
q.shutdown().await;
assert_eq!(RUNS.load(Ordering::SeqCst), 1);
let dl = captured.lock().unwrap();
assert_eq!((dl[0].attempts, dl[0].error.as_str()), (1, "again"));
}
#[tokio::test]
async fn dead_letter_callback_sees_the_enqueuer() {
use crate::audit::{current_source, with_source, AuditSource};
let q = InMemoryJobQueue::with_workers(1);
q.register::<AlwaysFail>().await;
let seen: Arc<Mutex<Option<String>>> = Arc::default();
let s = seen.clone();
q.on_dead_letter(move |_dl| {
let s = s.clone();
async move {
*s.lock().await = Some(current_source().as_token());
}
})
.await;
q.start().await;
with_source(AuditSource::User { id: "5".into() }, async {
q.dispatch(&AlwaysFail { fatal: true }).await.unwrap();
})
.await;
let deadline = std::time::Instant::now() + Duration::from_secs(5);
while seen.lock().await.is_none() && std::time::Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(10)).await;
}
q.shutdown().await;
assert_eq!(seen.lock().await.as_deref(), Some("user:5"));
}
#[tokio::test]
async fn retryable_succeeds_eventually() {
ATTEMPTS.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<EventuallyOk>().await;
q.start().await;
let marker = 12345;
q.dispatch(&EventuallyOk {
fail_n: 2,
success_marker_id: marker,
})
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(7000)).await;
let succ = SUCCESSES.lock().unwrap();
assert!(succ.contains(&marker), "expected marker, got {succ:?}");
drop(succ);
q.shutdown().await;
}
#[derive(Serialize, Deserialize, Debug)]
struct Panics;
#[async_trait::async_trait]
impl Job for Panics {
const NAME: &'static str = "test:panics";
const MAX_ATTEMPTS: u32 = 1;
async fn run(&self) -> Result<(), JobError> {
panic!("boom");
}
}
#[derive(Serialize, Deserialize, Debug)]
struct AfterPanic;
static AFTER_PANIC: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for AfterPanic {
const NAME: &'static str = "test:after_panic";
async fn run(&self) -> Result<(), JobError> {
AFTER_PANIC.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[tokio::test]
async fn a_panicking_job_keeps_the_worker() {
let q = InMemoryJobQueue::with_workers(1);
q.register::<Panics>().await;
q.register::<AfterPanic>().await;
let dead: Arc<Mutex<Vec<JobDeadLetter>>> = Arc::default();
let d = dead.clone();
q.on_dead_letter(move |dl| {
let d = d.clone();
async move { d.lock().await.push(dl) }
})
.await;
q.start().await;
q.dispatch(&Panics).await.unwrap();
q.dispatch(&AfterPanic).await.unwrap();
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(AFTER_PANIC.load(Ordering::SeqCst), 1, "worker still runs");
let dead = dead.lock().await;
assert!(dead[0].error.contains("job panicked: boom"), "{dead:?}");
assert_eq!(q.pending_count().await, 0);
q.shutdown().await;
}
#[tokio::test]
async fn unknown_job_is_logged_not_panic() {
#[derive(Serialize, Deserialize)]
struct UnregisteredJob;
#[async_trait::async_trait]
impl Job for UnregisteredJob {
const NAME: &'static str = "test:unregistered";
async fn run(&self) -> Result<(), JobError> {
Ok(())
}
}
let q = InMemoryJobQueue::with_workers(1);
q.start().await;
q.dispatch(&UnregisteredJob).await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
q.shutdown().await;
}
#[tokio::test]
async fn pending_count_tracks_in_flight() {
let q = InMemoryJobQueue::with_workers(0); q.register::<Increment>().await;
for _ in 0..3 {
q.dispatch(&Increment).await.unwrap();
}
assert_eq!(q.pending_count().await, 3);
}
#[derive(Serialize, Deserialize, Debug)]
struct SlowOnce {
ms: u64,
}
static DRAIN_RUNS: AtomicUsize = AtomicUsize::new(0);
static DRAINED: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for SlowOnce {
const NAME: &'static str = "test:slow_once";
async fn run(&self) -> Result<(), JobError> {
if DRAIN_RUNS.fetch_add(1, Ordering::SeqCst) == 0 {
tokio::time::sleep(Duration::from_millis(self.ms)).await;
}
DRAINED.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[derive(Serialize, Deserialize, Debug)]
struct RetryOnce;
static RETRY_RUNS: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for RetryOnce {
const NAME: &'static str = "test:retry_once";
async fn run(&self) -> Result<(), JobError> {
if RETRY_RUNS.fetch_add(1, Ordering::SeqCst) == 0 {
return Err(JobError::Retryable("once".into()));
}
Ok(())
}
}
async fn wait_until(what: &str, cond: impl Fn() -> bool) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while !cond() {
assert!(tokio::time::Instant::now() < deadline, "timed out: {what}");
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
static LIFECYCLE: Mutex<()> = Mutex::const_new(());
#[tokio::test]
async fn shutdown_drains_a_running_job() {
let _g = LIFECYCLE.lock().await;
DRAIN_RUNS.store(0, Ordering::SeqCst);
DRAINED.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<SlowOnce>().await;
q.start().await;
q.dispatch(&SlowOnce { ms: 300 }).await.unwrap();
wait_until("the run to start", || {
DRAIN_RUNS.load(Ordering::SeqCst) == 1
})
.await;
q.shutdown().await;
assert_eq!(DRAINED.load(Ordering::SeqCst), 1, "the job finished");
assert_eq!(q.pending_count().await, 0);
}
#[tokio::test]
async fn a_job_past_the_grace_is_requeued() {
let _g = LIFECYCLE.lock().await;
DRAIN_RUNS.store(0, Ordering::SeqCst);
DRAINED.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1).shutdown_grace(Duration::from_millis(50));
q.register::<SlowOnce>().await;
q.start().await;
q.dispatch(&SlowOnce { ms: 30_000 }).await.unwrap();
wait_until("the run to start", || {
DRAIN_RUNS.load(Ordering::SeqCst) == 1
})
.await;
q.shutdown().await;
assert_eq!(q.pending_count().await, 1, "aborted, not lost");
q.start().await;
wait_until("the rerun", || DRAINED.load(Ordering::SeqCst) == 1).await;
wait_until("the count to settle", || {
q.pending.load(Ordering::SeqCst) == 0
})
.await;
q.shutdown().await;
}
#[tokio::test]
async fn a_parked_retry_survives_shutdown() {
let _g = LIFECYCLE.lock().await;
RETRY_RUNS.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<RetryOnce>().await;
q.start().await;
q.dispatch(&RetryOnce).await.unwrap();
wait_until("the failed run", || RETRY_RUNS.load(Ordering::SeqCst) == 1).await;
q.shutdown().await;
assert_eq!(q.pending_count().await, 1, "the retry is still queued");
q.start().await;
wait_until("the retry", || RETRY_RUNS.load(Ordering::SeqCst) == 2).await;
wait_until("the count to reach 0", || {
q.pending.load(Ordering::SeqCst) == 0
})
.await;
q.shutdown().await;
}
#[tokio::test]
async fn start_after_shutdown_runs_jobs() {
let _g = LIFECYCLE.lock().await;
DRAIN_RUNS.store(1, Ordering::SeqCst); DRAINED.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<SlowOnce>().await;
q.start().await;
q.shutdown().await;
q.start().await;
q.dispatch(&SlowOnce { ms: 0 }).await.unwrap();
wait_until("the job after a restart", || {
DRAINED.load(Ordering::SeqCst) == 1
})
.await;
q.shutdown().await;
}
#[derive(Serialize, Deserialize, Debug)]
struct Marker;
static MARKED: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for Marker {
const NAME: &'static str = "test:marker";
async fn run(&self) -> Result<(), JobError> {
MARKED.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[derive(Serialize, Deserialize, Debug)]
struct HangOnce;
static HANG_RUNS: AtomicUsize = AtomicUsize::new(0);
#[async_trait::async_trait]
impl Job for HangOnce {
const NAME: &'static str = "test:hang_once";
const MAX_ATTEMPTS: u32 = 1;
async fn run(&self) -> Result<(), JobError> {
HANG_RUNS.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_secs(30)).await;
Ok(())
}
}
#[tokio::test]
async fn dispatch_during_drain_stays_queued() {
let _g = LIFECYCLE.lock().await;
DRAIN_RUNS.store(0, Ordering::SeqCst);
MARKED.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<SlowOnce>().await;
q.register::<Marker>().await;
q.start().await;
q.dispatch(&SlowOnce { ms: 300 }).await.unwrap();
wait_until("the run to start", || {
DRAIN_RUNS.load(Ordering::SeqCst) == 1
})
.await;
tokio::join!(q.shutdown(), async {
tokio::time::sleep(Duration::from_millis(50)).await;
q.dispatch(&Marker).await.unwrap();
});
assert_eq!(MARKED.load(Ordering::SeqCst), 0, "not run while stopping");
assert_eq!(q.pending_count().await, 1);
q.start().await;
wait_until("the queued job", || MARKED.load(Ordering::SeqCst) == 1).await;
q.shutdown().await;
}
#[tokio::test]
async fn a_parked_retry_keeps_its_backoff_across_restart() {
let _g = LIFECYCLE.lock().await;
RETRY_RUNS.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1);
q.register::<RetryOnce>().await;
q.start().await;
q.dispatch(&RetryOnce).await.unwrap();
wait_until("the failed run", || RETRY_RUNS.load(Ordering::SeqCst) == 1).await;
q.shutdown().await;
q.start().await;
tokio::time::sleep(Duration::from_millis(400)).await;
assert_eq!(
RETRY_RUNS.load(Ordering::SeqCst),
1,
"the 1s backoff still holds"
);
wait_until("the retry", || RETRY_RUNS.load(Ordering::SeqCst) == 2).await;
q.shutdown().await;
}
#[tokio::test]
async fn an_aborted_run_spends_its_attempt() {
let _g = LIFECYCLE.lock().await;
HANG_RUNS.store(0, Ordering::SeqCst);
let q = InMemoryJobQueue::with_workers(1).shutdown_grace(Duration::from_millis(50));
q.register::<HangOnce>().await;
let dead: Arc<Mutex<Vec<JobDeadLetter>>> = Arc::default();
let d = dead.clone();
q.on_dead_letter(move |dl| {
let d = d.clone();
async move { d.lock().await.push(dl) }
})
.await;
q.start().await;
q.dispatch(&HangOnce).await.unwrap();
wait_until("the run to start", || HANG_RUNS.load(Ordering::SeqCst) == 1).await;
q.shutdown().await;
q.start().await;
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while dead.lock().await.is_empty() {
assert!(tokio::time::Instant::now() < deadline, "no dead letter");
tokio::time::sleep(Duration::from_millis(10)).await;
}
assert_eq!(HANG_RUNS.load(Ordering::SeqCst), 1, "not run again");
assert_eq!(q.pending_count().await, 0);
q.shutdown().await;
}
#[cfg(feature = "config")]
#[test]
fn inmemory_from_settings_uses_configured_concurrency() {
let mut s = crate::config::JobsSettings::default();
s.concurrency = Some(8);
let q = inmemory_from_settings(&s);
assert_eq!(q.worker_count, 8);
}
#[cfg(feature = "config")]
#[test]
fn inmemory_from_settings_defaults_to_four_workers() {
let s = crate::config::JobsSettings::default();
let q = inmemory_from_settings(&s);
assert_eq!(q.worker_count, 4);
}
}