From 3df0617adf19d2730b488f584a08fc17a7eaa4b9 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:47:46 +0530 Subject: [PATCH 1/6] feat(storage): add a best-effort leader-election helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit try_lead renews an owner's live lock before trying to take a free one — an acquire-only caller would thrash, since acquire_lock refuses a live lock even for its owner. Load shedding, not mutual exclusion: the callers are idempotent, so a lost lock only means more than one worker sweeps. Adds the reaper/retention lock names and TTLs, and dead_worker_cutoff to share the staleness arithmetic. --- crates/taskito-core/src/storage/mod.rs | 98 ++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index 4b984202..89dde118 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -36,6 +36,54 @@ pub const DEAD_WORKER_THRESHOLD_MS: i64 = 30_000; /// threshold keeps one full failure-detection cycle of headroom. pub const EPHEMERAL_SUBSCRIPTION_GRACE_MS: i64 = 2 * DEAD_WORKER_THRESHOLD_MS; +/// The `last_heartbeat` at or below which a worker is dead. One definition so +/// the three backends' reaps and the orphan-recovery filter can never drift on +/// the arithmetic (they already share [`DEAD_WORKER_THRESHOLD_MS`]). +pub fn dead_worker_cutoff(now: i64) -> i64 { + now.saturating_sub(DEAD_WORKER_THRESHOLD_MS) +} + +/// Cluster-wide election lock for the dead-worker reap. Not namespaced — the +/// `workers` table is cluster-global, so one worker per cluster should reap. +pub const REAPER_LOCK: &str = "taskito:reaper"; + +/// Reaper lease. Longer than the 5s heartbeat cadence so the leader keeps the +/// lock across ticks, but under [`DEAD_WORKER_THRESHOLD_MS`] so a dead leader's +/// lock frees before the workers it should have reaped go stale-plus-a-cycle. +pub const REAPER_LOCK_TTL_MS: i64 = 15_000; + +/// Cluster-wide election lock for the retention purge. Separate from the reaper +/// lock so a long cleanup sweep can never starve worker-death detection. +pub const RETENTION_LOCK: &str = "taskito:retention"; + +/// Retention lease. Must exceed the worst-case sweep, which the per-tick batch +/// cap bounds — the first sweep after a retention default flip is the long one. +pub const RETENTION_LOCK_TTL_MS: i64 = 300_000; + +/// Hold `lock_name` as `owner_id` for another `ttl_ms`, renewing first and only +/// taking it when free. Returns whether this caller now holds it. +/// +/// The renew-then-acquire order is required, not stylistic: `acquire_lock` +/// refuses a live lock even for its current owner (it only checks +/// `expires_at > now`), so an acquire-first caller would lose leadership every +/// time its own lock is still valid and thrash on each expiry. +/// +/// This is best-effort load shedding, not mutual exclusion — the callers +/// (dead-worker reap, retention purge) are idempotent predicate deletes that +/// re-evaluate fresh state, so a lost lock degrades to the pre-election +/// behaviour of every worker doing the sweep, never to incorrectness. +pub fn try_lead( + storage: &impl Storage, + lock_name: &str, + owner_id: &str, + ttl_ms: i64, +) -> Result { + if storage.extend_lock(lock_name, owner_id, ttl_ms)? { + return Ok(true); + } + storage.acquire_lock(lock_name, owner_id, ttl_ms) +} + // ── Shared helper types ──────────────────────────────────────────────── #[derive(Debug, Clone, Default)] @@ -1460,3 +1508,53 @@ impl Storage for StorageBackend { delegate!(self, list_settings) } } + +#[cfg(test)] +mod election_tests { + use super::*; + use crate::storage::sqlite::SqliteStorage; + + fn store() -> SqliteStorage { + SqliteStorage::in_memory().unwrap() + } + + #[test] + fn try_lead_takes_a_free_lock() { + let s = store(); + assert!(try_lead(&s, "l", "worker-a", 10_000).unwrap()); + } + + #[test] + fn try_lead_renews_its_own_live_lock() { + // The load-bearing test: `acquire_lock` refuses a live lock even for its + // owner, so an acquire-first `try_lead` would return false on the second + // call and hand leadership to nobody. Renew-then-acquire keeps it. + let s = store(); + assert!(try_lead(&s, "l", "worker-a", 10_000).unwrap()); + assert!( + try_lead(&s, "l", "worker-a", 10_000).unwrap(), + "the current leader must keep the lock across ticks" + ); + } + + #[test] + fn try_lead_refuses_a_lock_held_by_another() { + let s = store(); + assert!(try_lead(&s, "l", "worker-a", 10_000).unwrap()); + assert!( + !try_lead(&s, "l", "worker-b", 10_000).unwrap(), + "a live lock held by another owner is not stealable" + ); + } + + #[test] + fn try_lead_takes_an_expired_lock() { + let s = store(); + assert!(try_lead(&s, "l", "worker-a", 1).unwrap()); + std::thread::sleep(std::time::Duration::from_millis(5)); + assert!( + try_lead(&s, "l", "worker-b", 10_000).unwrap(), + "an expired lock is stealable by a new owner" + ); + } +} From 17fd1a26ef97770fcdcd67f2b4299f0bef31c186 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:53:48 +0530 Subject: [PATCH 2/6] perf(storage): list live worker ids without loading full rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ephemeral-subscription reap and orphan recovery only need live worker ids, but both loaded every worker row — including the resource_health JSON blob — to extract them. Add a narrow list_live_worker_ids (pipelined on Redis), and share the staleness cutoff so the three backends and the recovery filter can't drift. --- .../taskito-core/src/scheduler/maintenance.rs | 10 +--- .../src/storage/diesel_common/workers.rs | 18 ++++++- crates/taskito-core/src/storage/mod.rs | 9 ++++ .../src/storage/postgres/workers.rs | 1 - .../src/storage/redis_backend/workers.rs | 53 ++++++++++++++++--- .../taskito-core/src/storage/sqlite/tests.rs | 28 ++++++++++ .../src/storage/sqlite/workers.rs | 1 - crates/taskito-core/src/storage/traits.rs | 4 ++ .../taskito-core/tests/rust/storage_tests.rs | 8 +++ crates/taskito-java/src/backend.rs | 7 +-- crates/taskito-node/src/queue/pubsub.rs | 8 +-- crates/taskito-python/src/py_queue/pubsub.rs | 7 +-- 12 files changed, 121 insertions(+), 33 deletions(-) diff --git a/crates/taskito-core/src/scheduler/maintenance.rs b/crates/taskito-core/src/scheduler/maintenance.rs index 2d5b42d3..32aaaca5 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, DEAD_WORKER_THRESHOLD_MS}; +use crate::storage::{dead_worker_cutoff, Storage}; use super::{JobResult, Scheduler}; @@ -76,13 +76,7 @@ impl Scheduler { // 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(); + let mut live = self.storage.list_live_worker_ids(dead_worker_cutoff(now))?; live.push(self.claim_owner.clone()); for (job, dead_owner) in self.storage.reap_orphaned_jobs(&live, now)? { diff --git a/crates/taskito-core/src/storage/diesel_common/workers.rs b/crates/taskito-core/src/storage/diesel_common/workers.rs index 4bbbc790..4de4f497 100644 --- a/crates/taskito-core/src/storage/diesel_common/workers.rs +++ b/crates/taskito-core/src/storage/diesel_common/workers.rs @@ -44,6 +44,22 @@ macro_rules! impl_diesel_worker_ops { Ok(rows) } + /// Ids of workers whose heartbeat is at or after `cutoff_ms`. + /// + /// Pushes both the liveness filter and the projection into the query + /// — a caller that only needs live ids must not load every worker's + /// `resource_health` blob just to discard it. + pub fn list_live_worker_ids(&self, cutoff_ms: i64) -> Result> { + let mut conn = self.conn()?; + + let ids = workers::table + .filter(workers::last_heartbeat.ge(cutoff_ms)) + .select(workers::worker_id) + .load(&mut conn)?; + + Ok(ids) + } + /// Remove workers that haven't sent a heartbeat within the threshold. /// Returns the IDs of the reaped workers. /// @@ -55,7 +71,7 @@ macro_rules! impl_diesel_worker_ops { /// orphan rescue, both of which tolerate the false positive. pub fn reap_dead_workers(&self) -> Result> { let mut conn = self.conn()?; - let cutoff = now_millis().saturating_sub(DEAD_WORKER_THRESHOLD_MS); + let cutoff = $crate::storage::dead_worker_cutoff(now_millis()); let dead_ids: Vec = workers::table .filter(workers::last_heartbeat.lt(cutoff)) diff --git a/crates/taskito-core/src/storage/mod.rs b/crates/taskito-core/src/storage/mod.rs index 89dde118..eb7b1df2 100644 --- a/crates/taskito-core/src/storage/mod.rs +++ b/crates/taskito-core/src/storage/mod.rs @@ -713,6 +713,12 @@ macro_rules! impl_storage { ) -> $crate::error::Result> { self.list_workers() } + fn list_live_worker_ids( + &self, + cutoff_ms: i64, + ) -> $crate::error::Result> { + self.list_live_worker_ids(cutoff_ms) + } fn reap_dead_workers(&self) -> $crate::error::Result> { self.reap_dead_workers() } @@ -1358,6 +1364,9 @@ impl Storage for StorageBackend { fn list_workers(&self) -> Result> { delegate!(self, list_workers) } + fn list_live_worker_ids(&self, cutoff_ms: i64) -> Result> { + delegate!(self, list_live_worker_ids, cutoff_ms) + } fn reap_dead_workers(&self) -> Result> { delegate!(self, reap_dead_workers) } diff --git a/crates/taskito-core/src/storage/postgres/workers.rs b/crates/taskito-core/src/storage/postgres/workers.rs index 9f9d4cae..5c38615e 100644 --- a/crates/taskito-core/src/storage/postgres/workers.rs +++ b/crates/taskito-core/src/storage/postgres/workers.rs @@ -2,7 +2,6 @@ use diesel::prelude::*; use super::super::models::*; use super::super::schema::{execution_claims, workers}; -use super::super::DEAD_WORKER_THRESHOLD_MS; use super::PostgresStorage; use crate::error::Result; use crate::job::now_millis; diff --git a/crates/taskito-core/src/storage/redis_backend/workers.rs b/crates/taskito-core/src/storage/redis_backend/workers.rs index 84910a73..9bbae675 100644 --- a/crates/taskito-core/src/storage/redis_backend/workers.rs +++ b/crates/taskito-core/src/storage/redis_backend/workers.rs @@ -4,7 +4,6 @@ use super::{map_err, RedisStorage}; use crate::error::Result; use crate::job::now_millis; use crate::storage::models::WorkerRow; -use crate::storage::DEAD_WORKER_THRESHOLD_MS; impl RedisStorage { #[allow(clippy::too_many_arguments)] @@ -119,24 +118,52 @@ impl RedisStorage { Ok(rows) } + /// Ids of workers whose heartbeat is at or after `cutoff_ms`. + /// + /// Pipelines the per-worker `HGET` so a live-set read costs one round trip + /// plus one pipeline flush, not one round trip per worker. + pub fn list_live_worker_ids(&self, cutoff_ms: i64) -> Result> { + let mut conn = self.conn()?; + let wall = self.key(&["workers", "all"]); + + let worker_ids: Vec = conn.smembers(&wall).map_err(map_err)?; + if worker_ids.is_empty() { + return Ok(Vec::new()); + } + + let heartbeats: Vec> = self.heartbeats_pipelined(&mut conn, &worker_ids)?; + + Ok(worker_ids + .into_iter() + .zip(heartbeats) + .filter_map(|(wid, hb)| match hb { + Some(last_hb) if last_hb >= cutoff_ms => Some(wid), + _ => None, + }) + .collect()) + } + pub fn reap_dead_workers(&self) -> Result> { let mut conn = self.conn()?; - let cutoff = now_millis().saturating_sub(DEAD_WORKER_THRESHOLD_MS); + let cutoff = crate::storage::dead_worker_cutoff(now_millis()); let wall = self.key(&["workers", "all"]); let worker_ids: Vec = conn.smembers(&wall).map_err(map_err)?; + if worker_ids.is_empty() { + return Ok(Vec::new()); + } - let mut reaped = Vec::new(); - for wid in worker_ids { - let wkey = self.key(&["worker", &wid]); - let hb: Option = conn.hget(&wkey, "last_heartbeat").map_err(map_err)?; + let heartbeats = self.heartbeats_pipelined(&mut conn, &worker_ids)?; + let mut reaped = Vec::new(); + for (wid, hb) in worker_ids.into_iter().zip(heartbeats) { let is_dead = match hb { Some(last_hb) => last_hb < cutoff, None => true, }; if is_dead { + let wkey = self.key(&["worker", &wid]); let pipe = &mut redis::pipe(); pipe.del(&wkey); pipe.srem(&wall, &wid); @@ -148,6 +175,20 @@ impl RedisStorage { Ok(reaped) } + /// Fetch every worker's `last_heartbeat` in one pipeline, preserving input + /// order so callers can zip the result back onto the id list. + fn heartbeats_pipelined( + &self, + conn: &mut redis::Connection, + worker_ids: &[String], + ) -> Result>> { + let mut pipe = redis::pipe(); + for wid in worker_ids { + pipe.hget(self.key(&["worker", wid]), "last_heartbeat"); + } + pipe.query(conn).map_err(map_err) + } + pub fn unregister_worker(&self, worker_id: &str) -> Result<()> { let mut conn = self.conn()?; let wkey = self.key(&["worker", worker_id]); diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index 3af50d3c..549c532d 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -457,6 +457,34 @@ fn test_reap_dead_workers_removes_stale_keeps_fresh() { assert!(!surviving.contains(&"stale".to_string())); } +#[test] +fn test_list_live_worker_ids_filters_stale() { + use diesel::prelude::*; + + use crate::storage::schema::workers; + + let storage = test_storage(); + storage + .register_worker("stale", "default", None, None, None, 1, None, None, None) + .unwrap(); + storage + .register_worker("fresh", "default", None, None, None, 1, None, None, None) + .unwrap(); + + let now = now_millis(); + let mut conn = storage.conn().unwrap(); + diesel::update(workers::table.filter(workers::worker_id.eq("stale"))) + .set(workers::last_heartbeat.eq(now - crate::storage::DEAD_WORKER_THRESHOLD_MS - 1_000)) + .execute(&mut conn) + .unwrap(); + drop(conn); + + let live = storage + .list_live_worker_ids(crate::storage::dead_worker_cutoff(now)) + .unwrap(); + assert_eq!(live, vec!["fresh".to_string()]); +} + #[test] fn test_stats() { let storage = test_storage(); diff --git a/crates/taskito-core/src/storage/sqlite/workers.rs b/crates/taskito-core/src/storage/sqlite/workers.rs index 8f581cf2..6951cbee 100644 --- a/crates/taskito-core/src/storage/sqlite/workers.rs +++ b/crates/taskito-core/src/storage/sqlite/workers.rs @@ -2,7 +2,6 @@ use diesel::prelude::*; use super::super::models::*; use super::super::schema::{execution_claims, workers}; -use super::super::DEAD_WORKER_THRESHOLD_MS; use super::SqliteStorage; use crate::error::Result; use crate::job::now_millis; diff --git a/crates/taskito-core/src/storage/traits.rs b/crates/taskito-core/src/storage/traits.rs index af8877d5..78596b61 100644 --- a/crates/taskito-core/src/storage/traits.rs +++ b/crates/taskito-core/src/storage/traits.rs @@ -261,6 +261,10 @@ pub trait Storage: Send + Sync + Clone { fn heartbeat(&self, worker_id: &str, resource_health: Option<&str>) -> Result<()>; fn update_worker_status(&self, worker_id: &str, status: &str) -> Result<()>; fn list_workers(&self) -> Result>; + /// Ids of workers whose heartbeat is at or after `cutoff_ms`. A narrow + /// projection of [`Self::list_workers`] for callers that only need the live + /// set and must not pay to load every worker's `resource_health` blob. + fn list_live_worker_ids(&self, cutoff_ms: i64) -> Result>; fn reap_dead_workers(&self) -> Result>; fn unregister_worker(&self, worker_id: &str) -> Result<()>; fn list_claims_by_worker(&self, worker_id: &str) -> Result>; diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index 4865e6e9..b68ecf38 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -400,6 +400,14 @@ fn test_workers(s: &impl Storage) { let w = workers.iter().find(|w| w.worker_id == "w-test-1").unwrap(); assert_eq!(w.status, "draining"); + // list_live_worker_ids applies the cutoff without loading the row: a fresh + // worker is live under a past cutoff and excluded under a future one. + let now = taskito_core::job::now_millis(); + let live = s.list_live_worker_ids(now - 10_000).unwrap(); + assert!(live.contains(&"w-test-1".to_string())); + let none_live = s.list_live_worker_ids(now + 10_000).unwrap(); + assert!(!none_live.contains(&"w-test-1".to_string())); + s.unregister_worker("w-test-1").unwrap(); } diff --git a/crates/taskito-java/src/backend.rs b/crates/taskito-java/src/backend.rs index f7b96d07..c0edd4ef 100644 --- a/crates/taskito-java/src/backend.rs +++ b/crates/taskito-java/src/backend.rs @@ -112,11 +112,8 @@ pub fn enqueue_batch_dedup( /// Callers prune dead workers first — a stale registry row keeps its /// subscriptions alive. Returns the count removed. pub fn reap_ephemeral_subscriptions(storage: &StorageBackend) -> Result { - let live: Vec = storage - .list_workers()? - .into_iter() - .map(|worker| worker.worker_id) - .collect(); + let cutoff = taskito_core::storage::dead_worker_cutoff(taskito_core::job::now_millis()); + let live = storage.list_live_worker_ids(cutoff)?; Ok(storage.reap_ephemeral_subscriptions(&live)?) } diff --git a/crates/taskito-node/src/queue/pubsub.rs b/crates/taskito-node/src/queue/pubsub.rs index c887202a..02be8fe7 100644 --- a/crates/taskito-node/src/queue/pubsub.rs +++ b/crates/taskito-node/src/queue/pubsub.rs @@ -122,12 +122,8 @@ impl JsQueue { // Prune stale worker rows first so a crashed owner doesn't keep its // ephemeral subscriptions alive; opportunistic like the heartbeat's. let _ = storage.reap_dead_workers(); - let live: Vec = storage - .list_workers() - .map_err(to_napi_err)? - .into_iter() - .map(|worker| worker.worker_id) - .collect(); + let cutoff = taskito_core::storage::dead_worker_cutoff(now_millis()); + let live = storage.list_live_worker_ids(cutoff).map_err(to_napi_err)?; storage .reap_ephemeral_subscriptions(&live) .map(|n| n as i64) diff --git a/crates/taskito-python/src/py_queue/pubsub.rs b/crates/taskito-python/src/py_queue/pubsub.rs index 1cfa0da4..bcb4fa59 100644 --- a/crates/taskito-python/src/py_queue/pubsub.rs +++ b/crates/taskito-python/src/py_queue/pubsub.rs @@ -149,11 +149,8 @@ impl PyQueue { // Release the GIL: the reap scans every worker + subscription, which // must not freeze other Python threads while it runs. py.detach(|| { - let live: Vec = storage - .list_workers()? - .into_iter() - .map(|w| w.worker_id) - .collect(); + let cutoff = taskito_core::storage::dead_worker_cutoff(now_millis()); + let live = storage.list_live_worker_ids(cutoff)?; storage.reap_ephemeral_subscriptions(&live) }) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) From 95e4a59fd3aed97afd2f69c62109f4a562347d55 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:09:31 +0530 Subject: [PATCH 3/6] perf(scaling): elect a single reaper for the dead-worker sweep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every worker scanned the whole registry for dead workers every 5s — O(N) per cluster — and each returned the same dead ids, so a WORKER_OFFLINE webhook fired once per live worker per death. The heartbeat now takes a cluster lease first and only the holder reaps; the ephemeral-subscription sweep rides the same lease. A non-leader reaps nothing and emits nothing. --- crates/taskito-java/src/worker.rs | 27 +++++++--- crates/taskito-node/src/queue/inspect.rs | 15 ++++++ crates/taskito-node/src/queue/pubsub.rs | 28 +++++++--- crates/taskito-python/src/py_queue/pubsub.rs | 24 ++++++++- crates/taskito-python/src/py_queue/worker.rs | 17 ++++++ sdks/node/src/worker.ts | 8 +-- sdks/python/taskito/_taskito.pyi | 2 +- sdks/python/taskito/mixins/lifecycle.py | 5 +- .../tests/worker/test_heartbeat_election.py | 52 +++++++++++++++++++ 9 files changed, 155 insertions(+), 23 deletions(-) create mode 100644 sdks/python/tests/worker/test_heartbeat_election.py diff --git a/crates/taskito-java/src/worker.rs b/crates/taskito-java/src/worker.rs index 9c36045b..40c0a0ed 100644 --- a/crates/taskito-java/src/worker.rs +++ b/crates/taskito-java/src/worker.rs @@ -476,14 +476,25 @@ fn spawn_lifecycle( if let Err(e) = storage.heartbeat(&worker_id, None) { log::warn!("[taskito-java] worker heartbeat failed: {e}"); } - // Prune stale workers first: the ephemeral-subscription reap - // treats every registry row as live, so a lingering dead - // worker would keep its subscriptions receiving deliveries. - if let Err(e) = storage.reap_dead_workers() { - log::warn!("[taskito-java] dead-worker reap failed: {e}"); - } - if let Err(e) = crate::backend::reap_ephemeral_subscriptions(&storage) { - log::warn!("[taskito-java] ephemeral subscription reap failed: {e}"); + // Elect a single reaper: without this, every worker's 5s + // sweep scans the whole registry, O(N) per cluster. The + // non-leader skips both reaps. Dead-worker reap first so a + // WORKER_OFFLINE peer sees the pruned registry; the ephemeral + // reap already filters the live set by heartbeat. + let leading = taskito_core::storage::try_lead( + &storage, + taskito_core::storage::REAPER_LOCK, + &worker_id, + taskito_core::storage::REAPER_LOCK_TTL_MS, + ) + .unwrap_or(false); + if leading { + if let Err(e) = storage.reap_dead_workers() { + log::warn!("[taskito-java] dead-worker reap failed: {e}"); + } + if let Err(e) = crate::backend::reap_ephemeral_subscriptions(&storage) { + log::warn!("[taskito-java] ephemeral subscription reap failed: {e}"); + } } } } diff --git a/crates/taskito-node/src/queue/inspect.rs b/crates/taskito-node/src/queue/inspect.rs index c571bbc4..09b47d1f 100644 --- a/crates/taskito-node/src/queue/inspect.rs +++ b/crates/taskito-node/src/queue/inspect.rs @@ -410,6 +410,11 @@ impl JsQueue { /// Record a heartbeat for a running worker, carrying its current resource /// health as a JSON object (`null` clears it). Called from the JS shell /// every 5s. Returns the ids of dead workers reaped as a side effect. + /// + /// Only the elected reaper sweeps: every worker scanning for dead workers + /// every 5s is O(N) per cluster, and each returns the same dead ids so a + /// `WORKER_OFFLINE` event fires N times per death. A non-leader reaps + /// nothing and returns an empty list. #[napi] pub async fn worker_heartbeat( &self, @@ -422,6 +427,16 @@ impl JsQueue { .heartbeat(&worker_id, resource_health.as_deref()) .map_err(to_napi_err)?; // Reaping is opportunistic — a failure must not fail the heartbeat. + let leading = taskito_core::storage::try_lead( + &storage, + taskito_core::storage::REAPER_LOCK, + &worker_id, + taskito_core::storage::REAPER_LOCK_TTL_MS, + ) + .unwrap_or(false); + if !leading { + return Ok(Vec::new()); + } Ok(storage.reap_dead_workers().unwrap_or_default()) }) .await diff --git a/crates/taskito-node/src/queue/pubsub.rs b/crates/taskito-node/src/queue/pubsub.rs index 02be8fe7..d38031ca 100644 --- a/crates/taskito-node/src/queue/pubsub.rs +++ b/crates/taskito-node/src/queue/pubsub.rs @@ -112,16 +112,30 @@ impl JsQueue { .map_err(join_to_napi_err)? } - /// Drop ephemeral subscriptions whose owning worker is gone. Prunes dead - /// worker rows first so the live set is current. Runs on the heartbeat - /// cadence. Returns the number of subscriptions removed. + /// Drop ephemeral subscriptions whose owning worker is gone. Runs on the + /// heartbeat cadence. Returns the number of subscriptions removed. + /// + /// Gated behind the same reaper election as the dead-worker reap when a + /// `workerId` is given: only the leader sweeps, so the scan runs once per + /// cluster. Without one (a manual call) it runs unconditionally. The live + /// set is filtered by heartbeat, so a stale worker is excluded whether or + /// not its registry row has been reaped yet. #[napi] - pub async fn reap_ephemeral_subscriptions(&self) -> Result { + pub async fn reap_ephemeral_subscriptions(&self, worker_id: Option) -> Result { let storage = self.storage.clone(); spawn_blocking(move || { - // Prune stale worker rows first so a crashed owner doesn't keep its - // ephemeral subscriptions alive; opportunistic like the heartbeat's. - let _ = storage.reap_dead_workers(); + if let Some(owner) = worker_id { + let leading = taskito_core::storage::try_lead( + &storage, + taskito_core::storage::REAPER_LOCK, + &owner, + taskito_core::storage::REAPER_LOCK_TTL_MS, + ) + .map_err(to_napi_err)?; + if !leading { + return Ok(0); + } + } let cutoff = taskito_core::storage::dead_worker_cutoff(now_millis()); let live = storage.list_live_worker_ids(cutoff).map_err(to_napi_err)?; storage diff --git a/crates/taskito-python/src/py_queue/pubsub.rs b/crates/taskito-python/src/py_queue/pubsub.rs index bcb4fa59..f0b5516d 100644 --- a/crates/taskito-python/src/py_queue/pubsub.rs +++ b/crates/taskito-python/src/py_queue/pubsub.rs @@ -144,11 +144,33 @@ impl PyQueue { /// Drop ephemeral subscriptions whose owning worker is gone. Runs on the /// heartbeat cadence, after `reap_dead_workers` has pruned the registry. - pub fn reap_ephemeral_subscriptions(&self, py: Python<'_>) -> PyResult { + /// + /// Gated behind the same reaper election as the dead-worker reap: only the + /// leader sweeps, so the scan runs once per cluster rather than once per + /// worker. A non-leader returns 0. `worker_id` is the election owner. + #[pyo3(signature = (worker_id=None))] + pub fn reap_ephemeral_subscriptions( + &self, + py: Python<'_>, + worker_id: Option<&str>, + ) -> PyResult { let storage = &self.storage; // Release the GIL: the reap scans every worker + subscription, which // must not freeze other Python threads while it runs. py.detach(|| { + // An explicit owner elects a single reaper; without one (a manual + // admin call) the sweep runs unconditionally. + if let Some(owner) = worker_id { + let leading = taskito_core::storage::try_lead( + storage, + taskito_core::storage::REAPER_LOCK, + owner, + taskito_core::storage::REAPER_LOCK_TTL_MS, + )?; + if !leading { + return Ok(0); + } + } let cutoff = taskito_core::storage::dead_worker_cutoff(now_millis()); let live = storage.list_live_worker_ids(cutoff)?; storage.reap_ephemeral_subscriptions(&live) diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 58f88a01..f29da7e8 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -824,6 +824,11 @@ impl PyQueue { /// Update the heartbeat for a running worker. Called from Python every 5s. /// Returns a list of worker IDs that were reaped as dead. + /// + /// Only the elected reaper sweeps: every worker running this every 5s makes + /// the dead-worker scan O(N) per cluster, and each returns the same dead ids + /// so a `WORKER_OFFLINE` webhook fires N times per death. A non-leader + /// returns an empty list and emits nothing. #[pyo3(signature = (worker_id, resource_health=None))] pub fn worker_heartbeat( &self, @@ -833,6 +838,18 @@ impl PyQueue { self.storage .heartbeat(worker_id, resource_health) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + + let leading = taskito_core::storage::try_lead( + &self.storage, + taskito_core::storage::REAPER_LOCK, + worker_id, + taskito_core::storage::REAPER_LOCK_TTL_MS, + ) + .unwrap_or(false); + if !leading { + return Ok(Vec::new()); + } + let reaped = self.storage.reap_dead_workers().unwrap_or_default(); Ok(reaped) } diff --git a/sdks/node/src/worker.ts b/sdks/node/src/worker.ts index 586882dd..d222aa0c 100644 --- a/sdks/node/src/worker.ts +++ b/sdks/node/src/worker.ts @@ -276,10 +276,10 @@ export class Worker { void queue.workerHeartbeat(native.id, snapshot && JSON.stringify(snapshot)).catch((error) => { log.debug(() => "worker heartbeat failed", error); }); - // Same cadence: prune ephemeral subscriptions whose owner is gone. - // Per-tick failures are swallowed like the heartbeat's — the next - // beat retries. - void queue.reapEphemeralSubscriptions().catch((error) => { + // Same cadence, same reaper election: passing this worker's id gates the + // sweep so only the leader runs it. Per-tick failures are swallowed like + // the heartbeat's — the next beat retries. + void queue.reapEphemeralSubscriptions(native.id).catch((error) => { log.debug(() => "ephemeral subscription reap failed", error); }); declareSubscriptions(); diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 0f3e2e95..0fc37d53 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -234,7 +234,7 @@ class PyQueue: def set_subscription_active( self, topic: str, subscription_name: str, active: bool ) -> bool: ... - def reap_ephemeral_subscriptions(self) -> int: ... + def reap_ephemeral_subscriptions(self, worker_id: str | None = None) -> int: ... def topic_backlog_stats( self, ) -> list[tuple[str, str, str, str, bool, bool, int, int, int, int | None]]: ... diff --git a/sdks/python/taskito/mixins/lifecycle.py b/sdks/python/taskito/mixins/lifecycle.py index 363ec9c7..fdd1b26e 100644 --- a/sdks/python/taskito/mixins/lifecycle.py +++ b/sdks/python/taskito/mixins/lifecycle.py @@ -336,8 +336,9 @@ def _run_heartbeat( for rid in reaped_ids: self._emit_event(EventType.WORKER_OFFLINE, {"worker_id": rid}) # type: ignore[attr-defined] # Dead workers gone from the registry → drop their ephemeral - # topic subscriptions on the same cadence. - self._inner.reap_ephemeral_subscriptions() + # topic subscriptions on the same cadence, under the same reaper + # election so only one worker sweeps. + self._inner.reap_ephemeral_subscriptions(worker_id) except Exception: logger.debug("Heartbeat failed", exc_info=True) diff --git a/sdks/python/tests/worker/test_heartbeat_election.py b/sdks/python/tests/worker/test_heartbeat_election.py new file mode 100644 index 00000000..90ca00f2 --- /dev/null +++ b/sdks/python/tests/worker/test_heartbeat_election.py @@ -0,0 +1,52 @@ +"""Reaper election: only one worker per cluster sweeps dead workers. + +Before the election, every worker reaped on every 5s tick, so a `WORKER_OFFLINE` +event fired once per live worker per death instead of once. The heartbeat now +takes a cluster-wide lease first and only the holder reaps; a non-holder returns +an empty reaped list and so emits nothing. + +These assert the mechanism through the lease itself rather than by simulating a +dead worker — a raw write to the queue's SQLite file is not reliably visible to +the queue's own connections (different journal mode), so the reap's *effect* +cannot be observed from a test, but the election that gates it can. +""" + +from __future__ import annotations + +from pathlib import Path + +from taskito import Queue + +# Mirror of `REAPER_LOCK` in crates/taskito-core/src/storage/mod.rs. +REAPER_LOCK = "taskito:reaper" + + +def test_reaper_lock_is_free_on_a_fresh_queue(tmp_path: Path) -> None: + # Guards the test below: the lease is observing the heartbeat taking it, not + # a lock that was already held. + queue = Queue(db_path=str(tmp_path / "fresh.db")) + assert queue._inner.acquire_lock(REAPER_LOCK, "someone", 15_000) is True + + +def test_heartbeat_takes_the_reaper_lease(tmp_path: Path) -> None: + db = str(tmp_path / "election.db") + q1 = Queue(db_path=db) + q2 = Queue(db_path=db) + + # A fresh heartbeat elects worker-a as the reaper for the lease window. + q1._inner.worker_heartbeat("worker-a") + + # worker-b cannot take the live lease, so its tick sweeps nothing — the + # dedup that keeps WORKER_OFFLINE from firing once per worker. + assert q2._inner.acquire_lock(REAPER_LOCK, "worker-b", 15_000) is False + + +def test_non_leader_heartbeat_reaps_nothing(tmp_path: Path) -> None: + db = str(tmp_path / "gated.db") + q = Queue(db_path=db) + + # A peer already holds the lease. + assert q._inner.acquire_lock(REAPER_LOCK, "peer", 15_000) is True + + # This worker is gated out, so it reaps nothing regardless of the registry. + assert q._inner.worker_heartbeat("worker-a") == [] From e08a30929618213a13f654b107377b3781e48c37 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:10:09 +0530 Subject: [PATCH 4/6] perf(scaling): elect a single retention cleaner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit auto_cleanup runs on every worker's maintenance tick, but the purge sweeps are cluster-wide — N workers draining the same tables is wasted work, and on a retention-default upgrade it is N simultaneous multi-million-row backfills. Gate it behind its own lease, separate from the reaper so a long sweep can't starve worker-death detection. --- .../taskito-core/src/scheduler/maintenance.rs | 20 +++++++++- crates/taskito-core/src/scheduler/mod.rs | 39 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/crates/taskito-core/src/scheduler/maintenance.rs b/crates/taskito-core/src/scheduler/maintenance.rs index 32aaaca5..a0543034 100644 --- a/crates/taskito-core/src/scheduler/maintenance.rs +++ b/crates/taskito-core/src/scheduler/maintenance.rs @@ -3,7 +3,9 @@ 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::{dead_worker_cutoff, Storage}; +use crate::storage::{ + dead_worker_cutoff, try_lead, Storage, RETENTION_LOCK, RETENTION_LOCK_TTL_MS, +}; use super::{JobResult, Scheduler}; @@ -132,6 +134,22 @@ impl Scheduler { /// deletes (and the side tables, which have no per-entry TTL) only run when a /// global `result_ttl` is set. pub(super) fn auto_cleanup(&self) -> Result<()> { + // Elect a single cleaner: every worker runs this tick, but the purge + // sweeps are cluster-wide, so N workers draining the same tables is + // wasted work — and on a retention-default upgrade, N simultaneous + // multi-million-row backfills. A separate lock from the reaper so a long + // sweep can't starve worker-death detection. + let leading = try_lead( + &self.storage, + RETENTION_LOCK, + &self.claim_owner, + RETENTION_LOCK_TTL_MS, + ) + .unwrap_or(false); + if !leading { + return Ok(()); + } + let now = now_millis(); let global_cutoff = self.global_retention_cutoff(now); // `i64::MIN` disables the methods' global-cutoff branch (`failed_at < MIN` diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 33dd2980..28f1c7e6 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -1966,6 +1966,45 @@ mod tests { assert!(fetched.is_none()); } + #[test] + fn test_auto_cleanup_skips_when_another_holds_the_lock() { + // A peer holds the retention lock, so this scheduler is not the elected + // cleaner and must purge nothing even though its own TTL has expired. + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); + let config = SchedulerConfig { + result_ttl_ms: Some(1), + ..SchedulerConfig::default() + }; + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config, None); + + let job = scheduler.storage.enqueue(make_job("gated")).unwrap(); + scheduler + .storage + .dequeue("default", now_millis() + 1000, None) + .unwrap(); + scheduler.storage.complete(&job.id, Some(vec![1])).unwrap(); + std::thread::sleep(Duration::from_millis(10)); + + // A different owner takes the lock first; the scheduler's own owner + // ("scheduler") then loses the election. + assert!(scheduler + .storage + .acquire_lock( + crate::storage::RETENTION_LOCK, + "another-worker", + crate::storage::RETENTION_LOCK_TTL_MS + ) + .unwrap()); + + scheduler.auto_cleanup().unwrap(); + + assert!( + scheduler.storage.get_job(&job.id).unwrap().is_some(), + "a non-leader must not purge" + ); + } + #[test] fn test_auto_cleanup_ignores_negative_result_ttl() { // A negative TTL inverts the cutoff — `now.saturating_sub(-ttl)` lands in From 2a1262ae8d694fcf831c10cf78828fff56305d07 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:37:49 +0530 Subject: [PATCH 5/6] fix(scaling): give each scheduler a unique default owner The default claim_owner was the shared constant "scheduler", so two un-configured schedulers on one store both renewed the same retention lease and both ran cleanup. Bindings always set a real worker_id, but the default must be safe on its own. --- crates/taskito-core/src/scheduler/mod.rs | 75 +++++++++++++++++++++--- 1 file changed, 68 insertions(+), 7 deletions(-) diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 28f1c7e6..e2318db4 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -278,7 +278,11 @@ impl Scheduler { shutdown: Arc::new(Notify::new()), paused_cache: Mutex::new((HashSet::new(), Instant::now())), namespace, - claim_owner: poller::SCHEDULER_CLAIM_OWNER.to_string(), + // Instance-unique by default so two un-configured schedulers sharing + // storage can't both win the retention election (`try_lead` renews a + // matching owner's lease). Bindings override this with the real + // `worker_id` via `set_claim_owner`. + claim_owner: format!("{}-{}", poller::SCHEDULER_CLAIM_OWNER, uuid::Uuid::now_v7()), in_flight: Mutex::new(InFlight::default()), dispatch_wake: Arc::new(Notify::new()), #[cfg(feature = "push-dispatch")] @@ -300,6 +304,11 @@ impl Scheduler { self.claim_owner = worker_id; } + /// The execution-claim / retention-election owner for this scheduler. + pub fn claim_owner(&self) -> &str { + &self.claim_owner + } + /// Free in-flight slots remaining before the dispatch cap, or `None` when /// no cap is configured (unbounded dispatch — legacy behavior). fn in_flight_remaining(&self) -> Option { @@ -802,7 +811,7 @@ mod tests { } let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert!(!claims.contains(&s0.id) && !claims.contains(&s1.id)); } @@ -1089,7 +1098,7 @@ mod tests { let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert_eq!( claims.len(), @@ -1150,7 +1159,7 @@ mod tests { let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert_eq!( claims.len(), @@ -1374,7 +1383,7 @@ mod tests { let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert_eq!( claims.len(), @@ -1576,7 +1585,7 @@ mod tests { let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert!( !claims.contains(&job.id), @@ -1610,7 +1619,7 @@ mod tests { let claims = scheduler .storage - .list_claims_by_worker("scheduler") + .list_claims_by_worker(scheduler.claim_owner()) .unwrap(); assert!(!claims.contains(&job.id)); @@ -2005,6 +2014,58 @@ mod tests { ); } + #[test] + fn test_two_default_schedulers_do_not_both_lead_retention() { + // Regression: the default claim_owner used to be the shared constant + // "scheduler", so two un-configured schedulers on one store both renewed + // the same lease and both ran retention. Unique default owners fix it. + // in_memory() pins a single-connection pool, so the clone shares the DB. + let storage = crate::storage::sqlite::SqliteStorage::in_memory().unwrap(); + let config = SchedulerConfig { + result_ttl_ms: Some(1), + ..SchedulerConfig::default() + }; + let s1 = Scheduler::new( + StorageBackend::Sqlite(storage.clone()), + vec!["default".to_string()], + config.clone(), + None, + ); + let s2 = Scheduler::new( + StorageBackend::Sqlite(storage), + vec!["default".to_string()], + config, + None, + ); + + assert_ne!( + s1.claim_owner(), + s2.claim_owner(), + "each scheduler must have a distinct default owner" + ); + + let lead1 = crate::storage::try_lead( + s1.storage(), + crate::storage::RETENTION_LOCK, + s1.claim_owner(), + crate::storage::RETENTION_LOCK_TTL_MS, + ) + .unwrap(); + let lead2 = crate::storage::try_lead( + s2.storage(), + crate::storage::RETENTION_LOCK, + s2.claim_owner(), + crate::storage::RETENTION_LOCK_TTL_MS, + ) + .unwrap(); + + assert!(lead1, "the first scheduler takes the free lease"); + assert!( + !lead2, + "the second must not also win — a shared owner would let it renew" + ); + } + #[test] fn test_auto_cleanup_ignores_negative_result_ttl() { // A negative TTL inverts the cutoff — `now.saturating_sub(-ttl)` lands in From 2975622c768370186af67c2531a73d110df3436c Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 22:38:04 +0530 Subject: [PATCH 6/6] fix(scaling): log election backend errors instead of swallowing try_lead(...).unwrap_or(false) treated a storage failure as lost leadership, so a backend outage could stall reaping or retention with no diagnostic. Log the error before skipping the tick; the reap/cleanup callers stay opportunistic. --- crates/taskito-core/src/scheduler/maintenance.rs | 7 ++++++- crates/taskito-java/src/worker.rs | 7 ++++++- crates/taskito-node/src/queue/inspect.rs | 7 ++++++- crates/taskito-python/src/py_queue/worker.rs | 7 ++++++- 4 files changed, 24 insertions(+), 4 deletions(-) diff --git a/crates/taskito-core/src/scheduler/maintenance.rs b/crates/taskito-core/src/scheduler/maintenance.rs index a0543034..70f93737 100644 --- a/crates/taskito-core/src/scheduler/maintenance.rs +++ b/crates/taskito-core/src/scheduler/maintenance.rs @@ -145,7 +145,12 @@ impl Scheduler { &self.claim_owner, RETENTION_LOCK_TTL_MS, ) - .unwrap_or(false); + .unwrap_or_else(|e| { + // A backend error is not the same as losing the election — surface it + // so a storage outage that stalls retention is diagnosable, not silent. + warn!("retention election failed: {e}"); + false + }); if !leading { return Ok(()); } diff --git a/crates/taskito-java/src/worker.rs b/crates/taskito-java/src/worker.rs index 40c0a0ed..a5f4f8c9 100644 --- a/crates/taskito-java/src/worker.rs +++ b/crates/taskito-java/src/worker.rs @@ -487,7 +487,12 @@ fn spawn_lifecycle( &worker_id, taskito_core::storage::REAPER_LOCK_TTL_MS, ) - .unwrap_or(false); + .unwrap_or_else(|e| { + // A backend error is not lost leadership — log it so a + // storage outage that stalls reaping is diagnosable. + log::warn!("[taskito-java] reaper election failed: {e}"); + false + }); if leading { if let Err(e) = storage.reap_dead_workers() { log::warn!("[taskito-java] dead-worker reap failed: {e}"); diff --git a/crates/taskito-node/src/queue/inspect.rs b/crates/taskito-node/src/queue/inspect.rs index 09b47d1f..d1ce3ac4 100644 --- a/crates/taskito-node/src/queue/inspect.rs +++ b/crates/taskito-node/src/queue/inspect.rs @@ -433,7 +433,12 @@ impl JsQueue { &worker_id, taskito_core::storage::REAPER_LOCK_TTL_MS, ) - .unwrap_or(false); + .unwrap_or_else(|e| { + // A backend error is not lost leadership — log it so a storage + // outage that stalls reaping is diagnosable, then skip this tick. + log::warn!("reaper election failed: {e}"); + false + }); if !leading { return Ok(Vec::new()); } diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index f29da7e8..92da8898 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -845,7 +845,12 @@ impl PyQueue { worker_id, taskito_core::storage::REAPER_LOCK_TTL_MS, ) - .unwrap_or(false); + .unwrap_or_else(|e| { + // A backend error is not lost leadership — log it so a storage outage + // that stalls reaping is diagnosable, then skip this tick. + log::warn!("reaper election failed: {e}"); + false + }); if !leading { return Ok(Vec::new()); }