mod registry;
use crate::backoff::{Backoff, retry_delay};
use crate::context::JournalEntry;
use crate::error::Error;
use crate::outcome::Outcome;
use crate::store::{JournalAppend, JournalOutcome, Settlement, Store};
use crate::task::Handler;
use chrono::{DateTime, Utc};
use registry::{Registry, RunInput, RunReport};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use ulid::Ulid;
const DEFAULT_BACKSTOP: u32 = 20;
const DEFAULT_JITTER_FRACTION: f64 = 0.5;
const DEFAULT_PRIORITY_RATIO: u32 = 4;
const MIN_LEASE: Duration = Duration::from_secs(1);
const MAX_LEASE: Duration = Duration::from_secs(365 * 24 * 60 * 60);
const MIN_POLL_MAX: Duration = Duration::from_millis(1);
fn clamp_lease(lease: Duration) -> Duration {
lease.clamp(MIN_LEASE, MAX_LEASE)
}
const FLOOR_ALL: i16 = 0;
const FLOOR_NORMAL: i16 = 1;
const FLOOR_LOW: i16 = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PanicPolicy {
Retry,
Dead,
}
#[derive(Debug, Clone)]
struct WorkerConfig {
concurrency: usize,
poll_max: Duration,
lease: Duration,
shutdown_timeout: Duration,
backoff: Backoff,
jitter_fraction: f64,
backstop: Option<u32>,
priority_ratio: Option<u32>,
panic_policy: PanicPolicy,
}
impl Default for WorkerConfig {
fn default() -> WorkerConfig {
WorkerConfig {
concurrency: default_concurrency(),
poll_max: Duration::from_secs(30),
lease: Duration::from_secs(15 * 60),
shutdown_timeout: Duration::from_secs(30),
backoff: Backoff::default(),
jitter_fraction: DEFAULT_JITTER_FRACTION,
backstop: Some(DEFAULT_BACKSTOP),
priority_ratio: Some(DEFAULT_PRIORITY_RATIO),
panic_policy: PanicPolicy::Retry,
}
}
}
fn default_concurrency() -> usize {
let cores = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
(cores / 2).clamp(1, 8)
}
pub struct WorkerBuilder<S> {
state: S,
store: Arc<dyn Store>,
registry: Registry<S>,
config: WorkerConfig,
}
impl<S> WorkerBuilder<S>
where
S: Send + Sync + 'static,
{
#[must_use]
pub fn register<T>(mut self) -> Self
where
T: Handler<S>,
{
self.registry.register::<T>(None);
self
}
#[must_use]
pub fn register_capped<T>(mut self, max: usize) -> Self
where
T: Handler<S>,
{
self.registry.register::<T>(Some(max.max(1)));
self
}
#[must_use]
pub fn concurrency(mut self, n: usize) -> Self {
self.config.concurrency = n.max(1);
self
}
#[must_use]
pub fn poll_max(mut self, d: Duration) -> Self {
self.config.poll_max = d.max(MIN_POLL_MAX);
self
}
#[must_use]
pub fn lease(mut self, d: Duration) -> Self {
self.config.lease = clamp_lease(d);
self
}
#[must_use]
pub fn shutdown_timeout(mut self, d: Duration) -> Self {
self.config.shutdown_timeout = d;
self
}
#[must_use]
pub fn backoff(mut self, backoff: Backoff) -> Self {
self.config.backoff = backoff;
self
}
#[must_use]
pub fn backstop(mut self, backstop: Option<u32>) -> Self {
self.config.backstop = backstop;
self
}
#[must_use]
pub fn priority_ratio(mut self, ratio: Option<u32>) -> Self {
self.config.priority_ratio = ratio;
self
}
#[must_use]
pub fn jitter_fraction(mut self, fraction: f64) -> Self {
self.config.jitter_fraction = fraction.clamp(0.0, 1.0);
self
}
#[must_use]
pub fn panic_policy(mut self, policy: PanicPolicy) -> Self {
self.config.panic_policy = policy;
self
}
pub fn build(self) -> Worker<S> {
Worker {
state: Arc::new(self.state),
store: self.store,
registry: Arc::new(self.registry),
config: self.config,
identity: worker_identity(),
}
}
}
pub struct Worker<S> {
state: Arc<S>,
store: Arc<dyn Store>,
registry: Arc<Registry<S>>,
config: WorkerConfig,
identity: String,
}
impl<S> Worker<S>
where
S: Send + Sync + 'static,
{
pub fn builder(state: S, store: Arc<dyn Store>) -> WorkerBuilder<S> {
WorkerBuilder {
state,
store,
registry: Registry::new(),
config: WorkerConfig::default(),
}
}
pub async fn run(self, shutdown: CancellationToken) {
let kinds = self.registry.kinds();
if kinds.is_empty() {
return;
}
let mut notifier = match self.store.notifier().await {
Ok(notifier) => notifier,
Err(error) => {
tracing::warn!(%error, "could not set up notifications; polling only");
Box::new(crate::store::NeverNotifier)
}
};
let mut running: JoinSet<FinishedRun> = JoinSet::new();
let mut inflight: HashMap<tokio::task::Id, InflightJob> = HashMap::new();
let mut by_kind: HashMap<String, usize> = HashMap::new();
let mut claim_counter: u64 = 0;
'outer: loop {
self.recover_stale().await;
while running.len() < self.config.concurrency {
if shutdown.is_cancelled() {
break 'outer;
}
let claimable = self.claimable_kinds(&kinds, &by_kind);
if claimable.is_empty() {
break;
}
claim_counter = claim_counter.wrapping_add(1);
let floor = floor_for(claim_counter, self.config.priority_ratio);
match self.claim_with_fallback(&claimable, floor).await {
Ok(Some(job)) => {
let wait = (Utc::now() - job.created_at)
.to_std()
.unwrap_or(Duration::ZERO);
crate::observability::claimed(&job.kind, wait);
self.apply_task_lease(&job).await;
let history = self.history_for(job.id).await;
*by_kind.entry(job.kind.clone()).or_insert(0) += 1;
let tracked = InflightJob {
id: job.id,
kind: job.kind.clone(),
run_no: job.run_count,
};
let task_id = self.spawn(&mut running, job, history, &shutdown);
inflight.insert(task_id, tracked);
}
Ok(None) => break,
Err(error) => {
tracing::warn!(%error, "claim failed; will retry");
break;
}
}
}
let wait = self.wait_duration(&kinds).await;
if running.is_empty() {
tokio::select! {
_ = shutdown.cancelled() => break,
_ = notifier.recv() => {}
_ = tokio::time::sleep(wait) => {}
}
} else {
tokio::select! {
_ = shutdown.cancelled() => break,
Some(joined) = running.join_next_with_id() => {
self.reap(joined, &mut inflight, &mut by_kind).await;
}
_ = notifier.recv() => {}
_ = tokio::time::sleep(wait) => {}
}
}
}
self.drain(running, inflight, by_kind).await;
}
fn claimable_kinds(&self, kinds: &[String], by_kind: &HashMap<String, usize>) -> Vec<String> {
kinds
.iter()
.filter(|kind| match self.registry.cap(kind) {
Some(cap) => by_kind.get(*kind).copied().unwrap_or(0) < cap,
None => true,
})
.cloned()
.collect()
}
async fn claim_with_fallback(
&self,
kinds: &[String],
floor: i16,
) -> Result<Option<crate::store::JobRecord>, Error> {
let claimed = self
.store
.claim_next(kinds, floor, self.config.lease, &self.identity)
.await?;
if claimed.is_some() || floor == FLOOR_ALL {
return Ok(claimed);
}
self.store
.claim_next(kinds, FLOOR_ALL, self.config.lease, &self.identity)
.await
}
async fn wait_duration(&self, kinds: &[String]) -> Duration {
let until_next = match self.store.next_visible_at(kinds).await {
Ok(Some(at)) => (at - Utc::now())
.to_std()
.unwrap_or(Duration::ZERO)
.max(Duration::from_millis(5)),
Ok(None) => self.config.poll_max,
Err(error) => {
tracing::warn!(%error, "could not compute next eligibility; using poll_max");
self.config.poll_max
}
};
until_next.min(self.config.poll_max)
}
async fn drain(
&self,
mut running: JoinSet<FinishedRun>,
mut inflight: HashMap<tokio::task::Id, InflightJob>,
mut by_kind: HashMap<String, usize>,
) {
crate::observability::shutdown_drain(running.len());
let deadline = tokio::time::sleep(self.config.shutdown_timeout);
tokio::pin!(deadline);
while !running.is_empty() {
tokio::select! {
joined = running.join_next_with_id() => match joined {
Some(result) => self.reap(result, &mut inflight, &mut by_kind).await,
None => break,
},
_ = &mut deadline => {
break;
}
}
}
self.force_finish(running, inflight).await;
}
async fn force_finish(
&self,
mut running: JoinSet<FinishedRun>,
mut inflight: HashMap<tokio::task::Id, InflightJob>,
) {
running.abort_all();
while let Some(joined) = running.join_next_with_id().await {
if let Ok((task_id, finished)) = joined {
inflight.remove(&task_id);
if let Err(error) = self.settle(finished).await {
tracing::warn!(%error, "settle failed during shutdown drain");
}
}
}
for (_, job) in inflight.drain() {
self.release(&job).await;
}
}
async fn history_for(&self, id: Ulid) -> Vec<JournalEntry> {
match self.store.journal(id).await {
Ok(records) => records.into_iter().map(JournalEntry::from_record).collect(),
Err(error) => {
tracing::warn!(%error, "could not load job history; proceeding with none");
Vec::new()
}
}
}
async fn recover_stale(&self) {
let stale = match self.store.find_stale().await {
Ok(stale) => stale,
Err(error) => {
tracing::warn!(%error, "could not scan for stale claims");
return;
}
};
let now = Utc::now();
for job in stale {
let next_failures = job.failure_count.saturating_add(1);
let attempt = next_failures.max(0) as u32;
let delay = retry_delay(
&self.config.backoff,
self.config.jitter_fraction,
attempt,
job.id,
);
let note = format!(
"lease expired; worker {} presumed dead",
job.claimed_by.as_deref().unwrap_or("unknown"),
);
let journal = JournalAppend {
kind: job.kind.clone(),
run_no: job.run_count,
recorded_at: now,
outcome: JournalOutcome::StaleRecovered,
note: Some(note),
attachment: None,
};
let kind = job.kind.clone();
match self
.store
.recover(
job.id,
add_duration(now, delay),
next_failures,
job.run_count,
journal,
)
.await
{
Ok(true) => crate::observability::recovered(&kind),
Ok(false) => {}
Err(error) => tracing::warn!(%error, "stale-claim recovery failed"),
}
}
}
async fn apply_task_lease(&self, job: &crate::store::JobRecord) {
let Some(lease) = self.registry.lease_for(&job.kind, &job.payload) else {
return;
};
let lease = clamp_lease(lease);
if lease == self.config.lease {
return;
}
if let Err(error) = self
.store
.extend_lease(job.id, &self.identity, job.run_count, lease)
.await
{
tracing::warn!(%error, "could not apply task lease override");
}
}
fn spawn(
&self,
running: &mut JoinSet<FinishedRun>,
job: crate::store::JobRecord,
history: Vec<JournalEntry>,
shutdown: &CancellationToken,
) -> tokio::task::Id {
let id = job.id;
let kind = job.kind.clone();
let run_no = job.run_count;
let failure_count = job.failure_count;
let input = RunInput {
payload: job.payload,
carry: job.carry,
run_count: job.run_count,
history,
state: Arc::clone(&self.state),
cancel: shutdown.clone(),
panic_policy: self.config.panic_policy,
};
let report = self.registry.dispatch(&kind, input);
running
.spawn(async move {
let report = match report {
Ok(future) => future.await,
Err(error) => Err(error),
};
FinishedRun {
id,
kind,
run_no,
failure_count,
report,
}
})
.id()
}
async fn reap(
&self,
joined: Result<(tokio::task::Id, FinishedRun), tokio::task::JoinError>,
inflight: &mut HashMap<tokio::task::Id, InflightJob>,
by_kind: &mut HashMap<String, usize>,
) {
match joined {
Ok((task_id, finished)) => {
if let Some(job) = inflight.remove(&task_id) {
release_slot(by_kind, &job.kind);
}
if let Err(error) = self.settle(finished).await {
tracing::warn!(%error, "settle failed");
}
}
Err(join_error) => {
if let Some(job) = inflight.remove(&join_error.id()) {
release_slot(by_kind, &job.kind);
}
tracing::error!(%join_error, "handler task did not complete cleanly");
}
}
}
async fn release(&self, job: &InflightJob) {
let now = Utc::now();
let journal = JournalAppend {
kind: job.kind.clone(),
run_no: job.run_no,
recorded_at: now,
outcome: JournalOutcome::Released,
note: Some("released by graceful shutdown".to_owned()),
attachment: None,
};
let settlement = Settlement::Release { visible_at: now };
if let Err(error) = self
.store
.settle(job.id, &self.identity, job.run_no, settlement, journal)
.await
{
tracing::warn!(%error, "release failed");
}
}
async fn settle(&self, finished: FinishedRun) -> Result<(), Error> {
let now = Utc::now();
let next_failures = finished.failure_count.saturating_add(1);
let (settlement, note, attachment, duration) = match finished.report {
Err(error) => {
tracing::error!(%error, "job could not be dispatched; marking dead");
let settlement = Settlement::Dead {
finished_at: now,
failure_count: next_failures,
};
(settlement, Some(error.to_string()), None, Duration::ZERO)
}
Ok(report) => {
let settlement = self.settlement_for(&report, finished.id, now, next_failures);
(
settlement,
note_for(&report.result),
report.attachment,
report.duration,
)
}
};
let outcome = journal_outcome_for(&settlement);
crate::observability::settled(&finished.kind, outcome, duration);
let journal = JournalAppend {
kind: finished.kind,
run_no: finished.run_no,
recorded_at: now,
outcome,
note,
attachment,
};
self.store
.settle(
finished.id,
&self.identity,
finished.run_no,
settlement,
journal,
)
.await?;
Ok(())
}
fn settlement_for(
&self,
report: &RunReport,
id: Ulid,
now: DateTime<Utc>,
next_failures: i32,
) -> Settlement {
match &report.result {
Ok(Outcome::Completed { .. }) => Settlement::Complete { finished_at: now },
Ok(Outcome::Pause { resume_in, .. }) => Settlement::Pause {
visible_at: add_duration(now, *resume_in),
carry: report.carry.clone(),
},
Err(error) if error.is_permanent() => Settlement::Dead {
finished_at: now,
failure_count: next_failures,
},
Err(_retryable) => self.retry_or_backstop(report, id, now, next_failures),
}
}
fn retry_or_backstop(
&self,
report: &RunReport,
id: Ulid,
now: DateTime<Utc>,
next_failures: i32,
) -> Settlement {
let attempt = next_failures.max(0) as u32;
if let Some(max) = self.config.backstop
&& attempt >= max
{
return Settlement::Dead {
finished_at: now,
failure_count: next_failures,
};
}
let backoff = report.backoff.unwrap_or(self.config.backoff);
let delay = retry_delay(&backoff, self.config.jitter_fraction, attempt, id);
Settlement::Retry {
visible_at: add_duration(now, delay),
failure_count: next_failures,
carry: report.carry.clone(),
}
}
}
struct FinishedRun {
id: Ulid,
kind: String,
run_no: u32,
failure_count: i32,
report: Result<RunReport, Error>,
}
struct InflightJob {
id: Ulid,
kind: String,
run_no: u32,
}
fn journal_outcome_for(settlement: &Settlement) -> JournalOutcome {
match settlement {
Settlement::Complete { .. } => JournalOutcome::Completed,
Settlement::Pause { .. } => JournalOutcome::Paused,
Settlement::Retry { .. } => JournalOutcome::Retried,
Settlement::Dead { .. } => JournalOutcome::Dead,
Settlement::Release { .. } => JournalOutcome::Released,
}
}
fn note_for(result: &Result<Outcome, crate::outcome::TaskError>) -> Option<String> {
match result {
Ok(Outcome::Completed { note }) | Ok(Outcome::Pause { note, .. }) => note.clone(),
Err(error) => Some(error.message().to_owned()),
}
}
fn floor_for(counter: u64, ratio: Option<u32>) -> i16 {
match ratio {
Some(r) if r >= 2 => {
let r = u64::from(r);
if counter.is_multiple_of(r * r) {
FLOOR_LOW
} else if counter.is_multiple_of(r) {
FLOOR_NORMAL
} else {
FLOOR_ALL
}
}
_ => FLOOR_ALL,
}
}
fn release_slot(by_kind: &mut HashMap<String, usize>, kind: &str) {
if let Some(count) = by_kind.get_mut(kind) {
*count = count.saturating_sub(1);
if *count == 0 {
by_kind.remove(kind);
}
}
}
fn add_duration(now: DateTime<Utc>, delta: Duration) -> DateTime<Utc> {
match chrono::Duration::from_std(delta) {
Ok(delta) => now
.checked_add_signed(delta)
.unwrap_or(DateTime::<Utc>::MAX_UTC),
Err(_) => DateTime::<Utc>::MAX_UTC,
}
}
fn worker_identity() -> String {
let host = std::env::var("HOSTNAME")
.ok()
.filter(|h| !h.is_empty())
.unwrap_or_else(|| "unknown".to_owned());
format!("{host}:{pid}", pid = std::process::id())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::Context;
use crate::outcome::TaskError;
use crate::queue::Queue;
use crate::test_support::FakeStore;
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
fn clamp_lease_floors_and_caps_degenerate_values() {
assert_eq!(clamp_lease(Duration::ZERO), MIN_LEASE);
assert_eq!(clamp_lease(Duration::from_secs(u64::MAX)), MAX_LEASE);
let mid = Duration::from_secs(60);
assert_eq!(clamp_lease(mid), mid);
}
#[test]
fn builder_clamps_degenerate_lease_and_poll_max() {
let store = Arc::new(FakeStore::new());
let clamped = Worker::builder((), store.clone())
.lease(Duration::ZERO)
.poll_max(Duration::ZERO);
assert_eq!(clamped.config.lease, MIN_LEASE);
assert_eq!(clamped.config.poll_max, MIN_POLL_MAX);
let capped = Worker::builder((), store).lease(Duration::from_secs(u64::MAX));
assert_eq!(capped.config.lease, MAX_LEASE);
}
#[derive(Clone)]
struct Counters {
active: Arc<AtomicUsize>,
peak: Arc<AtomicUsize>,
ran: Arc<AtomicUsize>,
}
impl Counters {
fn new() -> Counters {
Counters {
active: Arc::new(AtomicUsize::new(0)),
peak: Arc::new(AtomicUsize::new(0)),
ran: Arc::new(AtomicUsize::new(0)),
}
}
}
#[derive(Serialize, Deserialize)]
struct SlowJob;
impl crate::task::Task for SlowJob {
const KIND: &'static str = "slow";
type Carry = ();
}
impl Handler<Counters> for SlowJob {
async fn handle(
&self,
_ctx: &mut Context<()>,
state: &Counters,
) -> Result<Outcome, TaskError> {
let now = state.active.fetch_add(1, Ordering::SeqCst) + 1;
state.peak.fetch_max(now, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(40)).await;
state.active.fetch_sub(1, Ordering::SeqCst);
state.ran.fetch_add(1, Ordering::SeqCst);
Ok(Outcome::completed())
}
}
async fn wait_until(mut cond: impl FnMut() -> bool) {
for _ in 0..200 {
if cond() {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("condition not met within the deadline");
}
#[tokio::test]
async fn loop_completes_all_jobs_within_the_concurrency_bound() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
for _ in 0..6 {
queue.enqueue(SlowJob).await.expect("enqueue");
}
let counters = Counters::new();
let worker = Worker::builder(counters.clone(), Arc::new(store.clone()))
.register::<SlowJob>()
.concurrency(2)
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| counters.ran.load(Ordering::SeqCst) == 6).await;
shutdown.cancel();
handle.await.expect("worker loop joins");
assert_eq!(counters.ran.load(Ordering::SeqCst), 6);
assert!(
counters.peak.load(Ordering::SeqCst) <= 2,
"peak concurrency {} exceeded the bound of 2",
counters.peak.load(Ordering::SeqCst)
);
assert_eq!(store.count(crate::store::Status::Completed), 6);
}
#[tokio::test]
async fn shutdown_settles_a_handler_that_finished_before_the_abort() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(SlowJob).await.expect("enqueue");
let worker = Worker::builder(Counters::new(), Arc::new(store.clone()))
.register::<SlowJob>()
.build();
let claimed = store
.claim_next(
&["slow".to_owned()],
FLOOR_ALL,
worker.config.lease,
&worker.identity,
)
.await
.expect("claim succeeds")
.expect("a claimable job");
assert_eq!(claimed.id, id);
let shutdown = CancellationToken::new();
let mut running = JoinSet::new();
let run_no = claimed.run_count;
let kind = claimed.kind.clone();
let task_id = worker.spawn(&mut running, claimed, Vec::new(), &shutdown);
let mut inflight = HashMap::new();
inflight.insert(task_id, InflightJob { id, kind, run_no });
tokio::time::sleep(Duration::from_millis(120)).await;
worker.force_finish(running, inflight).await;
assert_eq!(
store.count(crate::store::Status::Completed),
1,
"a handler that finished before the abort must be settled, not re-pended"
);
assert_eq!(
store.count(crate::store::Status::Pending),
0,
"the completed job must not be released back to pending"
);
}
#[test]
fn add_duration_saturates_to_the_far_future_on_overflow() {
let now = Utc::now();
let parked = add_duration(now, Duration::from_secs(u64::MAX));
assert!(
parked > now + chrono::Duration::days(365_000),
"an overflowing delay must saturate to the far future, got {parked}"
);
let near_max = DateTime::<Utc>::MAX_UTC - chrono::Duration::days(1);
let pushed = add_duration(near_max, Duration::from_secs(7 * 24 * 60 * 60));
assert!(
pushed >= near_max,
"an add that overflows the representable range must not move backward"
);
}
#[tokio::test]
async fn settle_rejects_a_stale_run_after_same_identity_reclaim() {
use crate::store::{JournalAppend, JournalOutcome, Settlement, Status, Store};
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(SlowJob).await.expect("enqueue");
let kinds = vec!["slow".to_owned()];
let make_journal = |run_no: u32, outcome: JournalOutcome| JournalAppend {
kind: "slow".to_owned(),
run_no,
recorded_at: Utc::now(),
outcome,
note: None,
attachment: None,
};
let first = store
.claim_next(&kinds, FLOOR_ALL, Duration::from_millis(40), "worker-a")
.await
.expect("claim")
.expect("a job");
assert_eq!(first.run_count, 1);
tokio::time::sleep(Duration::from_millis(70)).await;
let recovered = store
.recover(
id,
Utc::now(),
1,
first.run_count,
make_journal(1, JournalOutcome::StaleRecovered),
)
.await
.expect("recover");
assert!(recovered);
let second = store
.claim_next(&kinds, FLOOR_ALL, Duration::from_secs(60), "worker-a")
.await
.expect("claim")
.expect("a job");
assert_eq!(second.run_count, 2);
let stale = store
.settle(
id,
"worker-a",
first.run_count,
Settlement::Complete {
finished_at: Utc::now(),
},
make_journal(1, JournalOutcome::Completed),
)
.await
.expect("settle stale");
assert!(
!stale,
"a stale run must not settle under a matching identity"
);
assert_eq!(store.job(id).expect("job").status, Status::Claimed);
let current = store
.settle(
id,
"worker-a",
second.run_count,
Settlement::Complete {
finished_at: Utc::now(),
},
make_journal(2, JournalOutcome::Completed),
)
.await
.expect("settle current");
assert!(current, "the live claim settles");
assert_eq!(store.count(Status::Completed), 1);
}
#[tokio::test]
async fn recover_rejects_a_stale_snapshot_after_the_claim_advanced() {
use crate::store::{JournalAppend, JournalOutcome, Status, Store};
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(SlowJob).await.expect("enqueue");
let kinds = vec!["slow".to_owned()];
let make_journal = |run_no: u32| JournalAppend {
kind: "slow".to_owned(),
run_no,
recorded_at: Utc::now(),
outcome: JournalOutcome::StaleRecovered,
note: None,
attachment: None,
};
let snapshot = store
.claim_next(&kinds, FLOOR_ALL, Duration::from_millis(40), "worker-a")
.await
.expect("claim")
.expect("a job");
assert_eq!(snapshot.run_count, 1);
tokio::time::sleep(Duration::from_millis(70)).await;
assert!(
store
.recover(id, Utc::now(), 1, snapshot.run_count, make_journal(1))
.await
.expect("first recover")
);
let reclaim = store
.claim_next(&kinds, FLOOR_ALL, Duration::from_millis(40), "worker-a")
.await
.expect("reclaim")
.expect("a job");
assert_eq!(reclaim.run_count, 2);
tokio::time::sleep(Duration::from_millis(70)).await;
let stale_applied = store
.recover(id, Utc::now(), 1, snapshot.run_count, make_journal(1))
.await
.expect("stale recover");
assert!(
!stale_applied,
"a recovery from a superseded snapshot epoch must not apply"
);
let live_applied = store
.recover(id, Utc::now(), 2, reclaim.run_count, make_journal(2))
.await
.expect("live recover");
assert!(live_applied, "the current claim epoch recovers");
assert_eq!(store.job(id).expect("job").status, Status::Pending);
}
#[derive(Serialize, Deserialize)]
struct Unencodable;
impl crate::task::Task for Unencodable {
const KIND: &'static str = "unencodable";
type Carry = std::collections::BTreeMap<(i32, i32), i32>;
}
impl Handler<()> for Unencodable {
async fn handle(
&self,
ctx: &mut Context<std::collections::BTreeMap<(i32, i32), i32>>,
_state: &(),
) -> Result<Outcome, TaskError> {
ctx.carry_mut().insert((1, 2), 3);
Ok(Outcome::completed())
}
}
#[tokio::test]
async fn completed_run_with_unencodable_carry_is_dead_with_an_accurate_note() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Unencodable).await.expect("enqueue");
let worker = Worker::builder((), Arc::new(store.clone()))
.register::<Unencodable>()
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| store.count(crate::store::Status::Dead) == 1).await;
shutdown.cancel();
handle.await.expect("worker joins");
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Dead);
let note = store
.journal(id)
.await
.expect("journal")
.last()
.expect("a journal entry")
.note
.clone()
.unwrap_or_default();
assert!(
note.contains("carry could not be serialized"),
"the note must describe the post-run carry serialization failure, got: {note:?}"
);
}
#[tokio::test]
async fn worker_with_no_kinds_returns_immediately() {
let store = FakeStore::new();
let worker: Worker<Counters> = Worker::builder(Counters::new(), Arc::new(store)).build();
worker.run(CancellationToken::new()).await;
}
proptest::proptest! {
#[test]
fn rotation_serves_every_tier_in_a_window(r in 2u32..8, start in 0u64..10_000) {
let window = u64::from(r) * u64::from(r);
let floors: Vec<i16> = (start..start + window)
.map(|c| floor_for(c, Some(r)))
.collect();
proptest::prop_assert!(floors.contains(&FLOOR_LOW), "Low never reserved");
proptest::prop_assert!(floors.contains(&FLOOR_ALL), "high-first never admitted");
}
}
#[test]
fn floor_rotation_reserves_lower_tiers() {
assert_eq!(floor_for(1, Some(2)), FLOOR_ALL);
assert_eq!(floor_for(2, Some(2)), FLOOR_NORMAL);
assert_eq!(floor_for(3, Some(2)), FLOOR_ALL);
assert_eq!(floor_for(4, Some(2)), FLOOR_LOW);
assert_eq!(floor_for(2, None), FLOOR_ALL);
assert_eq!(floor_for(4, None), FLOOR_ALL);
assert_eq!(floor_for(4, Some(1)), FLOOR_ALL);
}
#[derive(Clone)]
struct CapState {
active: Arc<AtomicUsize>,
peak: Arc<AtomicUsize>,
ran: Arc<AtomicUsize>,
}
#[derive(Serialize, Deserialize)]
struct Capped;
impl crate::task::Task for Capped {
const KIND: &'static str = "capped";
type Carry = ();
}
impl Handler<CapState> for Capped {
async fn handle(
&self,
_ctx: &mut Context<()>,
state: &CapState,
) -> Result<Outcome, TaskError> {
let now = state.active.fetch_add(1, Ordering::SeqCst) + 1;
state.peak.fetch_max(now, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(40)).await;
state.active.fetch_sub(1, Ordering::SeqCst);
state.ran.fetch_add(1, Ordering::SeqCst);
Ok(Outcome::completed())
}
}
#[tokio::test]
async fn per_kind_cap_bounds_in_flight_below_concurrency() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
for _ in 0..8 {
queue.enqueue(Capped).await.expect("enqueue");
}
let state = CapState {
active: Arc::new(AtomicUsize::new(0)),
peak: Arc::new(AtomicUsize::new(0)),
ran: Arc::new(AtomicUsize::new(0)),
};
let worker = Worker::builder(state.clone(), Arc::new(store.clone()))
.concurrency(8)
.register_capped::<Capped>(2)
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| state.ran.load(Ordering::SeqCst) == 8).await;
shutdown.cancel();
handle.await.expect("worker joins");
assert!(
state.peak.load(Ordering::SeqCst) <= 2,
"peak in-flight {} exceeded the cap of 2",
state.peak.load(Ordering::SeqCst)
);
}
#[derive(Clone)]
enum Mode {
AlwaysRetryable,
AlwaysPermanent,
PauseThenComplete,
}
#[derive(Serialize, Deserialize)]
struct Routed;
impl crate::task::Task for Routed {
const KIND: &'static str = "routed";
type Carry = u32;
}
impl Handler<Mode> for Routed {
async fn handle(&self, ctx: &mut Context<u32>, mode: &Mode) -> Result<Outcome, TaskError> {
match mode {
Mode::AlwaysRetryable => {
Err(TaskError::retryable(std::io::Error::other("transient")))
}
Mode::AlwaysPermanent => Err(TaskError::permanent("gone for good")),
Mode::PauseThenComplete => {
if *ctx.carry() == 0 {
*ctx.carry_mut() = 1;
Ok(Outcome::pause_in(Duration::ZERO))
} else {
Ok(Outcome::completed())
}
}
}
}
}
async fn run_until<F: FnMut() -> bool>(
store: &FakeStore,
mode: Mode,
backstop: Option<u32>,
cond: F,
) {
let worker = Worker::builder(mode, Arc::new(store.clone()))
.register::<Routed>()
.backstop(backstop)
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(cond).await;
shutdown.cancel();
handle.await.expect("worker joins");
}
#[tokio::test]
async fn permanent_failure_goes_straight_to_dead() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Routed).await.expect("enqueue");
let store_for_cond = store.clone();
run_until(&store, Mode::AlwaysPermanent, Some(20), move || {
store_for_cond.count(crate::store::Status::Dead) == 1
})
.await;
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Dead);
assert_eq!(job.failure_count, 1);
assert!(job.finished_at.is_some());
}
#[tokio::test]
async fn retryable_failures_reach_dead_at_the_backstop() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Routed).await.expect("enqueue");
let store_for_cond = store.clone();
run_until(&store, Mode::AlwaysRetryable, Some(2), move || {
store_for_cond.count(crate::store::Status::Dead) == 1
})
.await;
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Dead);
assert_eq!(job.failure_count, 2);
}
#[derive(Serialize, Deserialize)]
struct Cooperative;
impl crate::task::Task for Cooperative {
const KIND: &'static str = "cooperative";
type Carry = ();
}
impl Handler<()> for Cooperative {
async fn handle(&self, ctx: &mut Context<()>, _state: &()) -> Result<Outcome, TaskError> {
ctx.cancelled().await;
Ok(Outcome::pause_in(Duration::from_secs(1)))
}
}
#[tokio::test]
async fn cooperative_handler_settles_normally_on_shutdown() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Cooperative).await.expect("enqueue");
let worker = Worker::builder((), Arc::new(store.clone()))
.register::<Cooperative>()
.shutdown_timeout(Duration::from_secs(5))
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| {
store
.job(id)
.is_some_and(|j| j.status == crate::store::Status::Claimed)
})
.await;
shutdown.cancel();
handle.await.expect("worker joins");
let job = store.job(id).expect("job");
assert_eq!(job.status, crate::store::Status::Pending);
let journal = store
.journal(id)
.await
.expect("journal")
.into_iter()
.map(|e| e.outcome)
.collect::<Vec<_>>();
assert_eq!(journal, vec![crate::store::JournalOutcome::Paused]);
}
#[tokio::test]
async fn pause_repends_without_failure_and_persists_carry() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Routed).await.expect("enqueue");
let store_for_cond = store.clone();
run_until(&store, Mode::PauseThenComplete, Some(20), move || {
store_for_cond.count(crate::store::Status::Completed) == 1
})
.await;
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Completed);
assert_eq!(job.run_count, 2);
assert_eq!(job.failure_count, 0);
assert_eq!(job.carry, serde_json::json!(1));
}
#[derive(Serialize, Deserialize)]
struct Panicky;
impl crate::task::Task for Panicky {
const KIND: &'static str = "panicky";
type Carry = ();
}
impl Handler<Arc<AtomicUsize>> for Panicky {
async fn handle(
&self,
_ctx: &mut Context<()>,
state: &Arc<AtomicUsize>,
) -> Result<Outcome, TaskError> {
if state.fetch_add(1, Ordering::SeqCst) == 0 {
panic!("boom in handler");
}
Ok(Outcome::completed())
}
}
#[tokio::test]
async fn panicking_handler_settles_as_a_failed_execution_then_recovers() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Panicky).await.expect("enqueue");
let runs = Arc::new(AtomicUsize::new(0));
let worker = Worker::builder(runs.clone(), Arc::new(store.clone()))
.register::<Panicky>()
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| store.count(crate::store::Status::Completed) == 1).await;
shutdown.cancel();
handle.await.expect("worker joins");
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Completed);
assert_eq!(job.failure_count, 1);
let outcomes = store
.journal(id)
.await
.expect("journal")
.into_iter()
.map(|e| e.outcome)
.collect::<Vec<_>>();
assert_eq!(
outcomes,
vec![
crate::store::JournalOutcome::Retried,
crate::store::JournalOutcome::Completed,
]
);
}
#[tokio::test]
async fn panic_policy_dead_sends_a_panicking_job_straight_to_dead() {
let store = FakeStore::new();
let queue = Queue::new(Arc::new(store.clone()));
let id = queue.enqueue(Panicky).await.expect("enqueue");
let runs = Arc::new(AtomicUsize::new(0));
let worker = Worker::builder(runs.clone(), Arc::new(store.clone()))
.register::<Panicky>()
.panic_policy(PanicPolicy::Dead)
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| store.count(crate::store::Status::Dead) == 1).await;
shutdown.cancel();
handle.await.expect("worker joins");
let job = store.job(id).expect("job exists");
assert_eq!(job.status, crate::store::Status::Dead);
assert_eq!(job.failure_count, 1);
let outcomes = store
.journal(id)
.await
.expect("journal")
.into_iter()
.map(|e| e.outcome)
.collect::<Vec<_>>();
assert_eq!(outcomes, vec![crate::store::JournalOutcome::Dead]);
}
#[tokio::test]
async fn jitter_fraction_zero_schedules_the_exact_backoff_delay() {
let store = FakeStore::new();
let id = Ulid::from_string("01ARZ3NDEKTSV4RRFFQ69G5FAV").expect("valid ULID");
let now = Utc::now();
let job = crate::store::NewJob {
id,
kind: "routed".to_owned(),
payload: serde_json::Value::Null,
priority: 1,
created_at: now,
visible_at: now,
carry: serde_json::json!(0),
dedup_key: None,
};
store.enqueue(&job).await.expect("enqueue");
let backoff = Backoff::new(Duration::from_secs(100), Duration::from_secs(300));
let worker = Worker::builder(Mode::AlwaysRetryable, Arc::new(store.clone()))
.register::<Routed>()
.backoff(backoff)
.jitter_fraction(0.0)
.backstop(Some(50))
.build();
let shutdown = CancellationToken::new();
let handle = tokio::spawn(worker.run(shutdown.clone()));
wait_until(|| store.job(id).is_some_and(|j| j.failure_count >= 3)).await;
shutdown.cancel();
handle.await.expect("worker joins");
let job = store.job(id).expect("job exists");
let entries = store.journal(id).await.expect("journal");
let third = entries.last().expect("a third journal entry");
let scheduled = (job.visible_at - third.recorded_at)
.to_std()
.expect("non-negative delay");
let expected = crate::backoff::retry_delay(&backoff, 0.0, 3, id);
assert_eq!(scheduled, expected);
assert_eq!(expected, Duration::from_secs(100));
}
}