perf(scaling): tier 3 async backpressure (S16, S17, S19)

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

@@ -146,6 +146,55 @@ pub struct TaskConfig {
     pub rate_limit: Option<crate::resilience::rate_limiter::RateLimitConfig>,
     pub circuit_breaker: Option<CircuitBreakerConfig>,
     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
+    /// `max_concurrent`, which is the cluster-wide cap and costs a database read;
+    /// this one is in-process and free. `None` lets the task use the whole pool.
+    pub max_in_flight_per_task: Option<usize>,
+}
+
+/// In-flight bookkeeping behind the dispatch caps: which jobs are out, and how many
+/// of each task. Per-task counts are maintained alongside rather than derived so the
+/// per-task gate stays O(1) — the poller consults it on every dispatch, and the map
+/// is as large as `max_in_flight`.
+#[derive(Default)]
+struct InFlight {
+    task_by_job: HashMap<String, String>,
+    count_by_task: HashMap<String, usize>,
+}
+
+impl InFlight {
+    fn len(&self) -> usize {
+        self.task_by_job.len()
+    }
+
+    fn count_for(&self, task_name: &str) -> usize {
+        self.count_by_task.get(task_name).copied().unwrap_or(0)
+    }
+
+    fn insert(&mut self, job_id: &str, task_name: &str) {
+        if self
+            .task_by_job
+            .insert(job_id.to_string(), task_name.to_string())
+            .is_none()
+        {
+            *self.count_by_task.entry(task_name.to_string()).or_insert(0) += 1;
+        }
+    }
+
+    /// `true` if the job was tracked here (so a slot was actually freed).
+    fn remove(&mut self, job_id: &str) -> bool {
+        let Some(task_name) = self.task_by_job.remove(job_id) else {
+            return false;
+        };
+        if let Some(count) = self.count_by_task.get_mut(&task_name) {
+            *count -= 1;
+            if *count == 0 {
+                self.count_by_task.remove(&task_name);
+            }
+        }
+        true
+    }
 }
 
 /// Per-queue configuration for rate limiting and concurrency caps.
@@ -172,11 +221,11 @@ pub struct Scheduler {
     /// binding sets it to the process `worker_id` so dead-worker recovery can
     /// attribute orphaned claims. See [`Self::set_claim_owner`].
     claim_owner: String,
-    /// Ids of jobs dispatched to the worker channel but not yet finished, used
-    /// only when `config.max_in_flight` is set. Its length is the current
-    /// in-flight count; the poller stops dispatching once it hits the cap, and
-    /// [`Self::handle_result`] removes an id as each job finishes.
-    in_flight: Mutex<HashSet<String>>,
+    /// Jobs dispatched to the worker channel but not yet finished, used only when
+    /// `config.max_in_flight` is set. Its length is the current in-flight count;
+    /// the poller stops dispatching once it hits the cap, and [`Self::handle_result`]
+    /// removes a job as it finishes.
+    in_flight: Mutex<InFlight>,
     /// Signalled when a finished job frees an in-flight slot, so the poll loop
     /// can refill immediately instead of waiting out its backoff — keeping
     /// throughput up despite the in-flight cap.
@@ -225,7 +274,7 @@ impl Scheduler {
             paused_cache: Mutex::new((HashSet::new(), Instant::now())),
             namespace,
             claim_owner: poller::SCHEDULER_CLAIM_OWNER.to_string(),
-            in_flight: Mutex::new(HashSet::new()),
+            in_flight: Mutex::new(InFlight::default()),
             dispatch_wake: Arc::new(Notify::new()),
             #[cfg(feature = "push-dispatch")]
             wake_source: Mutex::new(None),
@@ -261,15 +310,23 @@ impl Scheduler {
 
     /// Record a job as in flight once it is handed to the worker channel.
     /// No-op when dispatch is unbounded.
-    fn track_in_flight(&self, job_id: &str) {
+    fn track_in_flight(&self, job_id: &str, task_name: &str) {
         if self.config.max_in_flight.is_some() {
             self.in_flight
                 .lock()
                 .unwrap_or_else(|p| p.into_inner())
-                .insert(job_id.to_string());
+                .insert(job_id, task_name);
         }
     }
 
+    /// How many of this task's jobs this worker currently has in flight.
+    fn task_in_flight(&self, task_name: &str) -> usize {
+        self.in_flight
+            .lock()
+            .unwrap_or_else(|p| p.into_inner())
+            .count_for(task_name)
+    }
+
     /// 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.
@@ -703,6 +760,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: None,
+                max_in_flight_per_task: None,
             },
         );
 
@@ -743,6 +801,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: None,
+                max_in_flight_per_task: None,
             },
         );
 
