feat(scaling): retry budget (S18)

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

@@ -145,6 +145,11 @@ pub struct TaskConfig {
     pub retry_policy: RetryPolicy,
     pub rate_limit: Option<crate::resilience::rate_limiter::RateLimitConfig>,
     pub circuit_breaker: Option<CircuitBreakerConfig>,
+    /// Cap on how fast this task may *retry*, across every job of the task.
+    /// Nothing else catches sustained retry storms: the circuit breaker trips on
+    /// hard failure, not on aggregate retry rate, and per-job `max_retries` is a
+    /// per-job budget, not a rate. Exhausting it dead-letters instead of retrying.
+    pub retry_budget: Option<crate::resilience::rate_limiter::RateLimitConfig>,
     pub max_concurrent: Option<i32>,
     /// Cap on this task's share of *this worker's* in-flight slots, so one slow
     /// task cannot occupy the whole pool and starve every other task. Complements
@@ -327,6 +332,35 @@ impl Scheduler {
             .count_for(task_name)
     }
 
+    /// Take a token from this task's retry budget. `true` when the retry may go
+    /// ahead — including when no budget is configured, which is the default.
+    ///
+    /// Keyed per namespace so tenants sharing a database cannot spend each
+    /// other's budget. The bucket is created on first use, so nothing needs
+    /// registering. A storage error must not strand the job, so it fails open:
+    /// the per-job `max_retries` still bounds the retry.
+    fn retry_budget_allows(&self, task_name: &str) -> bool {
+        let Some(config) = self
+            .task_configs
+            .get(task_name)
+            .and_then(|c| c.retry_budget.as_ref())
+        else {
+            return true;
+        };
+        let key = format!(
+            "retry::{}::{}",
+            self.namespace.as_deref().unwrap_or(""),
+            task_name
+        );
+        match self.rate_limiter.try_acquire(&key, config) {
+            Ok(allowed) => allowed,
+            Err(e) => {
+                error!("retry budget check failed for {task_name}: {e}");
+                true
+            }
+        }
+    }
+
     /// Release a job's in-flight slot when it finishes and wake the poller to
     /// refill. Only ids this scheduler dispatched are tracked, so results for
     /// recovered or foreign jobs (e.g. from maintenance) are a harmless no-op.
@@ -759,6 +793,7 @@ mod tests {
                 },
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: None,
             },
@@ -800,6 +835,7 @@ mod tests {
                 },
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: None,
             },
@@ -902,6 +938,7 @@ mod tests {
                     refill_rate: 0.0,
                 }),
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: None,
             },
@@ -947,6 +984,7 @@ mod tests {
                 retry_policy: RetryPolicy::default(),
                 rate_limit: None,
                 circuit_breaker: Some(cb_config),
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: None,
             },
@@ -986,6 +1024,7 @@ mod tests {
                 retry_policy: RetryPolicy::default(),
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: Some(2),
                 max_in_flight_per_task: None,
             },
@@ -1055,6 +1094,7 @@ mod tests {
                 retry_policy: RetryPolicy::default(),
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: Some(2),
                 max_in_flight_per_task: None,
             },
@@ -1156,6 +1196,7 @@ mod tests {
                 retry_policy: RetryPolicy::default(),
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: Some(per_task),
             },
@@ -1277,6 +1318,7 @@ mod tests {
                 retry_policy: RetryPolicy::default(),
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: Some(1),
                 max_in_flight_per_task: None,
             },
@@ -1442,6 +1484,7 @@ mod tests {
                 },
                 rate_limit: None,
                 circuit_breaker: None,
+                retry_budget: None,
                 max_concurrent: None,
                 max_in_flight_per_task: None,
             },
@@ -1479,6 +1522,7 @@ mod tests {
             },
             rate_limit: None,
             circuit_breaker: None,
+            retry_budget: None,
             max_concurrent: None,
             max_in_flight_per_task: None,
         }
