fix: storage transaction safety

crates/taskito-core/src/resilience/rate_limiter.rs

@@ -118,19 +118,22 @@ mod tests {
             handles.push(std::thread::spawn(move || {
                 barrier.wait();
                 // Retry on lock contention (SQLite busy)
+                let mut last_err = None;
                 for _ in 0..10 {
                     match limiter.try_acquire("concurrent_test", &config) {
                         Ok(true) => {
                             acquired.fetch_add(1, Ordering::Relaxed);
                             return;
                         }
                         Ok(false) => return,
-                        Err(_) => {
+                        Err(e) => {
+                            last_err = Some(e.to_string());
                             std::thread::sleep(std::time::Duration::from_millis(10));
                             continue;
                         }
                     }
                 }
+                panic!("exhausted retries in SQLite rate-limit test: {last_err:?}");
             }));
         }
 
@@ -141,4 +144,79 @@ mod tests {
         // With 10 tokens and no refill, exactly 10 should succeed
         assert_eq!(acquired.load(Ordering::Relaxed), 10);
     }
+
+    /// Postgres lost-update guard: under concurrency the row is seeded once and
+    /// read `FOR UPDATE`, so a fresh bucket must admit exactly `max_tokens` — not
+    /// more. Skips unless `TASKITO_POSTGRES_TEST_URL` is set (CI Postgres job).
+    #[cfg(feature = "postgres")]
+    #[test]
+    fn test_postgres_concurrent_token_acquisition_no_over_admit() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+        use std::sync::{Arc, Barrier};
+
+        let url = match std::env::var("TASKITO_POSTGRES_TEST_URL") {
+            Ok(u) => u,
+            Err(_) => {
+                eprintln!("Skipping Postgres rate-limit test (TASKITO_POSTGRES_TEST_URL not set)");
+                return;
+            }
+        };
+        let num_threads = 12;
+        let storage = match crate::storage::postgres::PostgresStorage::with_pool_size(
+            &url,
+            num_threads as u32 + 4,
+        ) {
+            Ok(s) => StorageBackend::Postgres(s),
+            Err(e) => {
+                eprintln!("Skipping Postgres rate-limit test (cannot connect): {e}");
+                return;
+            }
+        };
+        let limiter = Arc::new(RateLimiter::new(storage));
+        let config = RateLimitConfig {
+            max_tokens: 5.0,
+            refill_rate: 0.0,
+        };
+        // Unique key so repeated runs against the same DB start from a full bucket.
+        let key = format!("pg-conc-{}", crate::job::now_millis());
+
+        let acquired = Arc::new(AtomicUsize::new(0));
+        let barrier = Arc::new(Barrier::new(num_threads));
+        let mut handles = vec![];
+        for _ in 0..num_threads {
+            let limiter = limiter.clone();
+            let config = config.clone();
+            let acquired = acquired.clone();
+            let barrier = barrier.clone();
+            let key = key.clone();
+            handles.push(std::thread::spawn(move || {
+                barrier.wait();
+                let mut last_err = None;
+                for _ in 0..10 {
+                    match limiter.try_acquire(&key, &config) {
+                        Ok(true) => {
+                            acquired.fetch_add(1, Ordering::Relaxed);
+                            return;
+                        }
+                        Ok(false) => return,
+                        Err(e) => {
+                            last_err = Some(e.to_string());
+                            std::thread::sleep(std::time::Duration::from_millis(10));
+                            continue;
+                        }
+                    }
+                }
+                panic!("exhausted retries in Postgres rate-limit test: {last_err:?}");
+            }));
+        }
+        for h in handles {
+            h.join().unwrap();
+        }
+
+        assert_eq!(
+            acquired.load(Ordering::Relaxed),
+            5,
+            "FOR UPDATE + seed must admit exactly max_tokens, never over-admit"
+        );
+    }
 }

crates/taskito-core/src/storage/diesel_common/dead_letter.rs

