From 7cb5fc51f08262b58f1517a8c346a54ab1960b04 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 28 May 2026 03:56:42 +0530 Subject: [PATCH 1/5] refactor: extract dead_letter ops into diesel_common --- .../src/storage/diesel_common/dead_letter.rs | 156 ++++++++++++++++++ .../src/storage/diesel_common/mod.rs | 2 + .../src/storage/postgres/dead_letter.rs | 141 +--------------- .../src/storage/sqlite/dead_letter.rs | 145 +--------------- 4 files changed, 160 insertions(+), 284 deletions(-) create mode 100644 crates/taskito-core/src/storage/diesel_common/dead_letter.rs diff --git a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs new file mode 100644 index 00000000..093c1787 --- /dev/null +++ b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs @@ -0,0 +1,156 @@ +/// Generates shared dead-letter operation methods for Diesel-backed storage backends. +/// +/// All four methods (`move_to_dlq`, `list_dead`, `retry_dead`, `purge_dead`) are +/// identical between SQLite and Postgres, so the entire file is macro-generated. +macro_rules! impl_diesel_dead_letter_ops { + ($storage_type:ty) => { + impl $storage_type { + /// Move a job to the dead letter queue and cascade-cancel dependents. + pub fn move_to_dlq( + &self, + job: &Job, + error: &str, + metadata: Option<&str>, + ) -> Result<()> { + let now = now_millis(); + let dlq_id = uuid::Uuid::now_v7().to_string(); + let mut conn = self.conn()?; + let job_id = job.id.clone(); + + conn.transaction(|conn| { + let dlq_row = NewDeadLetterRow { + id: &dlq_id, + original_job_id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + error: Some(error), + retry_count: job.retry_count, + failed_at: now, + metadata, + notes: job.notes.as_deref(), + priority: job.priority, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + result_ttl_ms: job.result_ttl_ms, + namespace: job.namespace.as_deref(), + }; + + diesel::insert_into(dead_letter::table) + .values(&dlq_row) + .execute(conn)?; + + diesel::update(jobs::table) + .filter(jobs::id.eq(&job.id)) + .set(( + jobs::status.eq(JobStatus::Dead as i32), + jobs::error.eq(error), + jobs::completed_at.eq(now), + )) + .execute(conn)?; + + Ok::<(), diesel::result::Error>(()) + })?; + + // Drop connection before cascade (needed for single-connection pools). + drop(conn); + + // Cascade cancel dependents. + self.cascade_cancel(&job_id, "dependency failed")?; + + Ok(()) + } + + /// List dead letter entries. + pub fn list_dead(&self, limit: i64, offset: i64) -> Result> { + let mut conn = self.conn()?; + + let rows: Vec = dead_letter::table + .order(dead_letter::failed_at.desc()) + .limit(limit) + .offset(offset) + .select(DeadLetterRow::as_select()) + .load(&mut conn)?; + + Ok(rows.into_iter().map(DeadJob::from).collect()) + } + + /// Re-enqueue a dead letter job. Returns the new job ID. + pub fn retry_dead(&self, dead_id: &str) -> Result { + let mut conn = self.conn()?; + + let dead_row: DeadLetterRow = dead_letter::table + .find(dead_id) + .select(DeadLetterRow::as_select()) + .first(&mut conn) + .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; + + let new_job = NewJob { + queue: dead_row.queue, + task_name: dead_row.task_name, + payload: dead_row.payload, + priority: dead_row.priority, + scheduled_at: now_millis(), + max_retries: dead_row.max_retries, + timeout_ms: dead_row.timeout_ms, + unique_key: None, + metadata: dead_row.metadata, + notes: dead_row.notes, + depends_on: vec![], + expires_at: None, + result_ttl_ms: dead_row.result_ttl_ms, + namespace: dead_row.namespace, + }; + + let job = new_job.into_job(); + + conn.transaction(|conn| { + let row = super::super::models::NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, + payload: &job.payload, + status: job.status as i32, + priority: job.priority, + created_at: job.created_at, + scheduled_at: job.scheduled_at, + retry_count: job.retry_count, + max_retries: job.max_retries, + timeout_ms: job.timeout_ms, + unique_key: job.unique_key.as_deref(), + metadata: job.metadata.as_deref(), + notes: job.notes.as_deref(), + cancel_requested: 0, + expires_at: job.expires_at, + result_ttl_ms: job.result_ttl_ms, + namespace: job.namespace.as_deref(), + }; + + diesel::insert_into(jobs::table) + .values(&row) + .execute(conn)?; + + diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; + + Ok::<(), diesel::result::Error>(()) + })?; + + Ok(job.id) + } + + /// Purge dead letter entries older than the given timestamp. + pub fn purge_dead(&self, older_than_ms: i64) -> Result { + let mut conn = self.conn()?; + + let affected = diesel::delete( + dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms)), + ) + .execute(&mut conn)?; + + Ok(affected as u64) + } + } + }; +} + +pub(crate) use impl_diesel_dead_letter_ops; diff --git a/crates/taskito-core/src/storage/diesel_common/mod.rs b/crates/taskito-core/src/storage/diesel_common/mod.rs index 70eae41e..a350b989 100644 --- a/crates/taskito-core/src/storage/diesel_common/mod.rs +++ b/crates/taskito-core/src/storage/diesel_common/mod.rs @@ -3,12 +3,14 @@ //! These macros generate the identical method implementations that both //! Diesel-backed storage backends share, avoiding code duplication. +mod dead_letter; mod jobs; mod locks; mod logs; mod metrics; pub(crate) mod migrations; +pub(crate) use dead_letter::impl_diesel_dead_letter_ops; pub(crate) use jobs::impl_diesel_job_ops; pub(crate) use locks::impl_diesel_lock_ops; pub(crate) use logs::impl_diesel_log_ops; diff --git a/crates/taskito-core/src/storage/postgres/dead_letter.rs b/crates/taskito-core/src/storage/postgres/dead_letter.rs index 1c122267..24f3edc6 100644 --- a/crates/taskito-core/src/storage/postgres/dead_letter.rs +++ b/crates/taskito-core/src/storage/postgres/dead_letter.rs @@ -7,143 +7,4 @@ use crate::error::{QueueError, Result}; use crate::job::{now_millis, Job, JobStatus, NewJob}; use crate::storage::DeadJob; -impl PostgresStorage { - /// Move a job to the dead letter queue and cascade-cancel dependents. - pub fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { - let now = now_millis(); - let dlq_id = uuid::Uuid::now_v7().to_string(); - let mut conn = self.conn()?; - let job_id = job.id.clone(); - - conn.transaction(|conn| { - let dlq_row = NewDeadLetterRow { - id: &dlq_id, - original_job_id: &job.id, - queue: &job.queue, - task_name: &job.task_name, - payload: &job.payload, - error: Some(error), - retry_count: job.retry_count, - failed_at: now, - metadata, - notes: job.notes.as_deref(), - priority: job.priority, - max_retries: job.max_retries, - timeout_ms: job.timeout_ms, - result_ttl_ms: job.result_ttl_ms, - namespace: job.namespace.as_deref(), - }; - - diesel::insert_into(dead_letter::table) - .values(&dlq_row) - .execute(conn)?; - - diesel::update(jobs::table) - .filter(jobs::id.eq(&job.id)) - .set(( - jobs::status.eq(JobStatus::Dead as i32), - jobs::error.eq(error), - jobs::completed_at.eq(now), - )) - .execute(conn)?; - - Ok::<(), diesel::result::Error>(()) - })?; - - // Drop connection before cascade (needed for single-connection pools) - drop(conn); - - // Cascade cancel dependents - self.cascade_cancel(&job_id, "dependency failed")?; - - Ok(()) - } - - /// List dead letter entries. - pub fn list_dead(&self, limit: i64, offset: i64) -> Result> { - let mut conn = self.conn()?; - - let rows: Vec = dead_letter::table - .order(dead_letter::failed_at.desc()) - .limit(limit) - .offset(offset) - .select(DeadLetterRow::as_select()) - .load(&mut conn)?; - - Ok(rows.into_iter().map(DeadJob::from).collect()) - } - - /// Re-enqueue a dead letter job. Returns the new job ID. - pub fn retry_dead(&self, dead_id: &str) -> Result { - let mut conn = self.conn()?; - - let dead_row: DeadLetterRow = dead_letter::table - .find(dead_id) - .select(DeadLetterRow::as_select()) - .first(&mut conn) - .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; - - let new_job = NewJob { - queue: dead_row.queue, - task_name: dead_row.task_name, - payload: dead_row.payload, - priority: dead_row.priority, - scheduled_at: now_millis(), - max_retries: dead_row.max_retries, - timeout_ms: dead_row.timeout_ms, - unique_key: None, - metadata: dead_row.metadata, - notes: dead_row.notes, - depends_on: vec![], - expires_at: None, - result_ttl_ms: dead_row.result_ttl_ms, - namespace: dead_row.namespace, - }; - - let job = new_job.into_job(); - - conn.transaction(|conn| { - let row = super::super::models::NewJobRow { - id: &job.id, - queue: &job.queue, - task_name: &job.task_name, - payload: &job.payload, - status: job.status as i32, - priority: job.priority, - created_at: job.created_at, - scheduled_at: job.scheduled_at, - retry_count: job.retry_count, - max_retries: job.max_retries, - timeout_ms: job.timeout_ms, - unique_key: job.unique_key.as_deref(), - metadata: job.metadata.as_deref(), - notes: job.notes.as_deref(), - cancel_requested: 0, - expires_at: job.expires_at, - result_ttl_ms: job.result_ttl_ms, - namespace: job.namespace.as_deref(), - }; - - diesel::insert_into(jobs::table) - .values(&row) - .execute(conn)?; - - diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; - - Ok::<(), diesel::result::Error>(()) - })?; - - Ok(job.id) - } - - /// Purge dead letter entries older than the given timestamp. - pub fn purge_dead(&self, older_than_ms: i64) -> Result { - let mut conn = self.conn()?; - - let affected = - diesel::delete(dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms))) - .execute(&mut conn)?; - - Ok(affected as u64) - } -} +crate::storage::diesel_common::impl_diesel_dead_letter_ops!(PostgresStorage); diff --git a/crates/taskito-core/src/storage/sqlite/dead_letter.rs b/crates/taskito-core/src/storage/sqlite/dead_letter.rs index 6e3a3c33..8c1096c5 100644 --- a/crates/taskito-core/src/storage/sqlite/dead_letter.rs +++ b/crates/taskito-core/src/storage/sqlite/dead_letter.rs @@ -7,147 +7,4 @@ use crate::error::{QueueError, Result}; use crate::job::{now_millis, Job, JobStatus, NewJob}; use crate::storage::DeadJob; -impl SqliteStorage { - /// Move a job to the dead letter queue and cascade-cancel dependents. - pub fn move_to_dlq(&self, job: &Job, error: &str, metadata: Option<&str>) -> Result<()> { - let now = now_millis(); - let dlq_id = uuid::Uuid::now_v7().to_string(); - let mut conn = self.conn()?; - let job_id = job.id.clone(); - - conn.transaction(|conn| { - let dlq_row = NewDeadLetterRow { - id: &dlq_id, - original_job_id: &job.id, - queue: &job.queue, - task_name: &job.task_name, - payload: &job.payload, - error: Some(error), - retry_count: job.retry_count, - failed_at: now, - metadata, - notes: job.notes.as_deref(), - priority: job.priority, - max_retries: job.max_retries, - timeout_ms: job.timeout_ms, - result_ttl_ms: job.result_ttl_ms, - namespace: job.namespace.as_deref(), - }; - - diesel::insert_into(dead_letter::table) - .values(&dlq_row) - .execute(conn)?; - - diesel::update(jobs::table) - .filter(jobs::id.eq(&job.id)) - .set(( - jobs::status.eq(JobStatus::Dead as i32), - jobs::error.eq(error), - jobs::completed_at.eq(now), - )) - .execute(conn)?; - - Ok::<(), diesel::result::Error>(()) - })?; - - // Drop connection before cascade (needed for single-connection pools). - drop(conn); - - // Cascade cancel dependents. Errors propagate so callers can decide how - // to react; parity with the Postgres and Redis backends. Note: the DLQ - // row has already been committed, so a failure here leaves a partial - // state (DLQ entry present, dependents possibly uncancelled) — callers - // should log and alert, not silently retry `move_to_dlq`. - self.cascade_cancel(&job_id, "dependency failed")?; - - Ok(()) - } - - /// List dead letter entries. - pub fn list_dead(&self, limit: i64, offset: i64) -> Result> { - let mut conn = self.conn()?; - - let rows: Vec = dead_letter::table - .order(dead_letter::failed_at.desc()) - .limit(limit) - .offset(offset) - .select(DeadLetterRow::as_select()) - .load(&mut conn)?; - - Ok(rows.into_iter().map(DeadJob::from).collect()) - } - - /// Re-enqueue a dead letter job. Returns the new job ID. - pub fn retry_dead(&self, dead_id: &str) -> Result { - let mut conn = self.conn()?; - - let dead_row: DeadLetterRow = dead_letter::table - .find(dead_id) - .select(DeadLetterRow::as_select()) - .first(&mut conn) - .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; - - let new_job = NewJob { - queue: dead_row.queue, - task_name: dead_row.task_name, - payload: dead_row.payload, - priority: dead_row.priority, - scheduled_at: now_millis(), - max_retries: dead_row.max_retries, - timeout_ms: dead_row.timeout_ms, - unique_key: None, - metadata: dead_row.metadata, - notes: dead_row.notes, - depends_on: vec![], - expires_at: None, - result_ttl_ms: dead_row.result_ttl_ms, - namespace: dead_row.namespace, - }; - - let job = new_job.into_job(); - - conn.transaction(|conn| { - let row = super::super::models::NewJobRow { - id: &job.id, - queue: &job.queue, - task_name: &job.task_name, - payload: &job.payload, - status: job.status as i32, - priority: job.priority, - created_at: job.created_at, - scheduled_at: job.scheduled_at, - retry_count: job.retry_count, - max_retries: job.max_retries, - timeout_ms: job.timeout_ms, - unique_key: job.unique_key.as_deref(), - metadata: job.metadata.as_deref(), - notes: job.notes.as_deref(), - cancel_requested: 0, - expires_at: job.expires_at, - result_ttl_ms: job.result_ttl_ms, - namespace: job.namespace.as_deref(), - }; - - diesel::insert_into(jobs::table) - .values(&row) - .execute(conn)?; - - diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; - - Ok::<(), diesel::result::Error>(()) - })?; - - Ok(job.id) - } - - /// Purge dead letter entries older than the given timestamp. - pub fn purge_dead(&self, older_than_ms: i64) -> Result { - let mut conn = self.conn()?; - - let affected = - diesel::delete(dead_letter::table.filter(dead_letter::failed_at.lt(older_than_ms))) - .execute(&mut conn)?; - - Ok(affected as u64) - } -} +crate::storage::diesel_common::impl_diesel_dead_letter_ops!(SqliteStorage); From 6932c2ec656bc104cca02ce8e176ae20a725a2c0 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 28 May 2026 03:58:18 +0530 Subject: [PATCH 2/5] refactor: extract archival list_archived into diesel_common --- .../src/storage/diesel_common/archival.rs | 52 +++++++++++++++++++ .../src/storage/diesel_common/mod.rs | 2 + .../src/storage/postgres/archival.rs | 42 +-------------- .../src/storage/sqlite/archival.rs | 43 +-------------- 4 files changed, 58 insertions(+), 81 deletions(-) create mode 100644 crates/taskito-core/src/storage/diesel_common/archival.rs diff --git a/crates/taskito-core/src/storage/diesel_common/archival.rs b/crates/taskito-core/src/storage/diesel_common/archival.rs new file mode 100644 index 00000000..72d50d26 --- /dev/null +++ b/crates/taskito-core/src/storage/diesel_common/archival.rs @@ -0,0 +1,52 @@ +/// Generates the shared `list_archived` method for Diesel-backed storage backends. +/// +/// `archive_old_jobs` differs between SQLite (`INSERT OR IGNORE`) and Postgres +/// (`ON CONFLICT DO NOTHING` + `$1` bind syntax), so it remains backend-specific. +macro_rules! impl_diesel_archival_ops { + ($storage_type:ty) => { + impl $storage_type { + /// List archived jobs with pagination. + pub fn list_archived(&self, limit: i64, offset: i64) -> Result> { + let mut conn = self.conn()?; + + let rows: Vec = archived_jobs::table + .order(archived_jobs::completed_at.desc()) + .limit(limit) + .offset(offset) + .select(ArchivedJobRow::as_select()) + .load(&mut conn)?; + + Ok(rows + .into_iter() + .map(|row| Job { + id: row.id, + queue: row.queue, + task_name: row.task_name, + payload: row.payload, + status: JobStatus::from_i32(row.status).unwrap_or(JobStatus::Pending), + priority: row.priority, + created_at: row.created_at, + scheduled_at: row.scheduled_at, + started_at: row.started_at, + completed_at: row.completed_at, + retry_count: row.retry_count, + max_retries: row.max_retries, + result: row.result, + error: row.error, + timeout_ms: row.timeout_ms, + unique_key: row.unique_key, + progress: row.progress, + metadata: row.metadata, + notes: row.notes, + cancel_requested: row.cancel_requested != 0, + expires_at: row.expires_at, + result_ttl_ms: row.result_ttl_ms, + namespace: row.namespace, + }) + .collect()) + } + } + }; +} + +pub(crate) use impl_diesel_archival_ops; diff --git a/crates/taskito-core/src/storage/diesel_common/mod.rs b/crates/taskito-core/src/storage/diesel_common/mod.rs index a350b989..976e16e6 100644 --- a/crates/taskito-core/src/storage/diesel_common/mod.rs +++ b/crates/taskito-core/src/storage/diesel_common/mod.rs @@ -3,6 +3,7 @@ //! These macros generate the identical method implementations that both //! Diesel-backed storage backends share, avoiding code duplication. +mod archival; mod dead_letter; mod jobs; mod locks; @@ -10,6 +11,7 @@ mod logs; mod metrics; pub(crate) mod migrations; +pub(crate) use archival::impl_diesel_archival_ops; pub(crate) use dead_letter::impl_diesel_dead_letter_ops; pub(crate) use jobs::impl_diesel_job_ops; pub(crate) use locks::impl_diesel_lock_ops; diff --git a/crates/taskito-core/src/storage/postgres/archival.rs b/crates/taskito-core/src/storage/postgres/archival.rs index 241f303d..1fe64b99 100644 --- a/crates/taskito-core/src/storage/postgres/archival.rs +++ b/crates/taskito-core/src/storage/postgres/archival.rs @@ -6,6 +6,8 @@ use super::PostgresStorage; use crate::error::Result; use crate::job::{Job, JobStatus}; +crate::storage::diesel_common::impl_diesel_archival_ops!(PostgresStorage); + impl PostgresStorage { pub fn archive_old_jobs(&self, cutoff_ms: i64) -> Result { let mut conn = self.conn()?; @@ -35,44 +37,4 @@ impl PostgresStorage { Ok(affected as u64) }) } - - pub fn list_archived(&self, limit: i64, offset: i64) -> Result> { - let mut conn = self.conn()?; - - let rows: Vec = archived_jobs::table - .order(archived_jobs::completed_at.desc()) - .limit(limit) - .offset(offset) - .select(ArchivedJobRow::as_select()) - .load(&mut conn)?; - - Ok(rows - .into_iter() - .map(|row| Job { - id: row.id, - queue: row.queue, - task_name: row.task_name, - payload: row.payload, - status: JobStatus::from_i32(row.status).unwrap_or(JobStatus::Pending), - priority: row.priority, - created_at: row.created_at, - scheduled_at: row.scheduled_at, - started_at: row.started_at, - completed_at: row.completed_at, - retry_count: row.retry_count, - max_retries: row.max_retries, - result: row.result, - error: row.error, - timeout_ms: row.timeout_ms, - unique_key: row.unique_key, - progress: row.progress, - metadata: row.metadata, - notes: row.notes, - cancel_requested: row.cancel_requested != 0, - expires_at: row.expires_at, - result_ttl_ms: row.result_ttl_ms, - namespace: row.namespace, - }) - .collect()) - } } diff --git a/crates/taskito-core/src/storage/sqlite/archival.rs b/crates/taskito-core/src/storage/sqlite/archival.rs index b198ae5e..0936c7bf 100644 --- a/crates/taskito-core/src/storage/sqlite/archival.rs +++ b/crates/taskito-core/src/storage/sqlite/archival.rs @@ -6,6 +6,8 @@ use super::SqliteStorage; use crate::error::Result; use crate::job::{Job, JobStatus}; +crate::storage::diesel_common::impl_diesel_archival_ops!(SqliteStorage); + impl SqliteStorage { /// Archive completed/dead/cancelled jobs older than cutoff_ms. /// Moves them from `jobs` to `archived_jobs`. @@ -37,45 +39,4 @@ impl SqliteStorage { Ok(affected as u64) }) } - - /// List archived jobs with pagination. - pub fn list_archived(&self, limit: i64, offset: i64) -> Result> { - let mut conn = self.conn()?; - - let rows: Vec = archived_jobs::table - .order(archived_jobs::completed_at.desc()) - .limit(limit) - .offset(offset) - .select(ArchivedJobRow::as_select()) - .load(&mut conn)?; - - Ok(rows - .into_iter() - .map(|row| Job { - id: row.id, - queue: row.queue, - task_name: row.task_name, - payload: row.payload, - status: JobStatus::from_i32(row.status).unwrap_or(JobStatus::Pending), - priority: row.priority, - created_at: row.created_at, - scheduled_at: row.scheduled_at, - started_at: row.started_at, - completed_at: row.completed_at, - retry_count: row.retry_count, - max_retries: row.max_retries, - result: row.result, - error: row.error, - timeout_ms: row.timeout_ms, - unique_key: row.unique_key, - progress: row.progress, - metadata: row.metadata, - notes: row.notes, - cancel_requested: row.cancel_requested != 0, - expires_at: row.expires_at, - result_ttl_ms: row.result_ttl_ms, - namespace: row.namespace, - }) - .collect()) - } } From 0162b5fa2bab7682e40f24ec6fcdc8707c87336e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 28 May 2026 03:59:28 +0530 Subject: [PATCH 3/5] refactor: extract worker ops into diesel_common --- .../src/storage/diesel_common/mod.rs | 2 + .../src/storage/diesel_common/workers.rs | 91 +++++++++++++++++++ .../src/storage/postgres/workers.rs | 82 +---------------- .../src/storage/sqlite/workers.rs | 82 +---------------- 4 files changed, 97 insertions(+), 160 deletions(-) create mode 100644 crates/taskito-core/src/storage/diesel_common/workers.rs diff --git a/crates/taskito-core/src/storage/diesel_common/mod.rs b/crates/taskito-core/src/storage/diesel_common/mod.rs index 976e16e6..a7b8f9ad 100644 --- a/crates/taskito-core/src/storage/diesel_common/mod.rs +++ b/crates/taskito-core/src/storage/diesel_common/mod.rs @@ -10,6 +10,7 @@ mod locks; mod logs; mod metrics; pub(crate) mod migrations; +mod workers; pub(crate) use archival::impl_diesel_archival_ops; pub(crate) use dead_letter::impl_diesel_dead_letter_ops; @@ -17,3 +18,4 @@ pub(crate) use jobs::impl_diesel_job_ops; pub(crate) use locks::impl_diesel_lock_ops; pub(crate) use logs::impl_diesel_log_ops; pub(crate) use metrics::impl_diesel_metric_ops; +pub(crate) use workers::impl_diesel_worker_ops; diff --git a/crates/taskito-core/src/storage/diesel_common/workers.rs b/crates/taskito-core/src/storage/diesel_common/workers.rs new file mode 100644 index 00000000..bdcfd430 --- /dev/null +++ b/crates/taskito-core/src/storage/diesel_common/workers.rs @@ -0,0 +1,91 @@ +/// Generates shared worker operation methods for Diesel-backed storage backends. +/// +/// `register_worker` differs between SQLite (`replace_into`) and Postgres +/// (`insert_into...on_conflict.do_update`), so it remains backend-specific. +macro_rules! impl_diesel_worker_ops { + ($storage_type:ty) => { + impl $storage_type { + /// Update the heartbeat timestamp for a worker. + pub fn heartbeat(&self, worker_id: &str, resource_health: Option<&str>) -> Result<()> { + let mut conn = self.conn()?; + let now = now_millis(); + + diesel::update(workers::table) + .filter(workers::worker_id.eq(worker_id)) + .set(( + workers::last_heartbeat.eq(now), + workers::resource_health.eq(resource_health), + )) + .execute(&mut conn)?; + + Ok(()) + } + + /// Update the status of a worker. + pub fn update_worker_status(&self, worker_id: &str, status: &str) -> Result<()> { + let mut conn = self.conn()?; + + diesel::update(workers::table) + .filter(workers::worker_id.eq(worker_id)) + .set(workers::status.eq(status)) + .execute(&mut conn)?; + + Ok(()) + } + + /// List all workers with their heartbeat status. + pub fn list_workers(&self) -> Result> { + let mut conn = self.conn()?; + + let rows = workers::table + .select(WorkerRow::as_select()) + .load(&mut conn)?; + + Ok(rows) + } + + /// Remove workers that haven't sent a heartbeat within the threshold. + /// Returns the IDs of the reaped workers. + pub fn reap_dead_workers(&self) -> Result> { + let mut conn = self.conn()?; + let cutoff = now_millis().saturating_sub(DEAD_WORKER_THRESHOLD_MS); + + let dead_ids: Vec = workers::table + .filter(workers::last_heartbeat.lt(cutoff)) + .select(workers::worker_id) + .load(&mut conn)?; + + if !dead_ids.is_empty() { + diesel::delete(workers::table.filter(workers::worker_id.eq_any(&dead_ids))) + .execute(&mut conn)?; + } + + Ok(dead_ids) + } + + /// Unregister a worker (called on shutdown). + pub fn unregister_worker(&self, worker_id: &str) -> Result<()> { + let mut conn = self.conn()?; + + diesel::delete(workers::table.filter(workers::worker_id.eq(worker_id))) + .execute(&mut conn)?; + + Ok(()) + } + + /// List all job IDs currently claimed by a worker. + pub fn list_claims_by_worker(&self, worker_id: &str) -> Result> { + let mut conn = self.conn()?; + + let job_ids: Vec = execution_claims::table + .filter(execution_claims::worker_id.eq(worker_id)) + .select(execution_claims::job_id) + .load(&mut conn)?; + + Ok(job_ids) + } + } + }; +} + +pub(crate) use impl_diesel_worker_ops; diff --git a/crates/taskito-core/src/storage/postgres/workers.rs b/crates/taskito-core/src/storage/postgres/workers.rs index 943cc576..9f9d4cae 100644 --- a/crates/taskito-core/src/storage/postgres/workers.rs +++ b/crates/taskito-core/src/storage/postgres/workers.rs @@ -7,6 +7,8 @@ use super::PostgresStorage; use crate::error::Result; use crate::job::now_millis; +crate::storage::diesel_common::impl_diesel_worker_ops!(PostgresStorage); + impl PostgresStorage { /// Register a new worker or update an existing one. #[allow(clippy::too_many_arguments)] @@ -49,84 +51,4 @@ impl PostgresStorage { Ok(()) } - - /// Update the heartbeat timestamp for a worker. - pub fn heartbeat(&self, worker_id: &str, resource_health: Option<&str>) -> Result<()> { - let mut conn = self.conn()?; - let now = now_millis(); - - diesel::update(workers::table) - .filter(workers::worker_id.eq(worker_id)) - .set(( - workers::last_heartbeat.eq(now), - workers::resource_health.eq(resource_health), - )) - .execute(&mut conn)?; - - Ok(()) - } - - /// Update the status of a worker. - pub fn update_worker_status(&self, worker_id: &str, status: &str) -> Result<()> { - let mut conn = self.conn()?; - - diesel::update(workers::table) - .filter(workers::worker_id.eq(worker_id)) - .set(workers::status.eq(status)) - .execute(&mut conn)?; - - Ok(()) - } - - /// List all workers with their heartbeat status. - pub fn list_workers(&self) -> Result> { - let mut conn = self.conn()?; - - let rows = workers::table - .select(WorkerRow::as_select()) - .load(&mut conn)?; - - Ok(rows) - } - - /// Remove workers that haven't sent a heartbeat within the threshold. - /// Returns the IDs of the reaped workers. - pub fn reap_dead_workers(&self) -> Result> { - let mut conn = self.conn()?; - let cutoff = now_millis().saturating_sub(DEAD_WORKER_THRESHOLD_MS); - - let dead_ids: Vec = workers::table - .filter(workers::last_heartbeat.lt(cutoff)) - .select(workers::worker_id) - .load(&mut conn)?; - - if !dead_ids.is_empty() { - diesel::delete(workers::table.filter(workers::worker_id.eq_any(&dead_ids))) - .execute(&mut conn)?; - } - - Ok(dead_ids) - } - - /// Unregister a worker (called on shutdown). - pub fn unregister_worker(&self, worker_id: &str) -> Result<()> { - let mut conn = self.conn()?; - - diesel::delete(workers::table.filter(workers::worker_id.eq(worker_id))) - .execute(&mut conn)?; - - Ok(()) - } - - /// List all job IDs currently claimed by a worker. - pub fn list_claims_by_worker(&self, worker_id: &str) -> Result> { - let mut conn = self.conn()?; - - let job_ids: Vec = execution_claims::table - .filter(execution_claims::worker_id.eq(worker_id)) - .select(execution_claims::job_id) - .load(&mut conn)?; - - Ok(job_ids) - } } diff --git a/crates/taskito-core/src/storage/sqlite/workers.rs b/crates/taskito-core/src/storage/sqlite/workers.rs index 642cc997..8f581cf2 100644 --- a/crates/taskito-core/src/storage/sqlite/workers.rs +++ b/crates/taskito-core/src/storage/sqlite/workers.rs @@ -7,6 +7,8 @@ use super::SqliteStorage; use crate::error::Result; use crate::job::now_millis; +crate::storage::diesel_common::impl_diesel_worker_ops!(SqliteStorage); + impl SqliteStorage { /// Register a new worker or update an existing one. #[allow(clippy::too_many_arguments)] @@ -46,84 +48,4 @@ impl SqliteStorage { Ok(()) } - - /// Update the heartbeat timestamp for a worker. - pub fn heartbeat(&self, worker_id: &str, resource_health: Option<&str>) -> Result<()> { - let mut conn = self.conn()?; - let now = now_millis(); - - diesel::update(workers::table) - .filter(workers::worker_id.eq(worker_id)) - .set(( - workers::last_heartbeat.eq(now), - workers::resource_health.eq(resource_health), - )) - .execute(&mut conn)?; - - Ok(()) - } - - /// Update the status of a worker. - pub fn update_worker_status(&self, worker_id: &str, status: &str) -> Result<()> { - let mut conn = self.conn()?; - - diesel::update(workers::table) - .filter(workers::worker_id.eq(worker_id)) - .set(workers::status.eq(status)) - .execute(&mut conn)?; - - Ok(()) - } - - /// List all workers with their heartbeat status. - pub fn list_workers(&self) -> Result> { - let mut conn = self.conn()?; - - let rows = workers::table - .select(WorkerRow::as_select()) - .load(&mut conn)?; - - Ok(rows) - } - - /// Remove workers that haven't sent a heartbeat within the threshold. - /// Returns the IDs of the reaped workers. - pub fn reap_dead_workers(&self) -> Result> { - let mut conn = self.conn()?; - let cutoff = now_millis().saturating_sub(DEAD_WORKER_THRESHOLD_MS); - - let dead_ids: Vec = workers::table - .filter(workers::last_heartbeat.lt(cutoff)) - .select(workers::worker_id) - .load(&mut conn)?; - - if !dead_ids.is_empty() { - diesel::delete(workers::table.filter(workers::worker_id.eq_any(&dead_ids))) - .execute(&mut conn)?; - } - - Ok(dead_ids) - } - - /// Unregister a worker (called on shutdown). - pub fn unregister_worker(&self, worker_id: &str) -> Result<()> { - let mut conn = self.conn()?; - - diesel::delete(workers::table.filter(workers::worker_id.eq(worker_id))) - .execute(&mut conn)?; - - Ok(()) - } - - /// List all job IDs currently claimed by a worker. - pub fn list_claims_by_worker(&self, worker_id: &str) -> Result> { - let mut conn = self.conn()?; - - let job_ids: Vec = execution_claims::table - .filter(execution_claims::worker_id.eq(worker_id)) - .select(execution_claims::job_id) - .load(&mut conn)?; - - Ok(job_ids) - } } From 95da532b1d1dfa1a32e777b65a9fd2d69f1e7583 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 28 May 2026 11:17:11 +0530 Subject: [PATCH 4/5] fix(dead_letter): preserve real DB errors and detect retry races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit retry_dead previously masked every Diesel error as JobNotFound and did not verify the dead-letter delete affected a row. Two concurrent retries could both commit a fresh job for the same dead id. Match on diesel::result::Error::NotFound explicitly, propagate other errors as Storage, and assert the in-transaction delete removed exactly one row — otherwise abort the transaction so the freshly inserted job rolls back. --- .../src/storage/diesel_common/dead_letter.rs | 39 +++++++++++++++---- .../taskito-core/src/storage/sqlite/tests.rs | 10 +++++ 2 files changed, 42 insertions(+), 7 deletions(-) diff --git a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs index 093c1787..495da038 100644 --- a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs +++ b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs @@ -76,14 +76,27 @@ macro_rules! impl_diesel_dead_letter_ops { } /// Re-enqueue a dead letter job. Returns the new job ID. + /// + /// Returns `JobNotFound` only when the dead-letter row is absent or + /// is removed concurrently by another retry; other Diesel errors + /// propagate as `Storage` so real DB failures aren't masked. The + /// delete inside the transaction asserts a row was actually removed, + /// preventing two concurrent retries from both committing a fresh + /// enqueue against the same dead-letter id. pub fn retry_dead(&self, dead_id: &str) -> Result { let mut conn = self.conn()?; - let dead_row: DeadLetterRow = dead_letter::table + let dead_row: DeadLetterRow = match dead_letter::table .find(dead_id) .select(DeadLetterRow::as_select()) .first(&mut conn) - .map_err(|_| QueueError::JobNotFound(dead_id.to_string()))?; + { + Ok(row) => row, + Err(diesel::result::Error::NotFound) => { + return Err(QueueError::JobNotFound(dead_id.to_string())); + } + Err(err) => return Err(err.into()), + }; let new_job = NewJob { queue: dead_row.queue, @@ -104,7 +117,7 @@ macro_rules! impl_diesel_dead_letter_ops { let job = new_job.into_job(); - conn.transaction(|conn| { + let result = conn.transaction(|conn| { let row = super::super::models::NewJobRow { id: &job.id, queue: &job.queue, @@ -130,12 +143,24 @@ macro_rules! impl_diesel_dead_letter_ops { .values(&row) .execute(conn)?; - diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; + let deleted = diesel::delete(dead_letter::table.find(dead_id)).execute(conn)?; + if deleted == 0 { + // A concurrent retry won the race. Roll back the + // freshly-inserted job by returning an error from the + // transaction closure — Diesel aborts the txn. + return Err(diesel::result::Error::NotFound); + } Ok::<(), diesel::result::Error>(()) - })?; - - Ok(job.id) + }); + + match result { + Ok(()) => Ok(job.id), + Err(diesel::result::Error::NotFound) => { + Err(QueueError::JobNotFound(dead_id.to_string())) + } + Err(err) => Err(err.into()), + } } /// Purge dead letter entries older than the given timestamp. diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index b9cd35fd..a31c4b15 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::error::QueueError; use crate::job::{now_millis, JobStatus, NewJob}; fn test_storage() -> SqliteStorage { @@ -227,6 +228,15 @@ fn test_retry_dead() { assert!(dead.is_empty()); } +#[test] +fn test_retry_dead_missing_id_returns_not_found() { + let storage = test_storage(); + match storage.retry_dead("does-not-exist") { + Err(QueueError::JobNotFound(id)) => assert_eq!(id, "does-not-exist"), + other => panic!("expected JobNotFound, got {other:?}"), + } +} + #[test] fn test_stats() { let storage = test_storage(); From d86e050a3c03691cab7190971b0f316d551b0fa8 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Thu, 28 May 2026 11:17:42 +0530 Subject: [PATCH 5/5] fix(workers): re-check heartbeat in reap delete to avoid TOCTOU reap_dead_workers scanned for stale workers, then deleted by worker_id alone. A worker that heartbeated between the scan and the delete was still erased, deregistering a healthy fleet member. Apply the last_heartbeat < cutoff predicate to the DELETE as well so only workers that are still stale at delete time are removed. --- .../src/storage/diesel_common/workers.rs | 15 ++++++- .../taskito-core/src/storage/sqlite/tests.rs | 39 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/crates/taskito-core/src/storage/diesel_common/workers.rs b/crates/taskito-core/src/storage/diesel_common/workers.rs index bdcfd430..4bbbc790 100644 --- a/crates/taskito-core/src/storage/diesel_common/workers.rs +++ b/crates/taskito-core/src/storage/diesel_common/workers.rs @@ -46,6 +46,13 @@ macro_rules! impl_diesel_worker_ops { /// Remove workers that haven't sent a heartbeat within the threshold. /// Returns the IDs of the reaped workers. + /// + /// The delete re-applies the `last_heartbeat < cutoff` predicate so + /// a worker that heartbeats between the scan and the delete is not + /// reaped erroneously. The returned id list reflects the scan + /// snapshot and may slightly over-report the actual reaped set when + /// such a race occurs — callers use the list for event emission and + /// 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); @@ -56,8 +63,12 @@ macro_rules! impl_diesel_worker_ops { .load(&mut conn)?; if !dead_ids.is_empty() { - diesel::delete(workers::table.filter(workers::worker_id.eq_any(&dead_ids))) - .execute(&mut conn)?; + diesel::delete( + workers::table + .filter(workers::worker_id.eq_any(&dead_ids)) + .filter(workers::last_heartbeat.lt(cutoff)), + ) + .execute(&mut conn)?; } Ok(dead_ids) diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index a31c4b15..0dba4b13 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -237,6 +237,45 @@ fn test_retry_dead_missing_id_returns_not_found() { } } +#[test] +fn test_reap_dead_workers_removes_stale_keeps_fresh() { + 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(); + + // Backdate `stale` past the dead-worker threshold (30s) so it is reaped; + // `fresh` keeps its current heartbeat. This also exercises the new + // double-predicate in the DELETE — even if a worker_id ends up in the + // scan list, the DELETE only removes rows whose heartbeat is still stale. + let cutoff = now_millis() - crate::storage::DEAD_WORKER_THRESHOLD_MS - 1_000; + let mut conn = storage.conn().unwrap(); + diesel::update(workers::table.filter(workers::worker_id.eq("stale"))) + .set(workers::last_heartbeat.eq(cutoff)) + .execute(&mut conn) + .unwrap(); + drop(conn); + + let reaped = storage.reap_dead_workers().unwrap(); + assert_eq!(reaped, vec!["stale".to_string()]); + + let surviving: Vec = storage + .list_workers() + .unwrap() + .into_iter() + .map(|w| w.worker_id) + .collect(); + assert!(surviving.contains(&"fresh".to_string())); + assert!(!surviving.contains(&"stale".to_string())); +} + #[test] fn test_stats() { let storage = test_storage();