diff --git a/CHANGELOG.md b/CHANGELOG.md index 62ac6409..48d79be1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,12 @@ underlying Rust crates are released together, in lock-step. ### Added +- **Fast in-flight recovery on worker death.** When a worker process crashes mid-job, + a surviving worker now requeues its `Running` jobs within ~30s of the missed + heartbeat (via retry / dead-letter) instead of waiting the job's full timeout. The + scheduler claims execution under its real `worker_id` and a new maintenance step + reclaims orphaned claims atomically, so concurrent survivors never double-rescue. + Works across Python, Node, and Java; SQLite/Postgres/Redis. No API change. - **Indexes for dead-letter and filter paths.** `dead_letter` (previously unindexed) gains indexes on `failed_at` and `task_name`; `jobs` gains `(task_name, status)` and partial indexes on `expires_at` / `namespace`; `archived_jobs` gains diff --git a/crates/taskito-core/src/scheduler/maintenance.rs b/crates/taskito-core/src/scheduler/maintenance.rs index 687738e7..400dbb1c 100644 --- a/crates/taskito-core/src/scheduler/maintenance.rs +++ b/crates/taskito-core/src/scheduler/maintenance.rs @@ -3,7 +3,7 @@ use log::{error, info, warn}; use crate::error::Result; use crate::job::{now_millis, NewJob}; use crate::periodic::{next_cron_time, next_cron_time_tz}; -use crate::storage::Storage; +use crate::storage::{Storage, DEAD_WORKER_THRESHOLD_MS}; use super::{JobResult, Scheduler}; @@ -45,6 +45,64 @@ impl Scheduler { })?; } + // Fast-path recovery: requeue jobs whose worker died, without waiting + // for their full timeout. Log-and-continue so a failure here never + // aborts the rest of the sweep. + if let Err(e) = self.recover_orphaned_jobs(now) { + warn!("recover_orphaned_jobs error: {e}"); + } + + Ok(()) + } + + /// Requeue `Running` jobs whose claiming worker is no longer alive (crashed + /// without finishing). A surviving scheduler detects the dead owner via the + /// heartbeat table, atomically reclaims the orphaned claim — so exactly one + /// survivor rescues each job — then routes it through the normal + /// retry/dead-letter path. + fn recover_orphaned_jobs(&self, now: i64) -> Result<()> { + // Live owners = workers with a fresh heartbeat, plus self: a scheduler + // must never orphan its own in-flight jobs (covers the startup window + // before its first heartbeat row is written). + let mut live: Vec = self + .storage + .list_workers()? + .into_iter() + .filter(|w| w.last_heartbeat >= now - DEAD_WORKER_THRESHOLD_MS) + .map(|w| w.worker_id) + .collect(); + live.push(self.claim_owner.clone()); + + for (job, dead_owner) in self.storage.reap_orphaned_jobs(&live, now)? { + // Atomic election: only the survivor that wins the claim transfer + // requeues the job, so concurrent schedulers can't double-retry it. + match self + .storage + .reclaim_execution(&job.id, &dead_owner, &self.claim_owner) + { + Ok(true) => { + let error = format!("worker {dead_owner} died; recovering in-flight job"); + if let Err(e) = self.handle_result(JobResult::Failure { + job_id: job.id.clone(), + error, + retry_count: job.retry_count, + max_retries: job.max_retries, + task_name: job.task_name.clone(), + wall_time_ns: 0, + should_retry: true, + timed_out: false, + }) { + warn!( + "recover_orphaned_jobs: handle_result failed for {}: {e}", + job.id + ); + } + } + Ok(false) => {} // another survivor won the race + Err(e) => warn!("reclaim_execution failed for {}: {e}", job.id), + } + } + Ok(()) } diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 96ef771b..8df91c5a 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -149,6 +149,10 @@ pub struct Scheduler { shutdown: Arc, paused_cache: Mutex<(HashSet, Instant)>, namespace: Option, + /// Owner id recorded on execution claims. Defaults to a placeholder; the + /// binding sets it to the process `worker_id` so dead-worker recovery can + /// attribute orphaned claims. See [`Self::set_claim_owner`]. + claim_owner: String, /// Wake source for push-dispatch, installed before `run()`. Taken (moved /// out) once when the loop starts. `None` means the loop polls as today. #[cfg(feature = "push-dispatch")] @@ -192,6 +196,7 @@ impl Scheduler { shutdown: Arc::new(Notify::new()), paused_cache: Mutex::new((HashSet::new(), Instant::now())), namespace, + claim_owner: poller::SCHEDULER_CLAIM_OWNER.to_string(), #[cfg(feature = "push-dispatch")] wake_source: Mutex::new(None), #[cfg(feature = "push-dispatch")] @@ -203,6 +208,14 @@ impl Scheduler { &self.storage } + /// Set the execution-claim owner to this process's `worker_id`. Bindings + /// call this right after construction so dead-worker recovery can identify + /// which worker owns each in-flight job. Must match the id passed to + /// `register_worker`/`heartbeat`. + pub fn set_claim_owner(&mut self, worker_id: String) { + self.claim_owner = worker_id; + } + pub fn shutdown_handle(&self) -> Arc { self.shutdown.clone() } @@ -983,6 +996,108 @@ mod tests { assert_eq!(reaped.retry_count, 1); } + fn retry_task_config(max_retries: i32) -> TaskConfig { + TaskConfig { + retry_policy: RetryPolicy { + max_retries, + base_delay_ms: 100, + max_delay_ms: 1000, + custom_delays_ms: None, + }, + rate_limit: None, + circuit_breaker: None, + max_concurrent: None, + } + } + + /// A Running job whose claim owner is a dead worker (no live heartbeat) is + /// requeued by `recover_orphaned_jobs` without waiting for its timeout. + #[test] + fn test_recover_orphaned_jobs_requeues_dead_worker_job() { + let mut scheduler = test_scheduler(); + scheduler.set_claim_owner("survivor".to_string()); + scheduler.register_task("orphan_task".to_string(), retry_task_config(3)); + + let job = scheduler.storage.enqueue(make_job("orphan_task")).unwrap(); + // Move to Running and record a claim owned by a worker that never + // heartbeats (simulating a crashed worker). No workers are registered, + // so the live set is just the survivor → this claim is orphaned. + scheduler + .storage + .dequeue("default", now_millis() + 1000, None) + .unwrap() + .unwrap(); + assert!(scheduler + .storage + .claim_execution(&job.id, "dead-worker") + .unwrap()); + + scheduler.reap_stale().unwrap(); + + let recovered = scheduler.storage.get_job(&job.id).unwrap().unwrap(); + assert_eq!(recovered.status, JobStatus::Pending); + assert_eq!(recovered.retry_count, 1); + // The orphaned claim was reclaimed then cleared on requeue. + assert!(scheduler + .storage + .list_claims_by_worker("dead-worker") + .unwrap() + .is_empty()); + } + + /// A scheduler must never orphan its OWN in-flight jobs (self-rescue guard), + /// even before its first heartbeat row exists. + #[test] + fn test_recover_orphaned_jobs_skips_own_claims() { + let mut scheduler = test_scheduler(); + scheduler.set_claim_owner("survivor".to_string()); + scheduler.register_task("self_task".to_string(), retry_task_config(3)); + + let job = scheduler.storage.enqueue(make_job("self_task")).unwrap(); + scheduler + .storage + .dequeue("default", now_millis() + 1000, None) + .unwrap() + .unwrap(); + // Claim owned by this scheduler's own id. + assert!(scheduler + .storage + .claim_execution(&job.id, "survivor") + .unwrap()); + + scheduler.reap_stale().unwrap(); + + let still = scheduler.storage.get_job(&job.id).unwrap().unwrap(); + assert_eq!(still.status, JobStatus::Running); + assert_eq!(still.retry_count, 0); + } + + /// An orphan with no retries left is dead-lettered, not requeued. + #[test] + fn test_recover_orphaned_jobs_dead_letters_when_exhausted() { + let mut scheduler = test_scheduler(); + scheduler.set_claim_owner("survivor".to_string()); + scheduler.register_task("exhausted_task".to_string(), retry_task_config(0)); + + let mut nj = make_job("exhausted_task"); + nj.max_retries = 0; + let job = scheduler.storage.enqueue(nj).unwrap(); + scheduler + .storage + .dequeue("default", now_millis() + 1000, None) + .unwrap() + .unwrap(); + assert!(scheduler + .storage + .claim_execution(&job.id, "dead-worker") + .unwrap()); + + scheduler.reap_stale().unwrap(); + + let dead = scheduler.storage.list_dead(10, 0).unwrap(); + assert!(dead.iter().any(|d| d.original_job_id == job.id)); + } + #[test] fn test_check_periodic() { let scheduler = test_scheduler(); diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index 1454e98c..61c575b7 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/crates/taskito-core/src/scheduler/poller.rs @@ -23,8 +23,11 @@ const CONCURRENCY_RETRY_DELAY_MS: i64 = 500; /// catches up on the next tick rather than waiting for the stale-job reaper. const CHANNEL_BACKPRESSURE_RETRY_DELAY_MS: i64 = 100; -/// Worker identifier recorded on execution claims taken by the scheduler. -const SCHEDULER_CLAIM_OWNER: &str = "scheduler"; +/// Default execution-claim owner when the binding has not set the worker's id +/// (tests / embedders that don't register a worker). Production bindings call +/// [`Scheduler::set_claim_owner`] with the real `worker_id` so dead-worker +/// recovery can attribute claims. +pub(super) const SCHEDULER_CLAIM_OWNER: &str = "scheduler"; impl Scheduler { pub(super) fn try_dispatch(&self, job_tx: &tokio::sync::mpsc::Sender) -> Result { @@ -228,7 +231,7 @@ impl Scheduler { /// claim was actually taken; `Ok(false)` if it was already claimed by /// another scheduler **or** the claim attempt errored. fn claim_for_dispatch(&self, job: &Job) -> Result { - match self.storage.claim_execution(&job.id, SCHEDULER_CLAIM_OWNER) { + match self.storage.claim_execution(&job.id, &self.claim_owner) { Ok(true) => Ok(true), Ok(false) => Ok(false), Err(e) => { diff --git a/crates/taskito-core/src/storage/diesel_common/jobs.rs b/crates/taskito-core/src/storage/diesel_common/jobs.rs index d11e39bc..022a1f6a 100644 --- a/crates/taskito-core/src/storage/diesel_common/jobs.rs +++ b/crates/taskito-core/src/storage/diesel_common/jobs.rs @@ -1274,6 +1274,56 @@ macro_rules! impl_diesel_job_ops { .collect()) } + /// Running jobs whose execution-claim owner is no longer alive (the + /// worker that claimed them died). Read-only, like `reap_stale_jobs`: + /// the scheduler atomically reclaims and requeues each one. Two + /// indexed queries rather than a join, since `jobs` and + /// `execution_claims` are not declared joinable. + pub fn reap_orphaned_jobs( + &self, + live_owner_ids: &[String], + _now: i64, + ) -> Result> { + // Defensive: never treat every claim as orphaned. The caller + // always includes its own owner, so this is unreachable in + // practice but guards against a `NOT IN (empty)` sweep. + if live_owner_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut conn = self.conn()?; + + // Claims owned by a worker not in the live set. + let orphan_claims: Vec<(String, String)> = execution_claims::table + .filter(diesel::dsl::not( + execution_claims::worker_id.eq_any(live_owner_ids), + )) + .select((execution_claims::job_id, execution_claims::worker_id)) + .load(&mut conn)?; + if orphan_claims.is_empty() { + return Ok(Vec::new()); + } + + let job_ids: Vec = orphan_claims.iter().map(|(id, _)| id.clone()).collect(); + let owner_by_job: std::collections::HashMap = + orphan_claims.into_iter().collect(); + + // Of those, the ones still Running (blob-free narrow row). + let rows: Vec = jobs::table + .filter(jobs::id.eq_any(&job_ids)) + .filter(jobs::status.eq(JobStatus::Running as i32)) + .select(NarrowJobRow::as_select()) + .load(&mut conn)?; + + Ok(rows + .into_iter() + .map(|narrow| { + let owner = owner_by_job.get(&narrow.id).cloned().unwrap_or_default(); + (Job::from_narrow(narrow, Vec::new(), None), owner) + }) + .collect()) + } + /// Record an error for a job attempt. pub fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { let mut conn = self.conn()?; diff --git a/crates/taskito-core/src/storage/diesel_common/locks.rs b/crates/taskito-core/src/storage/diesel_common/locks.rs index 697063a6..0db5b2ba 100644 --- a/crates/taskito-core/src/storage/diesel_common/locks.rs +++ b/crates/taskito-core/src/storage/diesel_common/locks.rs @@ -75,6 +75,35 @@ macro_rules! impl_diesel_lock_ops { Ok(()) } + /// Atomically transfer an existing claim from `expected_owner` to + /// `new_owner`. The `job_id` PK plus the `worker_id = expected_owner` + /// filter serialize concurrent rescuers: the first UPDATE rewrites the + /// owner, every other rescuer's filter no longer matches → 0 rows. + /// `claim_execution` is INSERT-only and cannot reclaim, so this is a + /// distinct primitive. + pub fn reclaim_execution( + &self, + job_id: &str, + expected_owner: &str, + new_owner: &str, + ) -> Result { + let mut conn = self.conn()?; + let now = now_millis(); + + let affected = diesel::update( + execution_claims::table + .filter(execution_claims::job_id.eq(job_id)) + .filter(execution_claims::worker_id.eq(expected_owner)), + ) + .set(( + execution_claims::worker_id.eq(new_owner), + execution_claims::claimed_at.eq(now), + )) + .execute(&mut conn)?; + + Ok(affected > 0) + } + /// Purge old execution claims. Returns count removed. pub fn purge_execution_claims(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index 30722867..6bbc9207 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -212,6 +212,13 @@ macro_rules! impl_storage { fn reap_stale_jobs(&self, now: i64) -> $crate::error::Result> { self.reap_stale_jobs(now) } + fn reap_orphaned_jobs( + &self, + live_owner_ids: &[String], + now: i64, + ) -> $crate::error::Result> { + self.reap_orphaned_jobs(live_owner_ids, now) + } fn record_error( &self, job_id: &str, @@ -539,6 +546,14 @@ macro_rules! impl_storage { ) -> $crate::error::Result { self.purge_execution_claims(older_than_ms) } + fn reclaim_execution( + &self, + job_id: &str, + expected_owner: &str, + new_owner: &str, + ) -> $crate::error::Result { + self.reclaim_execution(job_id, expected_owner, new_owner) + } fn count_running_by_task( &self, task_name: &str, @@ -763,6 +778,13 @@ impl Storage for StorageBackend { fn reap_stale_jobs(&self, now: i64) -> Result> { delegate!(self, reap_stale_jobs, now) } + fn reap_orphaned_jobs( + &self, + live_owner_ids: &[String], + now: i64, + ) -> Result> { + delegate!(self, reap_orphaned_jobs, live_owner_ids, now) + } fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { delegate!(self, record_error, job_id, attempt, error) } @@ -1022,6 +1044,14 @@ impl Storage for StorageBackend { fn purge_execution_claims(&self, older_than_ms: i64) -> Result { delegate!(self, purge_execution_claims, older_than_ms) } + fn reclaim_execution( + &self, + job_id: &str, + expected_owner: &str, + new_owner: &str, + ) -> Result { + delegate!(self, reclaim_execution, job_id, expected_owner, new_owner) + } fn count_running_by_task(&self, task_name: &str) -> Result { delegate!(self, count_running_by_task, task_name) } diff --git a/crates/taskito-core/src/storage/postgres/jobs.rs b/crates/taskito-core/src/storage/postgres/jobs.rs index 10cdd9bc..3b1b3454 100644 --- a/crates/taskito-core/src/storage/postgres/jobs.rs +++ b/crates/taskito-core/src/storage/postgres/jobs.rs @@ -3,7 +3,8 @@ use diesel::prelude::*; use super::super::models::*; use super::super::schema::{ - archived_jobs, job_dependencies, job_errors, jobs, replay_history, task_logs, task_metrics, + archived_jobs, execution_claims, job_dependencies, job_errors, jobs, replay_history, task_logs, + task_metrics, }; use super::PostgresStorage; use crate::error::{QueueError, Result}; diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/maintenance.rs b/crates/taskito-core/src/storage/redis_backend/jobs/maintenance.rs index a111cf44..45eb6223 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/maintenance.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/maintenance.rs @@ -87,6 +87,48 @@ impl RedisStorage { Ok(stale) } + /// Running jobs whose execution-claim owner is not in `live_owner_ids` (its + /// worker died). Read-only — paired with the dead owner so the scheduler can + /// atomically reclaim before requeuing. Iterates the bounded Running set. + pub fn reap_orphaned_jobs( + &self, + live_owner_ids: &[String], + _now: i64, + ) -> Result> { + if live_owner_ids.is_empty() { + return Ok(Vec::new()); + } + let live: std::collections::HashSet<&str> = + live_owner_ids.iter().map(String::as_str).collect(); + + let mut conn = self.conn()?; + let status_key = self.key(&["jobs", "status", &(JobStatus::Running as i32).to_string()]); + let job_ids: Vec = conn.smembers(&status_key).map_err(map_err)?; + + let mut orphaned = Vec::new(); + for id in &job_ids { + let claim_key = self.key(&["exec_claim", id]); + let claim: Option = conn.get(&claim_key).map_err(map_err)?; + // Claim value is "{owner}:{ts}". The timestamp is a numeric suffix, so + // split on the LAST ':' — the owner itself may contain ':' (e.g. + // "host:pid"). No claim → time-based reap handles it. + let owner = match claim.as_deref() { + Some(v) => v.rsplit_once(':').map(|(o, _)| o).unwrap_or(v).to_string(), + None => continue, + }; + if live.contains(owner.as_str()) { + continue; + } + if let Some(job) = self.load_job(&mut conn, id)? { + if job.status == JobStatus::Running { + orphaned.push((job, owner)); + } + } + } + + Ok(orphaned) + } + pub fn expire_pending_jobs(&self, now: i64) -> Result { let mut conn = self.conn()?; let status_key = self.key(&["jobs", "status", &(JobStatus::Pending as i32).to_string()]); diff --git a/crates/taskito-core/src/storage/redis_backend/locks.rs b/crates/taskito-core/src/storage/redis_backend/locks.rs index 346e47d0..e58a9f61 100644 --- a/crates/taskito-core/src/storage/redis_backend/locks.rs +++ b/crates/taskito-core/src/storage/redis_backend/locks.rs @@ -5,6 +5,22 @@ use crate::error::Result; use crate::job::now_millis; use crate::storage::models::LockInfoRow; +/// Lua script: atomically transfer an execution claim from `expected_owner` +/// (ARGV[2]) to `new_owner` (ARGV[3]). The claim value is "{owner}:{ts}"; +/// returns 1 only if the current owner matches. KEYS[1] = claim key, +/// KEYS[2] = the by-time index; ARGV[1] = job_id, ARGV[4] = now. +const RECLAIM_CLAIM_SCRIPT: &str = r#" + local cur = redis.call('GET', KEYS[1]) + if not cur then return 0 end + -- Owner is everything before the LAST ':' (the timestamp is a numeric + -- suffix); the owner itself may contain ':' (e.g. "host:pid"). + local owner = string.match(cur, '^(.*):') or cur + if owner ~= ARGV[2] then return 0 end + redis.call('SET', KEYS[1], ARGV[3] .. ':' .. ARGV[4], 'PX', 86400000) + redis.call('ZADD', KEYS[2], ARGV[4], ARGV[1]) + return 1 +"#; + /// Lua script: release lock only if owner matches. const RELEASE_LOCK_SCRIPT: &str = r#" local key = KEYS[1] @@ -210,6 +226,33 @@ impl RedisStorage { Ok(()) } + /// Atomically transfer a claim from `expected_owner` to `new_owner`. Returns + /// `true` only if the claim was held by `expected_owner` — the single GET/SET + /// in the Lua script serializes concurrent rescuers so exactly one wins. + pub fn reclaim_execution( + &self, + job_id: &str, + expected_owner: &str, + new_owner: &str, + ) -> Result { + let mut conn = self.conn()?; + let now = now_millis(); + let ckey = self.key(&["exec_claim", job_id]); + let index_key = self.key(&["exec_claims", "by_time"]); + + let result: i32 = redis::Script::new(RECLAIM_CLAIM_SCRIPT) + .key(&ckey) + .key(&index_key) + .arg(job_id) + .arg(expected_owner) + .arg(new_owner) + .arg(now) + .invoke(&mut conn) + .map_err(map_err)?; + + Ok(result == 1) + } + pub fn purge_execution_claims(&self, older_than_ms: i64) -> Result { let mut conn = self.conn()?; let index_key = self.key(&["exec_claims", "by_time"]); diff --git a/crates/taskito-core/src/storage/sqlite/jobs.rs b/crates/taskito-core/src/storage/sqlite/jobs.rs index d798d65c..72daaa77 100644 --- a/crates/taskito-core/src/storage/sqlite/jobs.rs +++ b/crates/taskito-core/src/storage/sqlite/jobs.rs @@ -3,7 +3,8 @@ use diesel::sqlite::SqliteConnection; use super::super::models::*; use super::super::schema::{ - archived_jobs, job_dependencies, job_errors, jobs, replay_history, task_logs, task_metrics, + archived_jobs, execution_claims, job_dependencies, job_errors, jobs, replay_history, task_logs, + task_metrics, }; use super::SqliteStorage; use crate::error::{QueueError, Result}; diff --git a/crates/taskito-core/src/storage/traits.rs b/crates/taskito-core/src/storage/traits.rs index 51519b7e..23ee1bf1 100644 --- a/crates/taskito-core/src/storage/traits.rs +++ b/crates/taskito-core/src/storage/traits.rs @@ -70,6 +70,11 @@ pub trait Storage: Send + Sync + Clone { fn purge_completed(&self, older_than_ms: i64) -> Result; fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result; fn reap_stale_jobs(&self, now: i64) -> Result>; + /// Running jobs whose execution-claim owner is not in `live_owner_ids` (the + /// worker that claimed them has died). Read-only — paired with the dead + /// owner so the caller can atomically reclaim before requeuing. + fn reap_orphaned_jobs(&self, live_owner_ids: &[String], now: i64) + -> Result>; fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()>; fn get_job_errors(&self, job_id: &str) -> Result>; fn purge_job_errors(&self, older_than_ms: i64) -> Result; @@ -218,6 +223,16 @@ pub trait Storage: Send + Sync + Clone { fn claim_execution(&self, job_id: &str, worker_id: &str) -> Result; fn complete_execution(&self, job_id: &str) -> Result<()>; fn purge_execution_claims(&self, older_than_ms: i64) -> Result; + /// Atomically transfer an existing claim from `expected_owner` to + /// `new_owner`. Returns `true` only if the claim was held by + /// `expected_owner` — the `job_id` PK serializes concurrent rescuers so + /// exactly one wins. + fn reclaim_execution( + &self, + job_id: &str, + expected_owner: &str, + new_owner: &str, + ) -> Result; // ── Per-task concurrency ────────────────────────────────────── diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index 45851dc3..6ffefc80 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -445,6 +445,115 @@ fn test_execution_claims_purge(s: &impl Storage) { s.complete_execution(fresh_job).unwrap(); } +fn test_reap_stale_jobs(s: &impl Storage) { + // A running job past its timeout is reported by reap_stale_jobs (the + // scheduler then requeues it). Within-budget jobs are left alone. + let q = "q-reap-stale"; + let mut nj = make_job(q, "stale_task"); + nj.timeout_ms = 1; + let job = s.enqueue(nj).unwrap(); + let t0 = now_millis(); + s.dequeue(q, t0, None).unwrap().unwrap(); // Running, started_at = t0 + + let stale = s.reap_stale_jobs(t0 + 1000).unwrap(); + assert!( + stale.iter().any(|j| j.id == job.id), + "a running job past its timeout must be reaped" + ); + // Clean up so this Running job doesn't bleed into later shared-instance tests. + s.complete(&job.id, None).unwrap(); +} + +fn test_reclaim_execution(s: &impl Storage) { + // Atomic claim transfer: only the rescuer expecting the current owner wins. + let job = "reclaim-job-id"; + assert!(s.claim_execution(job, "dead").unwrap()); + assert!(s.reclaim_execution(job, "dead", "rescuer").unwrap()); + // A second rescuer still expecting "dead" loses — owner is now "rescuer". + assert!(!s.reclaim_execution(job, "dead", "other").unwrap()); + // The current owner can hand it on. + assert!(s.reclaim_execution(job, "rescuer", "rescuer2").unwrap()); + // No claim row → no-op. + assert!(!s.reclaim_execution("no-such-claim", "x", "y").unwrap()); + s.complete_execution(job).unwrap(); + + // Owners may contain ':' (e.g. "host:pid"). The numeric timestamp suffix is + // split off from the LAST ':', so the full owner must match — a truncated + // prefix must not. + let colon_job = "reclaim-colon-job"; + assert!(s.claim_execution(colon_job, "host:42").unwrap()); + assert!( + !s.reclaim_execution(colon_job, "host", "x").unwrap(), + "a truncated owner prefix must not match" + ); + assert!(s + .reclaim_execution(colon_job, "host:42", "rescuer") + .unwrap()); + s.complete_execution(colon_job).unwrap(); +} + +fn test_reap_orphaned_jobs(s: &impl Storage) { + // A running job whose claim owner is not in the live set is orphaned and + // paired with that dead owner; a live owner or an empty set yields nothing. + let q = "q-orphan-recovery"; + let job = s.enqueue(make_job(q, "orphan_task")).unwrap(); + s.dequeue(q, now_millis() + 1000, None).unwrap().unwrap(); + assert!(s.claim_execution(&job.id, "dead-worker").unwrap()); + + let orphans = s + .reap_orphaned_jobs(&["other".to_string()], now_millis()) + .unwrap(); + assert!( + orphans + .iter() + .any(|(j, owner)| j.id == job.id && owner == "dead-worker"), + "claim owned by a non-live worker must be reported as orphaned" + ); + + let live = s + .reap_orphaned_jobs(&["dead-worker".to_string()], now_millis()) + .unwrap(); + assert!( + !live.iter().any(|(j, _)| j.id == job.id), + "a live owner's job must not be orphaned" + ); + + // Empty live set is a defensive no-op (never sweeps). + assert!(s.reap_orphaned_jobs(&[], now_millis()).unwrap().is_empty()); + + // Once the job leaves Running it is no longer orphaned. + s.complete(&job.id, None).unwrap(); + let after = s + .reap_orphaned_jobs(&["other".to_string()], now_millis()) + .unwrap(); + assert!(!after.iter().any(|(j, _)| j.id == job.id)); + s.complete_execution(&job.id).unwrap(); + + // Owners containing ':' must be parsed whole (split on the LAST ':'), so a + // truncated prefix is neither reported as the owner nor matched as live. + let cq = "q-orphan-colon"; + let cjob = s.enqueue(make_job(cq, "orphan_colon_task")).unwrap(); + s.dequeue(cq, now_millis() + 1000, None).unwrap().unwrap(); + assert!(s.claim_execution(&cjob.id, "host:7").unwrap()); + let co = s + .reap_orphaned_jobs(&["other".to_string()], now_millis()) + .unwrap(); + assert!( + co.iter() + .any(|(j, owner)| j.id == cjob.id && owner == "host:7"), + "the full colon-containing owner must be reported" + ); + let cl = s + .reap_orphaned_jobs(&["host:7".to_string()], now_millis()) + .unwrap(); + assert!( + !cl.iter().any(|(j, _)| j.id == cjob.id), + "the full colon-containing owner being live means not orphaned" + ); + s.complete(&cjob.id, None).unwrap(); + s.complete_execution(&cjob.id).unwrap(); +} + fn test_dashboard_settings(s: &impl Storage) { // get on missing key assert!(s.get_setting("settings-nonexistent").unwrap().is_none()); @@ -765,6 +874,9 @@ fn run_storage_tests(s: &impl Storage) { test_periodic_crud(s); test_circuit_breakers(s); test_execution_claims_purge(s); + test_reap_stale_jobs(s); + test_reclaim_execution(s); + test_reap_orphaned_jobs(s); test_dashboard_settings(s); test_immediate_archival(s); test_enqueue_dep_on_completed_archived_job(s); diff --git a/crates/taskito-java/src/worker.rs b/crates/taskito-java/src/worker.rs index 11548b11..14008aa9 100644 --- a/crates/taskito-java/src/worker.rs +++ b/crates/taskito-java/src/worker.rs @@ -72,6 +72,9 @@ fn start_worker( let worker_id = format!("java-{}", uuid::Uuid::now_v7()); let mut scheduler = Scheduler::new(storage, queues, config, namespace); + // Claim execution under this worker's id so dead-worker recovery can + // attribute orphaned jobs (matches the register_worker id below). + scheduler.set_claim_owner(worker_id.clone()); register_task_policies(&mut scheduler, options.task_configs); let scheduler = Arc::new(scheduler); let shutdown = scheduler.shutdown_handle(); diff --git a/crates/taskito-node/src/worker.rs b/crates/taskito-node/src/worker.rs index 8044a7a4..c9c68175 100644 --- a/crates/taskito-node/src/worker.rs +++ b/crates/taskito-node/src/worker.rs @@ -97,6 +97,9 @@ pub fn start_worker( // Per-task/queue config must be registered before the scheduler is shared // (register_* take &mut self). let mut scheduler = Scheduler::new(storage, queues, config, namespace); + // Claim execution under this worker's id so dead-worker recovery can + // attribute orphaned jobs (matches the register_worker id below). + scheduler.set_claim_owner(worker_id.clone()); for input in options.task_configs.iter().flatten() { scheduler.register_task(input.name.clone(), crate::convert::task_config(input)?); } diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 35edfc14..0bb9b9f7 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -287,12 +287,17 @@ impl PyQueue { dlq_auto_retry_max: self.dlq_auto_retry_max, ..SchedulerConfig::default() }; + // Resolve the worker id up front so the scheduler claims execution under + // it — dead-worker recovery attributes orphaned claims by this id, which + // must match the `register_worker`/`heartbeat` id below. + let worker_id = worker_id.unwrap_or_else(|| uuid::Uuid::now_v7().to_string()); let mut scheduler = Scheduler::new( self.storage.clone(), queues, scheduler_config, self.namespace.clone(), ); + scheduler.set_claim_owner(worker_id.clone()); // Build retry filters dict from the Queue's _task_retry_filters let retry_filters = PyDict::new(py).into_any(); @@ -425,8 +430,8 @@ impl PyQueue { let scheduler_arc = Arc::new(scheduler); let scheduler_for_dispatch = scheduler_arc.clone(); - // Generate or use the provided worker ID and register - let worker_id = worker_id.unwrap_or_else(|| uuid::Uuid::now_v7().to_string()); + // Register the worker under the id resolved above (shared with the + // scheduler's claim owner). let hostname = gethostname::gethostname().to_string_lossy().to_string(); let pid = std::process::id() as i32; let _ = self.storage.register_worker( diff --git a/docs/content/docs/resources/changelog.mdx b/docs/content/docs/resources/changelog.mdx index 3e2010ee..988fcb25 100644 --- a/docs/content/docs/resources/changelog.mdx +++ b/docs/content/docs/resources/changelog.mdx @@ -22,6 +22,12 @@ underlying Rust crates are released together, in lock-step. ### Added +- **Fast in-flight recovery on worker death.** When a worker process crashes mid-job, + a surviving worker now requeues its `Running` jobs within ~30s of the missed + heartbeat (via retry / dead-letter) instead of waiting the job's full timeout. The + scheduler claims execution under its real `worker_id` and a new maintenance step + reclaims orphaned claims atomically, so concurrent survivors never double-rescue. + Works across Python, Node, and Java; SQLite/Postgres/Redis. No API change. - **Indexes for dead-letter and filter paths.** `dead_letter` (previously unindexed) gains indexes on `failed_at` and `task_name`; `jobs` gains `(task_name, status)` and partial indexes on `expires_at` / `namespace`; `archived_jobs` gains