@@ -1521,6 +1565,140 @@ mod tests {
 
     /// A scheduler must never orphan its OWN in-flight jobs (self-rescue guard),
     /// even before its first heartbeat row exists.
+    /// A task whose retries are capped at `budget` per minute.
+    fn budgeted_scheduler(task_name: &str, budget: &str) -> Scheduler {
+        let mut scheduler = test_scheduler();
+        let mut config = retry_task_config(5);
+        config.retry_budget = Some(RateLimitConfig::parse(budget).expect("valid rate"));
+        scheduler.register_task(task_name.to_string(), config);
+        scheduler
+    }
+
+    /// Fail `job_id` once and report what the scheduler decided.
+    fn fail_once(scheduler: &Scheduler, job_id: &str, task_name: &str) -> ResultOutcome {
+        scheduler
+            .handle_result(JobResult::Failure {
+                job_id: job_id.to_string(),
+                error: "boom".to_string(),
+                retry_count: 0,
+                max_retries: 5,
+                task_name: task_name.to_string(),
+                wall_time_ns: 1,
+                should_retry: true,
+                timed_out: false,
+            })
+            .unwrap()
+    }
+
+    #[test]
+    fn test_retry_budget_dead_letters_once_exhausted() {
+        // S18: the per-job max_retries is not a rate cap, so a broken dependency
+        // can retry forever across jobs. Two tokens, three failures: the third
+        // must dead-letter instead of retrying.
+        let scheduler = budgeted_scheduler("budgeted", "2/m");
+
+        let mut outcomes = Vec::new();
+        for _ in 0..3 {
+            let job = scheduler.storage.enqueue(make_job("budgeted")).unwrap();
+            outcomes.push(fail_once(&scheduler, &job.id, "budgeted"));
+        }
+
+        assert!(
+            matches!(outcomes[0], ResultOutcome::Retry { .. }),
+            "first failure is within budget"
+        );
+        assert!(
+            matches!(outcomes[1], ResultOutcome::Retry { .. }),
+            "second failure is within budget"
+        );
+        assert!(
+            matches!(outcomes[2], ResultOutcome::DeadLettered { .. }),
+            "third failure exceeds the budget and must dead-letter, got {:?}",
+            outcomes[2]
+        );
+    }
+
+    #[test]
+    fn test_retry_budget_marks_why_it_dead_lettered() {
+        // ResultOutcome can't say why, so the DLQ metadata is what distinguishes
+        // a budget kill from ordinary retry exhaustion.
+        let scheduler = budgeted_scheduler("marked", "1/m");
+        for _ in 0..2 {
+            let job = scheduler.storage.enqueue(make_job("marked")).unwrap();
+            fail_once(&scheduler, &job.id, "marked");
+        }
+
+        let dead = scheduler.storage.list_dead(10, 0).unwrap();
+        assert_eq!(dead.len(), 1);
+        assert_eq!(
+            dead[0].metadata.as_deref(),
+            Some(super::result_handler::RETRY_BUDGET_EXHAUSTED)
+        );
+    }
+
+    #[test]
+    fn test_retry_budget_is_per_task() {
+        // One task draining its budget must not stop another from retrying.
+        let mut scheduler = budgeted_scheduler("greedy", "1/m");
+        scheduler.register_task("innocent".to_string(), retry_task_config(5));
+
+        for _ in 0..2 {
+            let job = scheduler.storage.enqueue(make_job("greedy")).unwrap();
+            fail_once(&scheduler, &job.id, "greedy");
+        }
+
+        let job = scheduler.storage.enqueue(make_job("innocent")).unwrap();
+        let outcome = fail_once(&scheduler, &job.id, "innocent");
+        assert!(
+            matches!(outcome, ResultOutcome::Retry { .. }),
+            "an unbudgeted task must be unaffected, got {outcome:?}"
+        );
+    }
+
+    #[test]
+    fn test_no_retry_budget_configured_never_blocks() {
+        // The budget is opt-in; the default must keep retrying to max_retries.
+        let mut scheduler = test_scheduler();
+        scheduler.register_task("unbudgeted".to_string(), retry_task_config(5));
+        for _ in 0..5 {
+            let job = scheduler.storage.enqueue(make_job("unbudgeted")).unwrap();
+            let outcome = fail_once(&scheduler, &job.id, "unbudgeted");
+            assert!(matches!(outcome, ResultOutcome::Retry { .. }));
+        }
+    }
+
+    #[test]
+    fn test_retry_budget_not_spent_by_jobs_that_cannot_retry() {
+        // A job at its retry ceiling was never going to retry, so it must not
+        // take a token — otherwise it drains the budget for its siblings.
+        let scheduler = budgeted_scheduler("ceiling", "1/m");
+
+        let exhausted = scheduler.storage.enqueue(make_job("ceiling")).unwrap();
+        let outcome = scheduler
+            .handle_result(JobResult::Failure {
+                job_id: exhausted.id.clone(),
+                error: "boom".to_string(),
+                retry_count: 3,
+                max_retries: 3, // already at the ceiling
+                task_name: "ceiling".to_string(),
+                wall_time_ns: 1,
+                should_retry: true,
+                timed_out: false,
+            })
+            .unwrap();
+        assert!(matches!(outcome, ResultOutcome::DeadLettered { .. }));
+
+        // The single token must still be there for a job that can retry.
+        let fresh = scheduler.storage.enqueue(make_job("ceiling")).unwrap();
+        assert!(
+            matches!(
+                fail_once(&scheduler, &fresh.id, "ceiling"),
+                ResultOutcome::Retry { .. }
+            ),
+            "retry-exhausted job must not spend a budget token"
+        );
+    }
+
     #[test]
     fn test_recover_orphaned_jobs_skips_own_claims() {
         let mut scheduler = test_scheduler();

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

@@ -6,6 +6,11 @@ use crate::storage::Storage;
 
 use super::{JobResult, ResultOutcome, Scheduler};
 
+/// Dead-letter metadata marking a job the retry budget refused. `ResultOutcome`
+/// has no room to say *why* a job was dead-lettered, so this is what tells a
+/// budget kill apart from ordinary retry exhaustion when reading the DLQ.
+pub const RETRY_BUDGET_EXHAUSTED: &str = "retry_budget_exhausted";
+
 impl Scheduler {
     /// Handle a completed or failed job result from a worker.
     ///
@@ -62,19 +67,20 @@ impl Scheduler {
                 let job = self.storage.get_job(&job_id)?;
                 let queue = job.as_ref().map(|j| j.queue.clone()).unwrap_or_default();
 
-                let move_to_dlq = |job: Option<&crate::job::Job>| -> Result<()> {
-                    match job {
-                        Some(j) => self.dlq.move_to_dlq(j, &error, None),
-                        None => {
-                            warn!("job {job_id} disappeared before DLQ move");
-                            Ok(())
+                let move_to_dlq =
+                    |job: Option<&crate::job::Job>, metadata: Option<&str>| -> Result<()> {
+                        match job {
+                            Some(j) => self.dlq.move_to_dlq(j, &error, metadata),
+                            None => {
+                                warn!("job {job_id} disappeared before DLQ move");
+                                Ok(())
+                            }
                         }
-                    }
-                };
+                    };
 
                 // If should_retry is false (exception filtering), skip straight to DLQ
                 if !should_retry {
-                    move_to_dlq(job.as_ref())?;
+                    move_to_dlq(job.as_ref(), None)?;
                     return Ok(ResultOutcome::DeadLettered {
                         job_id,
                         task_name,
@@ -97,6 +103,20 @@ impl Scheduler {
                 let effective_max = max_retries;
 
                 if retry_count < effective_max {
+                    // Checked here, not before the per-job budget: a job that was
+                    // never going to retry must not spend a token, or a task at
+                    // its retry ceiling would drain the budget for its siblings.
+                    if !self.retry_budget_allows(&task_name) {
+                        warn!("retry budget exhausted for {task_name}; dead-lettering {job_id}");
+                        move_to_dlq(job.as_ref(), Some(RETRY_BUDGET_EXHAUSTED))?;
+                        return Ok(ResultOutcome::DeadLettered {
+                            job_id,
+                            task_name,
+                            queue,
+                            error,
+                            timed_out,
+                        });
+                    }
                     let next_at = policy.next_retry_at(retry_count);
                     self.storage.retry(&job_id, next_at)?;
                     #[cfg(feature = "push-dispatch")]
@@ -110,7 +130,7 @@ impl Scheduler {
                         timed_out,
                     })
                 } else {
-                    move_to_dlq(job.as_ref())?;
+                    move_to_dlq(job.as_ref(), None)?;
                     Ok(ResultOutcome::DeadLettered {
                         job_id,
                         task_name,

crates/taskito-java/src/worker.rs

@@ -277,6 +277,8 @@ fn register_task_policies(scheduler: &mut Scheduler, configs: Option<Vec<TaskRet
                 retry_policy,
                 rate_limit: None,
                 circuit_breaker,
+                // Not surfaced on this SDK yet; retries stay bounded per job.
+                retry_budget: None,
                 max_concurrent: None,
                 // Not surfaced on this SDK yet; the whole pool stays available.
                 max_in_flight_per_task: None,

crates/taskito-node/src/convert/task_config.rs

@@ -41,6 +41,8 @@ pub fn task_config(input: &TaskConfigInput) -> Result<TaskConfig> {
             .as_ref()
             .map(circuit_breaker_config)
             .transpose()?,
+        // Not surfaced on this SDK yet; retries stay bounded per job.
+        retry_budget: None,
         max_concurrent: input.max_concurrent,
         // Not surfaced on this SDK yet; the whole pool stays available.
         max_in_flight_per_task: None,

crates/taskito-python/src/py_config.rs

@@ -36,13 +36,16 @@ pub struct PyTaskConfig {
     pub circuit_breaker_half_open_success_rate: Option<f64>,
     #[pyo3(get, set)]
     pub max_in_flight_per_task: Option<i32>,
+    /// Retry-rate cap, same syntax as `rate_limit` (e.g. "100/m").
+    #[pyo3(get, set)]
+    pub retry_budget: Option<String>,
 }
 
 #[pymethods]
 #[allow(clippy::too_many_arguments)]
 impl PyTaskConfig {
     #[new]
-    #[pyo3(signature = (name, max_retries=3, retry_backoff=1.0, timeout=300, priority=0, rate_limit=None, queue="default".to_string(), circuit_breaker_threshold=None, circuit_breaker_window=None, circuit_breaker_cooldown=None, retry_delays=None, max_retry_delay=None, max_concurrent=None, circuit_breaker_half_open_probes=None, circuit_breaker_half_open_success_rate=None, max_in_flight_per_task=None))]
+    #[pyo3(signature = (name, max_retries=3, retry_backoff=1.0, timeout=300, priority=0, rate_limit=None, queue="default".to_string(), circuit_breaker_threshold=None, circuit_breaker_window=None, circuit_breaker_cooldown=None, retry_delays=None, max_retry_delay=None, max_concurrent=None, circuit_breaker_half_open_probes=None, circuit_breaker_half_open_success_rate=None, max_in_flight_per_task=None, retry_budget=None))]
     pub fn new(
         name: String,
         max_retries: i32,
@@ -60,6 +63,7 @@ impl PyTaskConfig {
         circuit_breaker_half_open_probes: Option<i32>,
         circuit_breaker_half_open_success_rate: Option<f64>,
         max_in_flight_per_task: Option<i32>,
+        retry_budget: Option<String>,
     ) -> Self {
         Self {
             name,
@@ -78,6 +82,7 @@ impl PyTaskConfig {
             circuit_breaker_half_open_probes,
             circuit_breaker_half_open_success_rate,
             max_in_flight_per_task,
+            retry_budget,
         }
     }
 }