feat(scheduler): recover in-flight jobs on worker death

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

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<String> = 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(())
     }
 

crates/taskito-core/src/scheduler/mod.rs

@@ -149,6 +149,10 @@ pub struct Scheduler {
     shutdown: Arc<Notify>,
     paused_cache: Mutex<(HashSet<String>, Instant)>,
     namespace: Option<String>,
+    /// 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<Notify> {
         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();

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<Job>) -> Result<bool> {
@@ -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<bool> {
-        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) => {

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<Vec<(Job, String)>> {
+                // 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<String> = orphan_claims.iter().map(|(id, _)| id.clone()).collect();
+                let owner_by_job: std::collections::HashMap<String, String> =
+                    orphan_claims.into_iter().collect();
+
+                // Of those, the ones still Running (blob-free narrow row).
+                let rows: Vec<NarrowJobRow> = 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()?;

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<bool> {
+                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<u64> {
                 let mut conn = self.conn()?;