diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 5df34012..d9355dc8 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -146,6 +146,55 @@ pub struct TaskConfig { pub rate_limit: Option, pub circuit_breaker: Option, pub max_concurrent: Option, + /// 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, +} + +/// 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, + count_by_task: HashMap, +} + +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>, + /// 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, /// 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, } } diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index 1224c55d..f42912f0 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/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) diff --git a/crates/taskito-java/src/worker.rs b/crates/taskito-java/src/worker.rs index a2570c08..1b67c709 100644 --- a/crates/taskito-java/src/worker.rs +++ b/crates/taskito-java/src/worker.rs @@ -278,6 +278,8 @@ fn register_task_policies(scheduler: &mut Scheduler, configs: Option Result { .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, }) } diff --git a/crates/taskito-python/src/lib.rs b/crates/taskito-python/src/lib.rs index b1df447a..7bec9cb1 100644 --- a/crates/taskito-python/src/lib.rs +++ b/crates/taskito-python/src/lib.rs @@ -36,7 +36,10 @@ fn _taskito(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; #[cfg(feature = "native-async")] - m.add_class::()?; + { + m.add_class::()?; + m.add_class::()?; + } #[cfg(feature = "workflows")] { m.add_class::()?; diff --git a/crates/taskito-python/src/native_async/mod.rs b/crates/taskito-python/src/native_async/mod.rs index 48cd9c93..c3cf3338 100644 --- a/crates/taskito-python/src/native_async/mod.rs +++ b/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; diff --git a/crates/taskito-python/src/native_async/permit.rs b/crates/taskito-python/src/native_async/permit.rs new file mode 100644 index 00000000..f4e33290 --- /dev/null +++ b/crates/taskito-python/src/native_async/permit.rs @@ -0,0 +1,59 @@ +use pyo3::prelude::*; +use tokio::sync::OwnedSemaphorePermit; + +/// An async-dispatch slot, held by the Python coroutine that is using it. +/// +/// The pool acquires a permit before handing a job to the Python executor, so the +/// number of async jobs in flight stays bounded. The permit has to outlive the +/// `submit_job` call — dispatch returns as soon as the coroutine is scheduled, long +/// before it runs — so ownership passes to the coroutine frame itself. CPython then +/// reclaims it on *every* exit path, including the ones that report nothing: an +/// escaping `CancelledError`, a loop stopped mid-flight, or a future dropped before +/// the coroutine ever ran. Releasing from the result-report path instead would leak a +/// slot permanently on each of those. +#[pyclass] +pub struct PyJobPermit(Option); + +impl PyJobPermit { + pub fn new(permit: OwnedSemaphorePermit) -> Self { + Self(Some(permit)) + } +} + +#[pymethods] +impl PyJobPermit { + /// Give the slot back now rather than waiting for the frame to be collected. + /// Idempotent; dropping the handle does the same thing. + fn release(&mut self) { + self.0.take(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use tokio::sync::Semaphore; + + #[tokio::test] + async fn drop_without_release_returns_the_permit() { + let sem = Arc::new(Semaphore::new(1)); + let handle = PyJobPermit::new(sem.clone().acquire_owned().await.unwrap()); + assert_eq!(sem.available_permits(), 0); + + drop(handle); + assert_eq!(sem.available_permits(), 1); + } + + #[tokio::test] + async fn release_is_idempotent() { + let sem = Arc::new(Semaphore::new(1)); + let mut handle = PyJobPermit::new(sem.clone().acquire_owned().await.unwrap()); + + handle.release(); + assert_eq!(sem.available_permits(), 1); + + handle.release(); + assert_eq!(sem.available_permits(), 1); + } +} diff --git a/crates/taskito-python/src/native_async/pool.rs b/crates/taskito-python/src/native_async/pool.rs index b1c9cdc5..6a51f2ce 100644 --- a/crates/taskito-python/src/native_async/pool.rs +++ b/crates/taskito-python/src/native_async/pool.rs @@ -11,12 +11,16 @@ use taskito_core::job::Job; use taskito_core::scheduler::JobResult; use taskito_core::worker::WorkerDispatcher; +use super::permit::PyJobPermit; use super::task_executor::execute_sync_task; /// Dual-dispatch worker pool: async tasks run natively on a Python event loop, -/// sync tasks use `spawn_blocking` (bounded by a semaphore). +/// sync tasks use `spawn_blocking`. Each branch is bounded by its own semaphore — +/// `num_workers` blocking threads, `async_concurrency` coroutines — so neither can +/// starve the other out of a shared budget. pub struct NativeAsyncPool { num_workers: usize, + async_concurrency: usize, task_registry: Arc>, retry_filters: Arc>, async_executor: Arc>, @@ -26,12 +30,14 @@ pub struct NativeAsyncPool { impl NativeAsyncPool { pub fn new( num_workers: usize, + async_concurrency: usize, task_registry: Arc>, retry_filters: Arc>, async_executor: Arc>, ) -> Self { Self { num_workers, + async_concurrency, task_registry, retry_filters, async_executor, @@ -47,7 +53,8 @@ impl WorkerDispatcher for NativeAsyncPool { mut job_rx: tokio::sync::mpsc::Receiver, result_tx: Sender, ) { - let semaphore = Arc::new(Semaphore::new(self.num_workers)); + let sync_semaphore = Arc::new(Semaphore::new(self.num_workers)); + let async_semaphore = Arc::new(Semaphore::new(self.async_concurrency)); while let Some(job) = job_rx.recv().await { if self.shutdown.load(Ordering::Relaxed) { @@ -68,22 +75,39 @@ impl WorkerDispatcher for NativeAsyncPool { }); if is_async { + // Take the slot before submitting, not inside the coroutine: dispatch + // returns as soon as the job is queued, so without this the loop would + // pull jobs from `job_rx` without bound and pile up Running-but-not-yet- + // executing work. Blocking here is the point — it is what bounds the + // in-flight count. The permit rides with the coroutine and is released + // when it finishes (or is collected). + let permit = match async_semaphore.clone().acquire_owned().await { + Ok(p) => p, + Err(_) => break, + }; + // Submit to the Python async executor — brief GIL hold, non-blocking let executor = self.async_executor.clone(); Python::attach(|py| { let exec = executor.bind(py); let payload = PyBytes::new(py, &job.payload); - if let Err(e) = exec.call_method1( - "submit_job", - ( - &job.id, - &job.task_name, - payload, - job.retry_count, - job.max_retries, - &job.queue, - ), - ) { + // Either failure drops the permit handle, returning the slot. + let submitted = Py::new(py, PyJobPermit::new(permit)).and_then(|handle| { + exec.call_method1( + "submit_job", + ( + &job.id, + &job.task_name, + payload, + job.retry_count, + job.max_retries, + &job.queue, + handle, + ), + ) + .map(|_| ()) + }); + if let Err(e) = submitted { log::error!( "[taskito] Failed to submit async task {}[{}]: {e}", job.task_name, @@ -103,7 +127,7 @@ impl WorkerDispatcher for NativeAsyncPool { }); } else { // Sync path: spawn_blocking bounded by semaphore - let permit = match semaphore.clone().acquire_owned().await { + let permit = match sync_semaphore.clone().acquire_owned().await { Ok(p) => p, Err(_) => break, }; diff --git a/crates/taskito-python/src/native_async/result_sender.rs b/crates/taskito-python/src/native_async/result_sender.rs index 72753560..e897d08c 100644 --- a/crates/taskito-python/src/native_async/result_sender.rs +++ b/crates/taskito-python/src/native_async/result_sender.rs @@ -1,13 +1,18 @@ -use crossbeam_channel::Sender; +use crossbeam_channel::{Sender, TrySendError}; use pyo3::prelude::*; use taskito_core::scheduler::JobResult; /// Python-accessible sender for reporting task results back to the Rust scheduler. /// -/// The `AsyncTaskExecutor` on the Python side calls `report_success`, `report_failure`, -/// or `report_cancelled` to push `JobResult` values into the same crossbeam channel -/// that sync tasks use. +/// The `AsyncTaskExecutor` on the Python side calls `try_report_success`, +/// `try_report_failure`, or `try_report_cancelled` to push `JobResult` values into the +/// same crossbeam channel that sync tasks use. +/// +/// Async callers must use the `try_report_*` variants. The blocking `report_*` methods +/// park the caller when the channel is full while still holding the GIL — fatal on the +/// executor's single event-loop thread, because the result drain loop re-acquires the GIL +/// each iteration and could never drain the channel that the sender is waiting on. #[pyclass] pub struct PyResultSender { tx: Sender, @@ -17,31 +22,37 @@ impl PyResultSender { pub fn new(tx: Sender) -> Self { Self { tx } } + + /// `true` when the result was handed off (or the channel is gone, so retrying is + /// futile); `false` only when the channel is full and the caller should back off. + fn try_hand_off(&self, result: JobResult) -> bool { + !matches!(self.tx.try_send(result), Err(TrySendError::Full(_))) + } } #[pymethods] impl PyResultSender { - /// Report a successful task execution. + /// Report a successful task execution. `false` means "channel full, retry". #[pyo3(signature = (job_id, task_name, result, wall_time_ns))] - fn report_success( + fn try_report_success( &self, job_id: String, task_name: String, result: Option>, wall_time_ns: i64, - ) { - let _ = self.tx.send(JobResult::Success { + ) -> bool { + self.try_hand_off(JobResult::Success { job_id, result, task_name, wall_time_ns, - }); + }) } - /// Report a failed task execution. + /// Report a failed task execution. `false` means "channel full, retry". #[pyo3(signature = (job_id, task_name, error, retry_count, max_retries, wall_time_ns, should_retry))] #[allow(clippy::too_many_arguments)] - fn report_failure( + fn try_report_failure( &self, job_id: String, task_name: String, @@ -50,8 +61,8 @@ impl PyResultSender { max_retries: i32, wall_time_ns: i64, should_retry: bool, - ) { - let _ = self.tx.send(JobResult::Failure { + ) -> bool { + self.try_hand_off(JobResult::Failure { job_id, error, retry_count, @@ -60,16 +71,16 @@ impl PyResultSender { wall_time_ns, should_retry, timed_out: false, - }); + }) } - /// Report a cancelled task. + /// Report a cancelled task. `false` means "channel full, retry". #[pyo3(signature = (job_id, task_name, wall_time_ns))] - fn report_cancelled(&self, job_id: String, task_name: String, wall_time_ns: i64) { - let _ = self.tx.send(JobResult::Cancelled { + fn try_report_cancelled(&self, job_id: String, task_name: String, wall_time_ns: i64) -> bool { + self.try_hand_off(JobResult::Cancelled { job_id, task_name, wall_time_ns, - }); + }) } } diff --git a/crates/taskito-python/src/py_config.rs b/crates/taskito-python/src/py_config.rs index b073ce70..0494ad8f 100644 --- a/crates/taskito-python/src/py_config.rs +++ b/crates/taskito-python/src/py_config.rs @@ -34,13 +34,15 @@ pub struct PyTaskConfig { pub circuit_breaker_half_open_probes: Option, #[pyo3(get, set)] pub circuit_breaker_half_open_success_rate: Option, + #[pyo3(get, set)] + pub max_in_flight_per_task: Option, } #[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))] + #[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))] pub fn new( name: String, max_retries: i32, @@ -57,6 +59,7 @@ impl PyTaskConfig { max_concurrent: Option, circuit_breaker_half_open_probes: Option, circuit_breaker_half_open_success_rate: Option, + max_in_flight_per_task: Option, ) -> Self { Self { name, @@ -74,6 +77,7 @@ impl PyTaskConfig { max_concurrent, circuit_breaker_half_open_probes, circuit_breaker_half_open_success_rate, + max_in_flight_per_task, } } } diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index f01e4a0f..91e14ef2 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -277,15 +277,29 @@ impl PyQueue { let queues = queues.unwrap_or_else(|| vec!["default".to_string()]); let queues_str = queues.join(","); + // Clamp once for every consumer: 0 would leave `asyncio.Semaphore(0)` holding + // every native async job forever, and a negative raises inside `start()`. + #[allow(unused_variables)] + let async_concurrency = async_concurrency.max(1) as usize; + + // Bound in-flight work to what this worker can actually run, so it never + // claims more and starves peers sharing the DB. + // + // Only the sync branch runs today: native async dispatch is dormant (see + // the note in decorators.py), so async tasks reach the pool as sync work + // and every job is bounded by `num_workers`. Adding `async_concurrency` + // here would claim jobs nothing can execute — they would sit Running in + // the channel waiting for a blocking thread. Widen this to the sum of the + // two budgets in the change that activates native dispatch, not before. + let max_in_flight = self.num_workers; + let scheduler_config = SchedulerConfig { poll_interval: std::time::Duration::from_millis(self.scheduler_poll_interval_ms), reap_interval: self.scheduler_reap_interval, cleanup_interval: self.scheduler_cleanup_interval, result_ttl_ms: self.result_ttl_ms, batch_size: self.scheduler_batch_size, - // Bound in-flight work to the pool size so this worker never claims - // more than it can run and starve peers sharing the DB. - max_in_flight: Some(self.num_workers), + max_in_flight: Some(max_in_flight), dlq_auto_retry_delay_ms: self.dlq_auto_retry_delay_ms, dlq_auto_retry_max: self.dlq_auto_retry_max, ..SchedulerConfig::default() @@ -369,6 +383,7 @@ impl PyQueue { rate_limit, circuit_breaker, max_concurrent: tc.max_concurrent, + max_in_flight_per_task: tc.max_in_flight_per_task.map(|n| n.max(1) as usize), }, ); } @@ -424,6 +439,9 @@ impl PyQueue { let shutdown = scheduler.shutdown_handle(); + // Headroom over the in-flight cap, so the cheap pre-claim gate binds before + // the channel does: the channel-full path claims and then rolls back, which + // churns storage where the gate does not. let (job_tx, job_rx) = tokio::sync::mpsc::channel(self.num_workers * 2); let (result_tx, result_rx) = crossbeam_channel::bounded(self.num_workers * 2); @@ -498,6 +516,7 @@ impl PyQueue { let pool_arc: Arc = Arc::new(crate::native_async::NativeAsyncPool::new( num_workers, + async_concurrency, registry_arc.clone(), filters_arc.clone(), async_executor, diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 714ec1b1..12bff94a 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -24,6 +24,7 @@ class PyTaskConfig: max_concurrent: int | None circuit_breaker_half_open_probes: int | None circuit_breaker_half_open_success_rate: float | None + max_in_flight_per_task: int | None def __init__( self, @@ -42,6 +43,7 @@ class PyTaskConfig: max_concurrent: int | None = None, circuit_breaker_half_open_probes: int | None = None, circuit_breaker_half_open_success_rate: float | None = None, + max_in_flight_per_task: int | None = None, ) -> None: ... class PyJob: @@ -524,14 +526,14 @@ class PyResultSender: Only available when built with the ``native-async`` feature. """ - def report_success( + def try_report_success( self, job_id: str, task_name: str, result: bytes | None, wall_time_ns: int, - ) -> None: ... - def report_failure( + ) -> bool: ... + def try_report_failure( self, job_id: str, task_name: str, @@ -540,13 +542,13 @@ class PyResultSender: max_retries: int, wall_time_ns: int, should_retry: bool, - ) -> None: ... - def report_cancelled( + ) -> bool: ... + def try_report_cancelled( self, job_id: str, task_name: str, wall_time_ns: int, - ) -> None: ... + ) -> bool: ... def _init_rust_logging() -> None: """Activate the Rust → Python `logging` bridge (idempotent).""" diff --git a/sdks/python/taskito/async_support/executor.py b/sdks/python/taskito/async_support/executor.py index e9806b45..45b51ce9 100644 --- a/sdks/python/taskito/async_support/executor.py +++ b/sdks/python/taskito/async_support/executor.py @@ -6,10 +6,12 @@ import logging import threading import time +from collections.abc import Callable from typing import TYPE_CHECKING, Any from taskito.async_support.context import clear_async_context, set_async_context from taskito.context import current_job +from taskito.events import EventType from taskito.exceptions import TaskCancelledError from taskito.interception.reconstruct import reconstruct_args from taskito.proxies import cleanup_proxies, reconstruct_proxies @@ -20,6 +22,10 @@ logger = logging.getLogger("taskito.async") +# Backoff bounds for handing a result to Rust while the result channel is full. +_REPORT_BACKOFF_START_S = 0.001 +_REPORT_BACKOFF_MAX_S = 0.05 + class AsyncTaskExecutor: """Runs async tasks natively on a dedicated event loop. @@ -63,15 +69,48 @@ def submit_job( retry_count: int, max_retries: int, queue_name: str, + permit: Any = None, ) -> None: - """Submit an async job for execution. Called from Rust — brief GIL hold.""" + """Submit an async job for execution. Called from Rust — brief GIL hold. + + ``permit`` is the dispatch slot the pool acquired for this job; holding it on + the coroutine frame is what releases it on every exit path. It is optional so + the executor can also be driven directly, without a pool. + """ if self._loop is None: raise RuntimeError("AsyncTaskExecutor not started") asyncio.run_coroutine_threadsafe( - self._execute(job_id, task_name, payload, retry_count, max_retries, queue_name), + self._execute( + job_id, task_name, payload, retry_count, max_retries, queue_name, permit + ), self._loop, ) + async def _hand_off(self, send: Callable[[], bool]) -> None: + """Push a result to Rust, backing off while the result channel is full. + + Never drops a result — a dropped one strands the job ``Running`` until the + reaper mislabels it a timeout. Awaiting between attempts is what keeps the + event loop responsive: a blocking send would freeze every other coroutine on + this thread, and would deadlock against the drain loop, which re-acquires the + GIL each iteration and so could never drain the channel being waited on. + """ + delay = _REPORT_BACKOFF_START_S + while not send(): + await asyncio.sleep(delay) + delay = min(delay * 2, _REPORT_BACKOFF_MAX_S) + + def _hand_off_now(self, send: Callable[[], bool], job_id: str) -> None: + """Single best-effort hand-off for a coroutine already being torn down. + + The cancellation path cannot use :meth:`_hand_off`: awaiting inside a + ``CancelledError`` handler risks being cancelled again on a loop that may be + stopping. The channel is drained continuously, so a full one is rare here; the + stale-job reaper remains the backstop if this does drop. + """ + if not send(): + logger.warning("result channel full while cancelling %s; job left to reaper", job_id) + async def _execute( self, job_id: str, @@ -80,9 +119,32 @@ async def _execute( retry_count: int, max_retries: int, queue_name: str, + permit: Any = None, ) -> None: """Execute a single async task with full lifecycle support.""" assert self._semaphore is not None + try: + await self._run_job( + job_id, task_name, payload_bytes, retry_count, max_retries, queue_name + ) + finally: + # Hand the dispatch slot back as soon as the job is done rather than + # waiting for the frame to be collected. Dropping the handle is the + # backstop for paths that never reach here. + if permit is not None: + permit.release() + + async def _run_job( + self, + job_id: str, + task_name: str, + payload_bytes: bytes, + retry_count: int, + max_retries: int, + queue_name: str, + ) -> None: + """Run one job: gate on the concurrency backstop, execute, report.""" + assert self._semaphore is not None async with self._semaphore: start_ns = time.monotonic_ns() token = set_async_context(job_id, task_name, retry_count, queue_name) @@ -91,6 +153,11 @@ async def _execute( result: Any = None error: Exception | None = None completed_mw: list[Any] = [] + # Set before the hand-off so a raising send can't also fire report_failure + # for the same job — the scheduler would then see two results for one job. + reported = False + # Cancellation has its own event, emitted by the Rust outcome loop. + cancelled = False try: queue = self._queue_ref @@ -161,27 +228,55 @@ async def _execute( queue._serialize_result(task_name, result) if result is not None else None ) wall_ns = time.monotonic_ns() - start_ns - self._sender.report_success(job_id, task_name, result_bytes, wall_ns) + reported = True + await self._hand_off( + lambda: self._sender.try_report_success( + job_id, task_name, result_bytes, wall_ns + ) + ) except TaskCancelledError: error = None # Don't treat cancellation as an error for middleware - wall_ns = time.monotonic_ns() - start_ns - self._sender.report_cancelled(job_id, task_name, wall_ns) + cancelled = True + if not reported: + wall_ns = time.monotonic_ns() - start_ns + reported = True + await self._hand_off( + lambda: self._sender.try_report_cancelled(job_id, task_name, wall_ns) + ) + + except asyncio.CancelledError: + # CancelledError is a BaseException, so `except Exception` misses it and + # the job would sit Running until the reaper mislabelled it a timeout. + # Report, then re-raise — swallowing a cancellation breaks loop shutdown. + cancelled = True + if not reported: + wall_ns = time.monotonic_ns() - start_ns + reported = True + self._hand_off_now( + lambda: self._sender.try_report_cancelled(job_id, task_name, wall_ns), + job_id, + ) + raise except Exception as exc: error = exc - wall_ns = time.monotonic_ns() - start_ns - error_msg = encode_task_error(exc) - should_retry = self._check_retry(task_name, exc) - self._sender.report_failure( - job_id, - task_name, - error_msg, - retry_count, - max_retries, - wall_ns, - should_retry, - ) + if not reported: + wall_ns = time.monotonic_ns() - start_ns + reported = True + error_msg = encode_task_error(exc) + should_retry = self._check_retry(task_name, exc) + await self._hand_off( + lambda: self._sender.try_report_failure( + job_id, + task_name, + error_msg, + retry_count, + max_retries, + wall_ns, + should_retry, + ) + ) finally: # Release task/request-scoped resources @@ -201,8 +296,36 @@ async def _execute( except Exception: logger.exception("middleware after() error") + # Emit job lifecycle events. The Rust outcome loop deliberately + # skips Success — the blocking wrapper emits it — and native + # dispatch bypasses that wrapper, so it has to emit here too or + # nothing downstream (workflow tracker, subscribers) ever learns + # the job finished. Cancellation is the exception — that outcome + # does have a Rust-side event. + if not cancelled: + self._emit_lifecycle_event(job_id, task_name, queue_name, error) + clear_async_context(token) + def _emit_lifecycle_event( + self, job_id: str, task_name: str, queue_name: str, error: Exception | None + ) -> None: + """Emit JOB_COMPLETED/JOB_FAILED, mirroring the blocking task wrapper.""" + payload: dict[str, Any] = { + "task_name": task_name, + "job_id": job_id, + "queue": queue_name, + } + if error is not None: + payload["error"] = str(error) + try: + self._queue_ref._emit_event( + EventType.JOB_FAILED if error is not None else EventType.JOB_COMPLETED, + payload, + ) + except Exception: + logger.exception("job lifecycle event error") + def _check_retry(self, task_name: str, exc: Exception) -> bool: """Check retry filters to decide if this exception should be retried.""" filters = self._queue_ref._task_retry_filters.get(task_name) diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 0f5616c2..b26340b0 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -385,6 +385,10 @@ def task( on_false: str = "defer", predicate_extras: dict[str, Any] | None = None, default_defer_seconds: float = 60.0, + # Appended, not slotted next to `max_concurrent`: this signature is not + # keyword-only, so inserting mid-list would silently rebind the positional + # arguments after it. + max_in_flight_per_task: int | None = None, ) -> Callable[[Callable[..., Any]], TaskWrapper]: """Decorator to register a function as a background task. @@ -421,6 +425,11 @@ def task( (5 minutes) if not set. max_concurrent: Maximum number of concurrent running instances of this task. ``None`` means no limit. + max_in_flight_per_task: Cap on this task's share of a single worker's + dispatch slots, so one slow task cannot occupy the whole pool and + starve the others. Unlike ``max_concurrent`` (a cluster-wide cap + that costs a database read), this is in-process and free. + ``None`` lets the task use the whole pool. compensates: Optional reference to a task that compensates this one. When this task runs as part of a workflow saga and a later step fails, the framework enqueues the compensation @@ -584,6 +593,14 @@ def decorator(fn: Callable) -> TaskWrapper: # Wrap the function with hooks, middleware, and context wrapped = self._wrap_task(fn, task_name, soft_timeout) + + # NOTE: `_taskito_is_async` is deliberately NOT set on `wrapped`. + # The pool reads that flag off this registry entry to choose native + # dispatch, so setting it here activates the native path — and that + # path reimplements the task lifecycle and still lacks this wrapper's + # queue hooks, saga compensation context, soft_timeout and per-item + # batch handling. Async tasks therefore run through this blocking + # wrapper (via run_maybe_async) until those gaps are closed. self._task_registry[task_name] = wrapped cb_threshold = None @@ -615,6 +632,7 @@ def decorator(fn: Callable) -> TaskWrapper: max_concurrent=max_concurrent, circuit_breaker_half_open_probes=cb_half_open_probes, circuit_breaker_half_open_success_rate=cb_half_open_success_rate, + max_in_flight_per_task=max_in_flight_per_task, ) self._task_configs.append(config) diff --git a/sdks/python/tests/worker/test_in_flight_bound.py b/sdks/python/tests/worker/test_in_flight_bound.py new file mode 100644 index 00000000..2f8a6f0f --- /dev/null +++ b/sdks/python/tests/worker/test_in_flight_bound.py @@ -0,0 +1,50 @@ +"""A worker must not claim more jobs than it can actually run.""" + +import threading +import time +from pathlib import Path +from typing import Any + +from taskito import Queue + +PollUntil = Any # the conftest fixture's runtime type + + +def test_worker_does_not_claim_more_than_it_can_run(tmp_path: Path) -> None: + """Claimed-but-unrunnable jobs sit Running, so peers sharing the DB skip them. + + Dequeue flips a job to Running, so over-claiming is invisible in results — every + job still completes, just after stranding Running for however long the worker + takes to reach it. Watch the Running count instead: with two worker threads it + must stay near two, no matter how deep the backlog. + """ + queue = Queue(db_path=str(tmp_path / "bound.db"), workers=2) + job_count = 40 + + @queue.task(name="slow") + def slow() -> int: + time.sleep(0.25) + return 1 + + for _ in range(job_count): + slow.delay() + + thread = threading.Thread(target=queue.run_worker, daemon=True) + thread.start() + try: + peak_running = 0 + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + stats = queue.stats() + peak_running = max(peak_running, stats.get("running", 0)) + if stats.get("completed", 0) >= job_count: + break + time.sleep(0.02) + finally: + queue.shutdown() + thread.join(timeout=5) + + assert queue.stats().get("completed", 0) == job_count + # Two threads can be executing, plus one in hand-off; the bound is generous + # because the failure mode is the whole backlog going Running at once. + assert peak_running <= 8, f"worker claimed {peak_running} jobs but can only run 2" diff --git a/sdks/python/tests/worker/test_native_async.py b/sdks/python/tests/worker/test_native_async.py index e4e1a211..299ce572 100644 --- a/sdks/python/tests/worker/test_native_async.py +++ b/sdks/python/tests/worker/test_native_async.py @@ -167,11 +167,11 @@ class FakeWrapper: payload = cloudpickle.dumps(((2, 3), {})) executor.submit_job("job-1", "test_mod.my_task", payload, 0, 3, "default") - poll_until(lambda: sender.report_success.called, message="job-1 result not reported") + poll_until(lambda: sender.try_report_success.called, message="job-1 result not reported") executor.stop() - sender.report_success.assert_called_once() - call_args = sender.report_success.call_args + sender.try_report_success.assert_called_once() + call_args = sender.try_report_success.call_args assert call_args[0][0] == "job-1" assert call_args[0][1] == "test_mod.my_task" result = cloudpickle.loads(call_args[0][2]) @@ -211,11 +211,11 @@ class FakeWrapper: payload = cloudpickle.dumps(((), {})) executor.submit_job("job-2", "mod.failing_task", payload, 0, 3, "default") - poll_until(lambda: sender.report_failure.called, message="job-2 failure not reported") + poll_until(lambda: sender.try_report_failure.called, message="job-2 failure not reported") executor.stop() - sender.report_failure.assert_called_once() - call_args = sender.report_failure.call_args + sender.try_report_failure.assert_called_once() + call_args = sender.try_report_failure.call_args assert call_args[0][0] == "job-2" assert "boom" in call_args[0][2] assert call_args[0][6] is True # should_retry @@ -254,11 +254,13 @@ class FakeWrapper: payload = cloudpickle.dumps(((), {})) executor.submit_job("job-3", "mod.cancelling_task", payload, 0, 3, "default") - poll_until(lambda: sender.report_cancelled.called, message="job-3 cancellation not reported") + poll_until( + lambda: sender.try_report_cancelled.called, message="job-3 cancellation not reported" + ) executor.stop() - sender.report_cancelled.assert_called_once() - assert sender.report_cancelled.call_args[0][0] == "job-3" + sender.try_report_cancelled.assert_called_once() + assert sender.try_report_cancelled.call_args[0][0] == "job-3" def test_async_retry_filter(poll_until: PollUntil) -> None: @@ -297,11 +299,11 @@ class FakeWrapper: payload = cloudpickle.dumps(((), {})) executor.submit_job("job-4", "mod.flaky_task", payload, 0, 3, "default") - poll_until(lambda: sender.report_failure.called, message="job-4 failure not reported") + poll_until(lambda: sender.try_report_failure.called, message="job-4 failure not reported") executor.stop() - sender.report_failure.assert_called_once() - assert sender.report_failure.call_args[0][6] is False # should_retry = False + sender.try_report_failure.assert_called_once() + assert sender.try_report_failure.call_args[0][6] is False # should_retry = False def test_async_concurrency_limit(poll_until: PollUntil) -> None: @@ -350,14 +352,14 @@ class FakeWrapper: executor.submit_job(f"job-{i}", "mod.slow_task", payload, 0, 3, "default") poll_until( - lambda: sender.report_success.call_count >= 5, + lambda: sender.try_report_success.call_count >= 5, timeout=10, message="not all 5 slow_task jobs reported success", ) executor.stop() assert max_concurrent <= 2 - assert sender.report_success.call_count == 5 + assert sender.try_report_success.call_count == 5 def test_async_middleware_hooks(poll_until: PollUntil) -> None: @@ -449,11 +451,11 @@ class FakeWrapper: payload = cloudpickle.dumps(((), {})) executor.submit_job("inj-job", "mod.db_task", payload, 0, 3, "default") - poll_until(lambda: sender.report_success.called, message="inj-job result not reported") + poll_until(lambda: sender.try_report_success.called, message="inj-job result not reported") executor.stop() - sender.report_success.assert_called_once() - result = cloudpickle.loads(sender.report_success.call_args[0][2]) + sender.try_report_success.assert_called_once() + result = cloudpickle.loads(sender.try_report_success.call_args[0][2]) assert result == "got-fake-conn" @@ -548,10 +550,245 @@ class FakeWrapper: payload = serializer.dumps(([2, 3], {})) executor.submit_job("ser-job", "mod.add", payload, 0, 3, "default") - poll_until(lambda: sender.report_success.called, message="ser-job result not reported") + poll_until(lambda: sender.try_report_success.called, message="ser-job result not reported") executor.stop() queue_ref._deserialize_payload.assert_called_once_with("mod.add", payload) queue_ref._serialize_result.assert_called_once_with("mod.add", 5) - result = serializer.loads(sender.report_success.call_args[0][2]) + result = serializer.loads(sender.try_report_success.call_args[0][2]) assert result == 5 + + +# ── Backpressure: permits, result hand-off, cancellation (S16/S17) ── + + +class FakePermit: + """Stands in for the Rust PyJobPermit handle the pool hands to the executor.""" + + def __init__(self) -> None: + self.releases = 0 + + def release(self) -> None: + self.releases += 1 + + +def _backpressure_executor(sender: Any, fn: Any, task_name: str) -> Any: + """An executor wired to a single async `fn`, with everything else mocked out.""" + import cloudpickle + + from taskito.async_support.executor import AsyncTaskExecutor + + class FakeWrapper: + _taskito_async_fn = staticmethod(fn) + + queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) + queue_ref._serialize_result.side_effect = lambda _name, result: cloudpickle.dumps(result) + queue_ref._interceptor = None + queue_ref._proxy_registry = None + queue_ref._test_mode_active = False + queue_ref._resource_runtime = None + queue_ref._task_inject_map = {} + queue_ref._task_retry_filters = {} + queue_ref._get_middleware_chain.return_value = [] + queue_ref._proxy_metrics = None + + executor = AsyncTaskExecutor(sender, {task_name: FakeWrapper()}, queue_ref, max_concurrency=10) + executor.start() + return executor + + +def _payload() -> bytes: + import cloudpickle + + payload: bytes = cloudpickle.dumps(((), {})) + return payload + + +def test_cancelled_coroutine_reports_cancelled(poll_until: PollUntil) -> None: + """asyncio.CancelledError is a BaseException, so `except Exception` misses it. + + Unreported, the job sits Running until the reaper mislabels it a timeout. + """ + sender = MagicMock() + sender.try_report_cancelled.return_value = True + + async def cancelled_task() -> None: + raise asyncio.CancelledError + + executor = _backpressure_executor(sender, cancelled_task, "mod.cancelled_task") + executor.submit_job("job-c", "mod.cancelled_task", _payload(), 0, 3, "default") + poll_until( + lambda: sender.try_report_cancelled.called, + message="asyncio.CancelledError was not reported", + ) + executor.stop() + + assert sender.try_report_cancelled.call_args[0][0] == "job-c" + assert not sender.try_report_failure.called, "a cancellation is not a failure" + + +def test_no_double_report_when_report_raises(poll_until: PollUntil) -> None: + """A raising success hand-off must not also fire a failure for the same job.""" + sender = MagicMock() + sender.try_report_success.side_effect = RuntimeError("channel exploded") + sender.try_report_failure.return_value = True + + async def ok_task() -> int: + return 1 + + executor = _backpressure_executor(sender, ok_task, "mod.ok_task") + executor.submit_job("job-d", "mod.ok_task", _payload(), 0, 3, "default") + poll_until(lambda: sender.try_report_success.called, message="success never attempted") + executor.stop() + + assert not sender.try_report_failure.called, "one job must yield at most one result" + + +def test_permit_released_on_success(poll_until: PollUntil) -> None: + sender = MagicMock() + sender.try_report_success.return_value = True + permit = FakePermit() + + async def ok_task() -> int: + return 1 + + executor = _backpressure_executor(sender, ok_task, "mod.ok_task") + executor.submit_job("job-p", "mod.ok_task", _payload(), 0, 3, "default", permit) + poll_until(lambda: permit.releases == 1, message="permit not released on success") + executor.stop() + + +def test_permit_released_on_exception(poll_until: PollUntil) -> None: + sender = MagicMock() + sender.try_report_failure.return_value = True + permit = FakePermit() + + async def boom_task() -> None: + raise ValueError("boom") + + executor = _backpressure_executor(sender, boom_task, "mod.boom_task") + executor.submit_job("job-p", "mod.boom_task", _payload(), 0, 3, "default", permit) + poll_until(lambda: permit.releases == 1, message="permit not released on exception") + executor.stop() + + +def test_permit_released_on_cancellation(poll_until: PollUntil) -> None: + """The path that leaks a slot forever if release rides on the result report.""" + sender = MagicMock() + sender.try_report_cancelled.return_value = True + permit = FakePermit() + + async def cancelled_task() -> None: + raise asyncio.CancelledError + + executor = _backpressure_executor(sender, cancelled_task, "mod.cancelled_task") + executor.submit_job("job-p", "mod.cancelled_task", _payload(), 0, 3, "default", permit) + poll_until(lambda: permit.releases == 1, message="permit not released on cancellation") + executor.stop() + + +def test_report_retries_when_channel_full_without_stalling_the_loop( + poll_until: PollUntil, +) -> None: + """A full result channel must back off, not block — and never drop the result. + + `job-a` is refused until `job-b` reports, so `job-b` can only get through if the + event loop stayed responsive while `job-a` was backing off. A blocking send would + wedge both and time this test out. + """ + sender = MagicMock() + reported_b = threading.Event() + + def send(job_id: str, task_name: str, result: Any, wall_ns: int) -> bool: + if job_id == "job-b": + reported_b.set() + return True + return reported_b.is_set() + + sender.try_report_success.side_effect = send + + async def ok_task() -> int: + return 1 + + executor = _backpressure_executor(sender, ok_task, "mod.ok_task") + executor.submit_job("job-a", "mod.ok_task", _payload(), 0, 3, "default") + executor.submit_job("job-b", "mod.ok_task", _payload(), 0, 3, "default") + + poll_until( + lambda: ( + sum(1 for c in sender.try_report_success.call_args_list if c[0][0] == "job-a") > 1 + and reported_b.is_set() + ), + message="job-a never retried, or job-b starved behind it", + ) + executor.stop() + + a_calls = [c for c in sender.try_report_success.call_args_list if c[0][0] == "job-a"] + assert len(a_calls) > 1, "job-a must retry rather than drop its result" + + +def test_native_dispatch_emits_job_completed(poll_until: PollUntil) -> None: + """Native dispatch must emit JOB_COMPLETED itself. + + The Rust outcome loop skips Success on the grounds that the blocking task + wrapper emits it — but native dispatch bypasses that wrapper. Without this, + nothing downstream (notably the workflow tracker) ever learns the job + finished, and a workflow with an async step hangs in `running` forever. + """ + from taskito.events import EventType + + sender = MagicMock() + sender.try_report_success.return_value = True + + async def ok_task() -> int: + return 1 + + executor = _backpressure_executor(sender, ok_task, "mod.ok_task") + queue_ref = executor._queue_ref + executor.submit_job("job-e", "mod.ok_task", _payload(), 0, 3, "default") + poll_until(lambda: queue_ref._emit_event.called, message="no lifecycle event emitted") + executor.stop() + + event_type, payload = queue_ref._emit_event.call_args[0] + assert event_type == EventType.JOB_COMPLETED + assert payload["job_id"] == "job-e" + assert payload["task_name"] == "mod.ok_task" + assert payload["queue"] == "default" + + +def test_native_dispatch_emits_job_failed(poll_until: PollUntil) -> None: + from taskito.events import EventType + + sender = MagicMock() + sender.try_report_failure.return_value = True + + async def boom_task() -> None: + raise ValueError("boom") + + executor = _backpressure_executor(sender, boom_task, "mod.boom_task") + queue_ref = executor._queue_ref + executor.submit_job("job-f", "mod.boom_task", _payload(), 0, 3, "default") + poll_until(lambda: queue_ref._emit_event.called, message="no lifecycle event emitted") + executor.stop() + + event_type, payload = queue_ref._emit_event.call_args[0] + assert event_type == EventType.JOB_FAILED + assert payload["error"] == "boom" + + +def test_native_dispatch_does_not_emit_completed_on_cancel(poll_until: PollUntil) -> None: + """Cancellation has its own Rust-side event; emitting COMPLETED would be a lie.""" + sender = MagicMock() + sender.try_report_cancelled.return_value = True + + async def cancelling_task() -> None: + raise TaskCancelledError("nope") + + executor = _backpressure_executor(sender, cancelling_task, "mod.cancelling_task") + queue_ref = executor._queue_ref + executor.submit_job("job-g", "mod.cancelling_task", _payload(), 0, 3, "default") + poll_until(lambda: sender.try_report_cancelled.called, message="cancellation not reported") + executor.stop() + + assert not queue_ref._emit_event.called, "a cancelled job did not complete"