@@ -18,15 +18,15 @@ macro_rules! impl_diesel_dead_letter_ops {
 
                 // Write-priority transaction: this reads the job row then writes
                 // it to `archived_jobs`, which would deadlock under SQLite's
-                // deferred lock-upgrade. See `archive_transaction`.
+                // deferred lock-upgrade. See `write_transaction`.
                 let dlq_retry_count = job
                     .metadata
                     .as_deref()
                     .and_then(|m| serde_json::from_str::<serde_json::Value>(m).ok())
                     .and_then(|v| v.get("__dlq_retry_count")?.as_i64())
                     .unwrap_or(0) as i32;
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let dlq_row = NewDeadLetterRow {
                         id: &dlq_id,
                         original_job_id: &job.id,

crates/taskito-core/src/storage/diesel_common/jobs.rs

@@ -164,9 +164,8 @@ macro_rules! impl_diesel_job_ops {
             pub fn enqueue(&self, new_job: NewJob) -> Result<Job> {
                 let depends_on = new_job.depends_on.clone();
                 let job = new_job.into_job();
-                let mut conn = self.conn()?;
 
-                conn.transaction(|conn| {
+                self.write_transaction(|conn| {
                     // Validate dependencies exist and aren't dead/cancelled.
                     // Terminal deps live in `archived_jobs`, so a missing live
                     // row falls back to the archive.
@@ -226,10 +225,12 @@ macro_rules! impl_diesel_job_ops {
                     Ok(())
                 })
                 .map_err(|e| match e {
-                    diesel::result::Error::RollbackTransaction => QueueError::DependencyNotFound(
-                        "dependency not found or already dead/cancelled".to_string(),
-                    ),
-                    other => QueueError::Storage(other),
+                    QueueError::Storage(diesel::result::Error::RollbackTransaction) => {
+                        QueueError::DependencyNotFound(
+                            "dependency not found or already dead/cancelled".to_string(),
+                        )
+                    }
+                    other => other,
                 })?;
 
                 Ok(job)
@@ -243,10 +244,9 @@ macro_rules! impl_diesel_job_ops {
                 // chunk size keeps the macro-generated code identical.
                 const BATCH_INSERT_CHUNK: usize = 50;
 
-                let mut conn = self.conn()?;
                 let jobs: Vec<Job> = new_jobs.into_iter().map(|nj| nj.into_job()).collect();
 
-                conn.transaction(|conn| {
+                self.write_transaction(|conn| {
                     let rows: Vec<NewJobRow> = jobs
                         .iter()
                         .map(|job| NewJobRow {
@@ -304,7 +304,6 @@ macro_rules! impl_diesel_job_ops {
             pub fn enqueue_unique(&self, new_job: NewJob) -> Result<Job> {
                 let depends_on = new_job.depends_on.clone();
                 let job = new_job.into_job();
-                let mut conn = self.conn()?;
 
                 // A UniqueViolation means a concurrent insert won the unique slot.
                 // If that job is still active we return it; if it has since gone
@@ -313,7 +312,7 @@ macro_rules! impl_diesel_job_ops {
                 // an error instead of a phantom job that was never persisted.
                 const MAX_ENQUEUE_ATTEMPTS: usize = 3;
                 for _ in 0..MAX_ENQUEUE_ATTEMPTS {
-                    let result = conn.transaction(|conn| {
+                    let result = self.write_transaction(|conn| {
                         // Return any existing active job with the same unique_key.
                         if let Some(ref uk) = job.unique_key {
                             let existing: Option<JobRow> =
@@ -387,13 +386,17 @@ macro_rules! impl_diesel_job_ops {
 
                     match result {
                         Ok(j) => return Ok(j),
-                        Err(diesel::result::Error::DatabaseError(
+                        Err(QueueError::Storage(diesel::result::Error::DatabaseError(
                             diesel::result::DatabaseErrorKind::UniqueViolation,
                             _,
-                        )) => {
+                        ))) => {
                             // Concurrent winner: return it if still active, else the
                             // slot was freed by a terminal transition — retry insert.
+                            // Acquire the connection here (not before the loop) so
+                            // it never overlaps `write_transaction`'s own checkout —
+                            // the in-memory test pool has a single connection.
                             if let Some(ref uk) = job.unique_key {
+                                let mut conn = self.conn()?;
                                 let existing: Option<JobRow> = jobs::table
                                     .filter(jobs::unique_key.eq(uk))
                                     .filter(jobs::status.eq_any([
@@ -409,12 +412,12 @@ macro_rules! impl_diesel_job_ops {
                             }
                             continue;
                         }
-                        Err(diesel::result::Error::RollbackTransaction) => {
+                        Err(QueueError::Storage(diesel::result::Error::RollbackTransaction)) => {
                             return Err(QueueError::DependencyNotFound(
                                 "dependency not found or already dead/cancelled".to_string(),
                             ));
                         }
-                        Err(e) => return Err(QueueError::Storage(e)),
+                        Err(e) => return Err(e),
                     }
                 }
 
@@ -432,9 +435,7 @@ macro_rules! impl_diesel_job_ops {
                 now: i64,
                 namespace: Option<&str>,
             ) -> Result<Option<Job>> {
-                let mut conn = self.conn()?;
-
-                conn.transaction(|conn| {
+                self.write_transaction(|conn| {
                     let mut query = jobs::table
                         .filter(jobs::queue.eq(queue_name))
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
@@ -559,9 +560,7 @@ macro_rules! impl_diesel_job_ops {
                 // to keep the loaded set small.
                 let scan_limit = (max.saturating_mul(4)).min(400) as i64;
 
-                let mut conn = self.conn()?;
-
-                conn.transaction(|conn| {
+                self.write_transaction(|conn| {
                     let mut query = jobs::table
                         .filter(jobs::queue.eq(queue_name))
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
@@ -661,7 +660,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn complete(&self, id: &str, result_bytes: Option<Vec<u8>>) -> Result<()> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let mut row: JobRow = match jobs::table
                         .find(id)
                         .filter(jobs::status.eq(JobStatus::Running as i32))
@@ -686,7 +685,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn fail(&self, id: &str, error: &str) -> Result<()> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let mut row: JobRow = match jobs::table
                         .find(id)
                         .filter(jobs::status.eq(JobStatus::Running as i32))
@@ -758,7 +757,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn cancel_job(&self, id: &str) -> Result<bool> {
                 let now = now_millis();
 
-                let archived = self.archive_transaction(|conn| {
+                let archived = self.write_transaction(|conn| {
                     let mut row: JobRow = match jobs::table
                         .find(id)
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
@@ -814,7 +813,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn mark_cancelled(&self, id: &str) -> Result<()> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let mut row: JobRow = match jobs::table
                         .find(id)
                         .select(JobRow::as_select())
@@ -870,7 +869,7 @@ macro_rules! impl_diesel_job_ops {
 
                 if !queue.is_empty() {
                     let error_msg = format!("{reason}: {failed_job_id}");
-                    self.archive_transaction(|conn| {
+                    self.write_transaction(|conn| {
                         let rows: Vec<JobRow> = jobs::table
                             .filter(jobs::id.eq_any(&queue))
                             .filter(jobs::status.eq(JobStatus::Pending as i32))
@@ -1283,7 +1282,7 @@ macro_rules! impl_diesel_job_ops {
             /// Purge completed jobs older than the given timestamp. Terminal
             /// jobs live in `archived_jobs`, so the purge targets that table.
             pub fn purge_completed(&self, older_than_ms: i64) -> Result<u64> {
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let job_ids: Vec<String> = archived_jobs::table
                         .filter(archived_jobs::status.eq(JobStatus::Complete as i32))
                         .filter(archived_jobs::completed_at.lt(older_than_ms))
@@ -1306,7 +1305,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn purge_completed_with_ttl(&self, global_cutoff_ms: i64) -> Result<u64> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let global_ids: Vec<String> = archived_jobs::table
                         .filter(archived_jobs::status.eq(JobStatus::Complete as i32))
                         .filter(archived_jobs::result_ttl_ms.is_null())
@@ -1418,7 +1417,7 @@ macro_rules! impl_diesel_job_ops {
 
             /// Expire pending jobs that have passed their expires_at.
             pub fn expire_pending_jobs(&self, now: i64) -> Result<u64> {
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let rows: Vec<JobRow> = jobs::table
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
                         .filter(jobs::expires_at.is_not_null())
@@ -1434,7 +1433,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn cancel_pending_by_queue(&self, queue: &str) -> Result<u64> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let rows: Vec<JobRow> = jobs::table
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
                         .filter(jobs::queue.eq(queue))
@@ -1449,7 +1448,7 @@ macro_rules! impl_diesel_job_ops {
             pub fn cancel_pending_by_task(&self, task_name: &str) -> Result<u64> {
                 let now = now_millis();
 
-                self.archive_transaction(|conn| {
+                self.write_transaction(|conn| {
                     let rows: Vec<JobRow> = jobs::table
                         .filter(jobs::status.eq(JobStatus::Pending as i32))
                         .filter(jobs::task_name.eq(task_name))

crates/taskito-core/src/storage/models.rs

@@ -156,7 +156,17 @@ pub struct NewDeadLetterRow<'a> {
 }
 
 /// A row in the `rate_limits` table.
-#[derive(Queryable, Selectable, Insertable, AsChangeset, Debug, Clone, Serialize, Deserialize)]
+#[derive(
+    Queryable,
+    QueryableByName,
+    Selectable,
+    Insertable,
+    AsChangeset,
+    Debug,
+    Clone,
+    Serialize,
+    Deserialize,
+)]
 #[diesel(table_name = rate_limits)]
 pub struct RateLimitRow {
     pub key: String,

crates/taskito-core/src/storage/postgres/jobs.rs

@@ -16,9 +16,9 @@ crate::storage::diesel_common::impl_diesel_job_ops!(PostgresStorage, PgConnectio
 impl PostgresStorage {
     /// Run a read-then-write unit of work in a transaction. Postgres serializes
     /// row-level writes without the SQLite lock-upgrade deadlock, so a regular
-    /// transaction suffices; this mirrors [`SqliteStorage::archive_transaction`]
+    /// transaction suffices; this mirrors [`SqliteStorage::write_transaction`]
     /// so the shared job-ops macro can call one name on both backends.
-    pub(crate) fn archive_transaction<T, F>(&self, f: F) -> Result<T>
+    pub(crate) fn write_transaction<T, F>(&self, f: F) -> Result<T>
     where
         F: FnOnce(&mut PgConnection) -> std::result::Result<T, QueueError>,
     {

crates/taskito-core/src/storage/postgres/mod.rs

@@ -22,12 +22,13 @@ use diesel::r2d2::{ConnectionManager, Pool};
 use crate::error::Result;
 use crate::storage::diesel_common::migrations as common_migrations;
 
-/// Run an ALTER TABLE migration, logging a warning on any failure.
-/// Postgres migrations use IF NOT EXISTS, so any error is genuinely unexpected.
-fn migration_alter(conn: &mut PgConnection, sql: &str) {
-    if let Err(e) = diesel::sql_query(sql).execute(conn) {
-        log::warn!("migration failed for '{sql}': {e}");
-    }
+/// Run an ALTER TABLE migration. Postgres statements use `ADD COLUMN IF NOT
+/// EXISTS`, so the already-applied case succeeds silently and any error is a
+/// genuine failure (locked table, permissions, disk) that must be propagated
+/// rather than leaving the schema missing a column.
+fn migration_alter(conn: &mut PgConnection, sql: &str) -> Result<()> {
+    diesel::sql_query(sql).execute(conn)?;
+    Ok(())
 }
 
 type PgPool = Pool<ConnectionManager<PgConnection>>;
@@ -151,7 +152,7 @@ impl PostgresStorage {
             diesel::sql_query(*sql).execute(&mut conn)?;
         }
         for sql in common_migrations::alter_statements(&common_migrations::POSTGRES) {
-            migration_alter(&mut conn, &sql);
+            migration_alter(&mut conn, &sql)?;
         }
         // Data backfills must fail loudly — a swallowed failure would leave
         // has_deps wrong and let dequeue bypass dependency enforcement.