diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 472958c6..64d565e5 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -565,6 +565,46 @@ mod tests { assert_eq!(retried.retry_count, 1); } + #[test] + fn test_explicit_zero_max_retries_no_retry() { + // The task policy would retry 3×, but an explicit per-job max_retries=0 + // (at-most-once) must override it and go straight to the DLQ. + let mut scheduler = test_scheduler(); + scheduler.register_task( + "noretry_task".to_string(), + TaskConfig { + retry_policy: RetryPolicy { + max_retries: 3, + base_delay_ms: 100, + max_delay_ms: 1000, + custom_delays_ms: None, + }, + rate_limit: None, + circuit_breaker: None, + max_concurrent: None, + }, + ); + + let job = enqueue_and_run(&scheduler, "noretry_task"); + + let outcome = scheduler + .handle_result(JobResult::Failure { + job_id: job.id.clone(), + error: "boom".to_string(), + retry_count: 0, + max_retries: 0, + task_name: "noretry_task".to_string(), + wall_time_ns: 500_000, + should_retry: true, + timed_out: false, + }) + .unwrap(); + + assert!(matches!(outcome, ResultOutcome::DeadLettered { .. })); + let dead = scheduler.storage.list_dead(10, 0).unwrap(); + assert!(dead.iter().any(|d| d.original_job_id == job.id)); + } + #[test] fn test_handle_failure_exhausted() { let scheduler = test_scheduler(); diff --git a/crates/taskito-core/src/scheduler/result_handler.rs b/crates/taskito-core/src/scheduler/result_handler.rs index 3d077a3d..62bb9c78 100644 --- a/crates/taskito-core/src/scheduler/result_handler.rs +++ b/crates/taskito-core/src/scheduler/result_handler.rs @@ -104,11 +104,11 @@ impl Scheduler { .map(|c| c.retry_policy.clone()) .unwrap_or_default(); - let effective_max = if max_retries > 0 { - max_retries - } else { - policy.max_retries - }; + // job.max_retries is the budget resolved at enqueue (the queue + // default or the caller's explicit value), so honor it exactly — + // treating a stored 0 as "unset" would silently re-run a job the + // caller marked at-most-once. policy is still used for backoff. + let effective_max = max_retries; if retry_count < effective_max { let next_at = policy.next_retry_at(retry_count); diff --git a/crates/taskito-core/src/storage/diesel_common/jobs.rs b/crates/taskito-core/src/storage/diesel_common/jobs.rs index 40f9f2fe..16f161ce 100644 --- a/crates/taskito-core/src/storage/diesel_common/jobs.rs +++ b/crates/taskito-core/src/storage/diesel_common/jobs.rs @@ -298,90 +298,103 @@ macro_rules! impl_diesel_job_ops { }) } - /// Enqueue with unique_key deduplication. Returns existing job if duplicate. + /// Enqueue with unique_key deduplication. Returns the existing active + /// job on a duplicate, validates dependencies exactly like `enqueue`, + /// and never returns a job whose insert was rolled back. pub fn enqueue_unique(&self, new_job: NewJob) -> Result { let depends_on = new_job.depends_on.clone(); let job = new_job.into_job(); let mut conn = self.conn()?; - let result = conn.transaction(|conn| { - // Check for existing active job with same unique_key - if let Some(ref uk) = job.unique_key { - let existing: Option = jobs::table - .filter(jobs::unique_key.eq(uk)) - .filter( - jobs::status - .eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), - ) - .select(JobRow::as_select()) - .first(conn) - .optional()?; - - if let Some(row) = existing { - return Ok(Job::from(row)); + // A UniqueViolation means a concurrent insert won the unique slot. + // If that job is still active we return it; if it has since gone + // terminal — freeing the partial unique index — we retry the + // insert. Bound the attempts so persistent contention surfaces as + // 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| { + // Return any existing active job with the same unique_key. + if let Some(ref uk) = job.unique_key { + let existing: Option = + jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter(jobs::status.eq_any([ + JobStatus::Pending as i32, + JobStatus::Running as i32, + ])) + .select(JobRow::as_select()) + .first(conn) + .optional()?; + if let Some(row) = existing { + return Ok(Job::from(row)); + } } - } - - let row = 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(), - has_deps: job.has_deps, - }; - diesel::insert_into(jobs::table) - .values(&row) - .execute(conn)?; + // Validate dependencies exist and aren't dead/cancelled, + // matching `enqueue` (RollbackTransaction → DependencyNotFound). + for dep_id in &depends_on { + Self::validate_dependency(conn, dep_id)?; + } - diesel::insert_into(job_payloads::table) - .values(&NewJobPayloadRow { - job_id: &job.id, + let row = NewJobRow { + id: &job.id, + queue: &job.queue, + task_name: &job.task_name, payload: &job.payload, - result: None, - }) - .execute(conn)?; - - // Insert dependency rows - for dep_id in &depends_on { - let dep_row = NewJobDependencyRow { - id: &uuid::Uuid::now_v7().to_string(), - job_id: &job.id, - depends_on_job_id: dep_id, + 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(), + has_deps: job.has_deps, }; - diesel::insert_into(job_dependencies::table) - .values(&dep_row) + + diesel::insert_into(jobs::table) + .values(&row) .execute(conn)?; - } - Ok(job.clone()) - }); + diesel::insert_into(job_payloads::table) + .values(&NewJobPayloadRow { + job_id: &job.id, + payload: &job.payload, + result: None, + }) + .execute(conn)?; - // Handle unique constraint violation by returning existing job - match result { - Ok(j) => Ok(j), - Err(diesel::result::Error::DatabaseError( - diesel::result::DatabaseErrorKind::UniqueViolation, - _, - )) => { - if let Some(ref uk) = job.unique_key { - let mut conn = self.conn()?; - let existing: Option = - jobs::table + for dep_id in &depends_on { + let dep_row = NewJobDependencyRow { + id: &uuid::Uuid::now_v7().to_string(), + job_id: &job.id, + depends_on_job_id: dep_id, + }; + diesel::insert_into(job_dependencies::table) + .values(&dep_row) + .execute(conn)?; + } + + Ok(job.clone()) + }); + + match result { + Ok(j) => return Ok(j), + Err(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. + if let Some(ref uk) = job.unique_key { + let existing: Option = jobs::table .filter(jobs::unique_key.eq(uk)) .filter(jobs::status.eq_any([ JobStatus::Pending as i32, @@ -390,14 +403,24 @@ macro_rules! impl_diesel_job_ops { .select(JobRow::as_select()) .first(&mut conn) .optional()?; - if let Some(row) = existing { - return Ok(Job::from(row)); + if let Some(row) = existing { + return Ok(Job::from(row)); + } } + continue; } - Ok(job) + Err(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) => Err(QueueError::Storage(e)), } + + Err(QueueError::Other( + "enqueue_unique: unique key contended across retries".to_string(), + )) } /// Atomically dequeue the highest-priority ready job from the given queue. diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index 9be18451..2b4ce633 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -464,6 +464,64 @@ fn test_unique_key_dedup() { assert_eq!(j1.id, j2.id); } +#[test] +fn test_enqueue_unique_rejects_missing_dependency() { + // enqueue_unique must validate dependencies like enqueue (it previously + // inserted dep rows blind, treating a bogus dep as satisfied). + let storage = test_storage(); + let mut job = make_job("unique_orphan"); + job.unique_key = Some("uk-missing-dep".to_string()); + job.depends_on = vec!["nonexistent-id".to_string()]; + assert!(matches!( + storage.enqueue_unique(job), + Err(QueueError::DependencyNotFound(_)) + )); +} + +#[test] +fn test_enqueue_unique_rejects_dead_dependency() { + let storage = test_storage(); + // A cancelled (archived, non-Complete) dependency must be rejected. + let dep = storage.enqueue(make_job("dep_to_cancel")).unwrap(); + assert!(storage.cancel_job(&dep.id).unwrap()); + + let mut job = make_job("unique_blocked"); + job.unique_key = Some("uk-dead-dep".to_string()); + job.depends_on = vec![dep.id.clone()]; + assert!(matches!( + storage.enqueue_unique(job), + Err(QueueError::DependencyNotFound(_)) + )); +} + +#[test] +fn test_enqueue_unique_after_dup_completes() { + // Once the prior holder of a unique_key completes (archived, freeing the + // partial index), a fresh enqueue_unique must return a real, persisted job — + // never the phantom (rolled-back) job the old fallback could return. + let storage = test_storage(); + let mut a = make_job("unique_reuse"); + a.unique_key = Some("uk-reuse".to_string()); + let a = storage.enqueue_unique(a).unwrap(); + storage + .dequeue("default", now_millis() + 1000, None) + .unwrap(); + storage.complete(&a.id, None).unwrap(); + + let mut b = make_job("unique_reuse"); + b.unique_key = Some("uk-reuse".to_string()); + let b = storage.enqueue_unique(b).unwrap(); + + assert_ne!( + a.id, b.id, + "freed unique key must yield a new job, not dedup to A" + ); + assert!( + storage.get_job(&b.id).unwrap().is_some(), + "returned job must actually be persisted (no phantom)" + ); +} + #[test] fn test_enqueue_batch() { let storage = test_storage(); diff --git a/crates/taskito-core/tests/rust/storage_tests.rs b/crates/taskito-core/tests/rust/storage_tests.rs index 27e8d72f..6b101007 100644 --- a/crates/taskito-core/tests/rust/storage_tests.rs +++ b/crates/taskito-core/tests/rust/storage_tests.rs @@ -210,6 +210,18 @@ fn test_unique_key_dedup(s: &impl Storage) { assert_eq!(j1.id, j2.id); } +fn test_enqueue_unique_validates_deps(s: &impl Storage) { + // enqueue_unique must reject a missing dependency on every backend, matching + // enqueue (Redis already validated; the Diesel backends did not). + let mut job = make_job("q-unique-deps", "unique_dep_task"); + job.unique_key = Some("unique-dep-key".to_string()); + job.depends_on = vec!["nonexistent-dep".to_string()]; + assert!(matches!( + s.enqueue_unique(job), + Err(taskito_core::error::QueueError::DependencyNotFound(_)) + )); +} + fn test_enqueue_batch(s: &impl Storage) { let jobs: Vec = (0..5) .map(|i| { @@ -603,6 +615,7 @@ fn run_storage_tests(s: &impl Storage) { test_stats(s); test_stats_by_queue_and_task(s); test_unique_key_dedup(s); + test_enqueue_unique_validates_deps(s); test_enqueue_batch(s); test_dead_letter_queue(s); test_delete_dead(s); diff --git a/pyproject.toml b/pyproject.toml index e8387138..d715365d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ Homepage = "https://github.com/ByteVeda/taskito" Documentation = "https://docs.byteveda.org/taskito" Repository = "https://github.com/ByteVeda/taskito" -Changelog = "https://docs.byteveda.org/taskito/docs/more/changelog" +Changelog = "https://docs.byteveda.org/taskito/more/changelog" Issues = "https://github.com/ByteVeda/taskito/issues" [project.optional-dependencies] diff --git a/tests/core/test_dlq.py b/tests/core/test_dlq.py index f02697d3..86abb98f 100644 --- a/tests/core/test_dlq.py +++ b/tests/core/test_dlq.py @@ -67,3 +67,26 @@ def test_dlq_retry_count_exposed(queue: Queue, poll_until: PollUntil) -> None: dead = queue.dead_letters() assert dead[0]["dlq_retry_count"] == 0 + + +def test_explicit_zero_max_retries_runs_once(queue: Queue, poll_until: PollUntil) -> None: + """A per-call max_retries=0 overrides the task's default and runs at most once.""" + call_count = 0 + + @queue.task(max_retries=3, retry_backoff=0.1) + def at_most_once() -> None: + nonlocal call_count + call_count += 1 + raise RuntimeError("permanent failure") + + at_most_once.apply_async(max_retries=0) + + worker = threading.Thread(target=queue.run_worker, daemon=True) + worker.start() + + poll_until( + lambda: len(queue.dead_letters()) >= 1, + timeout=15, + message="max_retries=0 job did not reach the DLQ", + ) + assert call_count == 1, "max_retries=0 must run the task exactly once"