@@ -844,6 +903,7 @@ mod tests {
                 }),
                 circuit_breaker: None,
                 max_concurrent: None,
+                max_in_flight_per_task: None,
             },
         );
 
@@ -888,6 +948,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: Some(cb_config),
                 max_concurrent: None,
+                max_in_flight_per_task: None,
             },
         );
 
@@ -926,6 +987,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: Some(2),
+                max_in_flight_per_task: None,
             },
         );
 
@@ -994,6 +1056,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: Some(2),
+                max_in_flight_per_task: None,
             },
         );
 
@@ -1078,6 +1141,129 @@ mod tests {
         );
     }
 
+    /// Scheduler with an in-flight cap and a per-task bulkhead on `task_name`.
+    fn bulkhead_scheduler(max_in_flight: usize, task_name: &str, per_task: usize) -> Scheduler {
+        let storage =
+            StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap());
+        let config = SchedulerConfig {
+            max_in_flight: Some(max_in_flight),
+            ..SchedulerConfig::default()
+        };
+        let mut scheduler = Scheduler::new(storage, vec!["default".to_string()], config, None);
+        scheduler.register_task(
+            task_name.to_string(),
+            TaskConfig {
+                retry_policy: RetryPolicy::default(),
+                rate_limit: None,
+                circuit_breaker: None,
+                max_concurrent: None,
+                max_in_flight_per_task: Some(per_task),
+            },
+        );
+        scheduler
+    }
+
+    #[test]
+    fn test_per_task_in_flight_cap_rejects_and_reschedules() {
+        // S19: a task capped at one in-flight job must not take a second slot
+        // even though the pool has room. The surplus rolls back to Pending with
+        // its claim cleared, rather than stranding Running.
+        let scheduler = bulkhead_scheduler(8, "hog", 1);
+        for _ in 0..3 {
+            scheduler.storage.enqueue(make_job("hog")).unwrap();
+        }
+
+        let (tx, mut rx) = make_channel(16);
+        while scheduler.tick_dispatch(&tx) {}
+
+        let mut dispatched = 0;
+        while rx.try_recv().is_ok() {
+            dispatched += 1;
+        }
+        assert_eq!(dispatched, 1, "per-task bulkhead limits dispatch to one");
+
+        let pending = scheduler
+            .storage
+            .list_jobs(Some(JobStatus::Pending as i32), None, None, 10, 0, None)
+            .unwrap();
+        assert_eq!(pending.len(), 2, "over-cap jobs must return to Pending");
+
+        let claims = scheduler
+            .storage
+            .list_claims_by_worker("scheduler")
+            .unwrap();
+        assert_eq!(
+            claims.len(),
+            1,
+            "rejected jobs must not keep a stale claim row"
+        );
+    }
+
+    #[test]
+    fn test_per_task_cap_does_not_block_other_tasks() {
+        // S19's whole point: the bulkhead is scoped to the task that tripped it.
+        // A saturated `hog` must not consume slots that `free` could use — the
+        // starvation a global semaphore (or a blocking per-task acquire on the
+        // dispatch loop) would cause.
+        let scheduler = bulkhead_scheduler(8, "hog", 1);
+        let (tx, mut rx) = make_channel(16);
+
+        scheduler.storage.enqueue(make_job("hog")).unwrap();
+        while scheduler.tick_dispatch(&tx) {}
+        assert_eq!(scheduler.task_in_flight("hog"), 1, "hog is at its bulkhead");
+
+        // With hog saturated, an unrelated task still dispatches up to the pool cap.
+        for _ in 0..3 {
+            scheduler.storage.enqueue(make_job("free")).unwrap();
+        }
+        while scheduler.tick_dispatch(&tx) {}
+
+        let mut hog = 0;
+        let mut free = 0;
+        while let Ok(job) = rx.try_recv() {
+            match job.task_name.as_str() {
+                "hog" => hog += 1,
+                "free" => free += 1,
+                other => panic!("unexpected task {other}"),
+            }
+        }
+        assert_eq!(hog, 1, "capped task is held to its bulkhead");
+        assert_eq!(free, 3, "uncapped task dispatches past the saturated one");
+    }
+
+    #[test]
+    fn test_release_in_flight_frees_per_task_slot() {
+        // The per-task count must decrement as jobs finish, or the bulkhead would
+        // leak slots and wedge the task at its cap forever.
+        let scheduler = bulkhead_scheduler(8, "hog", 1);
+        let (tx, mut rx) = make_channel(16);
+
+        scheduler.storage.enqueue(make_job("hog")).unwrap();
+        while scheduler.tick_dispatch(&tx) {}
+        let first = rx.try_recv().expect("one job dispatched");
+        assert_eq!(scheduler.task_in_flight("hog"), 1);
+
+        scheduler
+            .handle_result(JobResult::Success {
+                job_id: first.id.clone(),
+                result: None,
+                task_name: "hog".to_string(),
+                wall_time_ns: 1,
+            })
+            .unwrap();
+        assert_eq!(
+            scheduler.task_in_flight("hog"),
+            0,
+            "slot released on finish"
+        );
+
+        // The freed slot admits the next job of the same task.
+        scheduler.storage.enqueue(make_job("hog")).unwrap();
+        while scheduler.tick_dispatch(&tx) {}
+        assert!(rx.try_recv().is_ok(), "freed slot admits the next job");
+        assert_eq!(scheduler.task_in_flight("hog"), 1);
+    }
+
     #[test]
     fn test_try_dispatch_per_task_max_one_dispatches_one() {
         // Regression: `max_concurrent = 1` must allow exactly one job to
@@ -1092,6 +1278,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: Some(1),
+                max_in_flight_per_task: None,
             },
         );
 
@@ -1256,6 +1443,7 @@ mod tests {
                 rate_limit: None,
                 circuit_breaker: None,
                 max_concurrent: None,
+                max_in_flight_per_task: None,
             },
         );
 
@@ -1292,6 +1480,7 @@ mod tests {
             rate_limit: None,
             circuit_breaker: None,
             max_concurrent: None,
+            max_in_flight_per_task: None,
         }
     }
 

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

@@ -290,9 +290,10 @@ impl Scheduler {
         // the stale-reaper times it out, which surfaces as a *timeout* in
         // metrics and middleware (wrong outcome for a job that never ran).
         let job_id = job.id.clone();
+        let task_name = job.task_name.clone();
         match job_tx.try_send(job) {
             Ok(()) => {
-                self.track_in_flight(&job_id);
+                self.track_in_flight(&job_id, &task_name);
                 Ok(true)
             }
             Err(TrySendError::Full(job)) => {
@@ -417,6 +418,16 @@ impl Scheduler {
                     return Ok(false);
                 }
             }
+
+            // In-process bulkhead: keep one task from taking every slot in this
+            // worker's pool. Unlike the running-count above, the in-flight map does
+            // not yet include this job — it's only tracked once the send succeeds —
+            // so this compares with `>=` where that one uses a strict `>`.
+            if let Some(max_in_flight) = config.max_in_flight_per_task {
+                if self.task_in_flight(&job.task_name) >= max_in_flight {
+                    return Ok(false);
+                }
+            }
         }
 
         Ok(true)

crates/taskito-java/src/worker.rs

@@ -278,6 +278,8 @@ fn register_task_policies(scheduler: &mut Scheduler, configs: Option<Vec<TaskRet
                 rate_limit: None,
                 circuit_breaker,
                 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

@@ -42,6 +42,8 @@ pub fn task_config(input: &TaskConfigInput) -> Result<TaskConfig> {
             .map(circuit_breaker_config)
             .transpose()?,
         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/lib.rs

@@ -36,7 +36,10 @@ fn _taskito(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyJob>()?;
     m.add_class::<PyTaskConfig>()?;
     #[cfg(feature = "native-async")]
-    m.add_class::<native_async::PyResultSender>()?;
+    {
+        m.add_class::<native_async::PyResultSender>()?;
+        m.add_class::<native_async::PyJobPermit>()?;
+    }
     #[cfg(feature = "workflows")]
     {
         m.add_class::<py_workflow::PyWorkflowBuilder>()?;

crates/taskito-python/src/native_async/mod.rs

@@ -4,9 +4,11 @@
 //! because every line is Python-coupled (GIL, cloudpickle, the Python async
 //! executor). Compiled only under the `native-async` feature.
 
+mod permit;
 mod pool;
 mod result_sender;
 mod task_executor;
 
+pub use permit::PyJobPermit;
 pub use pool::NativeAsyncPool;
 pub use result_sender::PyResultSender;