From d998d32bd8f499a00834931064c1701dcb5fe97e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:41:37 +0530 Subject: [PATCH 1/7] fix(async): correct the native executor's result path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four defects, all in how a native async job reports its outcome: The blocking send parked the caller while holding the GIL, and the result drain loop re-acquires the GIL every iteration, so a full channel wedged both threads for good. Hand off with try_send and back off with await instead. A coroutine cancelled via asyncio.CancelledError reported nothing, because that is a BaseException and escaped `except Exception` — the job sat Running until the reaper mislabelled it a timeout. A success send that raised also reported a failure for the same job. And nothing emitted JOB_COMPLETED, which the Rust outcome loop skips on the grounds that the blocking wrapper sends it, so a workflow with an async step would hang in running forever. Also releases the dispatch permit that the pool now attaches to each job. --- .../src/native_async/result_sender.rs | 47 ++++-- sdks/python/taskito/_taskito.pyi | 12 +- sdks/python/taskito/async_support/executor.py | 157 ++++++++++++++++-- 3 files changed, 175 insertions(+), 41 deletions(-) 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/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 714ec1b1..12b2b443 100644 --- a/sdks/python/taskito/_taskito.pyi +++ b/sdks/python/taskito/_taskito.pyi @@ -524,14 +524,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 +540,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) From 6ecb1e1738a16b334ceb22df3ec077960e7735b0 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:42:11 +0530 Subject: [PATCH 2/7] perf(async): bound async jobs in flight with a dispatch permit The async branch submitted to the Python executor without taking a permit, so the dispatch loop pulled from the job channel unbounded and piled up work that was Running but not yet executing. Take a permit before submitting and let the coroutine frame own it, so it comes back even on the paths that report nothing (an escaping CancelledError, a loop stopped mid-flight). The permit is sized to async_concurrency but stays inert until native dispatch is activated, since every task reaches the pool as sync work today. Note there why max_in_flight cannot widen to cover it before then. --- crates/taskito-python/src/lib.rs | 5 +- crates/taskito-python/src/native_async/mod.rs | 2 + .../taskito-python/src/native_async/permit.rs | 59 +++++++++++++++++++ .../taskito-python/src/native_async/pool.rs | 52 +++++++++++----- crates/taskito-python/src/py_queue/worker.rs | 19 +++++- 5 files changed, 119 insertions(+), 18 deletions(-) create mode 100644 crates/taskito-python/src/native_async/permit.rs 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/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index f01e4a0f..a3dcfec0 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -277,15 +277,24 @@ impl PyQueue { let queues = queues.unwrap_or_else(|| vec!["default".to_string()]); let queues_str = queues.join(","); + // 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() @@ -424,6 +433,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 +510,7 @@ impl PyQueue { let pool_arc: Arc = Arc::new(crate::native_async::NativeAsyncPool::new( num_workers, + async_concurrency.max(1) as usize, registry_arc.clone(), filters_arc.clone(), async_executor, From fb7bbb4fd2d117652d1be0f4161260feee01fc6e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:42:30 +0530 Subject: [PATCH 3/7] perf(scheduler): cap per-task in-flight jobs One slow task could occupy every dispatch slot on a worker and starve the rest. Track in-flight jobs by task and reject past a per-task cap, reusing the existing rollback-and-reschedule path, so a saturated task never blocks its neighbours. Complements max_concurrent, which is the cluster-wide cap and costs a database read; this one is in-process. --- crates/taskito-core/src/scheduler/mod.rs | 205 +++++++++++++++++- crates/taskito-core/src/scheduler/poller.rs | 13 +- crates/taskito-java/src/worker.rs | 2 + .../taskito-node/src/convert/task_config.rs | 2 + crates/taskito-python/src/py_config.rs | 6 +- crates/taskito-python/src/py_queue/worker.rs | 1 + sdks/python/taskito/_taskito.pyi | 2 + sdks/python/taskito/mixins/decorators.py | 7 + 8 files changed, 228 insertions(+), 10 deletions(-) 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/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 a3dcfec0..55675546 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -378,6 +378,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), }, ); } diff --git a/sdks/python/taskito/_taskito.pyi b/sdks/python/taskito/_taskito.pyi index 12b2b443..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: diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 0f5616c2..44b82e7f 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -378,6 +378,7 @@ def task( codecs: list[str] | None = None, max_retry_delay: int | None = None, max_concurrent: int | None = None, + max_in_flight_per_task: int | None = None, idempotent: bool = False, compensates: TaskWrapper | str | None = None, batch: bool | dict[str, Any] | None = None, @@ -421,6 +422,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 @@ -615,6 +621,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) From 8d3c99847083b4b191536a64d2914776aef44a9a Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:43:14 +0530 Subject: [PATCH 4/7] test(async): cover backpressure, permits and result reporting Regression tests for the four result-path defects and the per-task bulkhead. The channel-full test refuses one job until a second reports, so it only passes if the loop stayed responsive while the first backed off. --- sdks/python/tests/worker/test_native_async.py | 275 ++++++++++++++++-- 1 file changed, 256 insertions(+), 19 deletions(-) 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" From 3ce09aea7b1f31668ad8dd713ee806619d4665e7 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:43:26 +0530 Subject: [PATCH 5/7] test(worker): assert a worker never over-claims MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Over-claiming is invisible in results — every job still completes — so this watches the Running count against a deep backlog instead. Verified it fails (40 <= 8) when the in-flight bound is widened past what the pool can run. --- .../tests/worker/test_in_flight_bound.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 sdks/python/tests/worker/test_in_flight_bound.py 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" From 35412aaf653eb568b6c3e175dcdf80a25fa5e53e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:43:32 +0530 Subject: [PATCH 6/7] docs: record that native async dispatch never runs The pool reads _taskito_is_async off the registry entry, but the decorator sets it on the TaskWrapper it returns, so the flag is always missing and every async task takes the sync branch. Setting it here would activate a path that still lacks this wrapper's hooks, saga context, soft_timeout and per-item batch handling, so note why it stays off rather than leave the mismatch looking like an oversight. --- sdks/python/taskito/mixins/decorators.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 44b82e7f..4d468f23 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -590,6 +590,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 From de978fd792905840633f76628bcaa76ab9689164 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Wed, 15 Jul 2026 23:40:09 +0530 Subject: [PATCH 7/7] fix: clamp async_concurrency and stop shifting task() arguments The pool clamped async_concurrency but the Python executor got the raw value, so async_concurrency=0 built asyncio.Semaphore(0) and held every native async job forever, and a negative raised inside start(). Clamp once for both consumers. task() is not keyword-only, so max_in_flight_per_task landing next to max_concurrent rebound every positional argument after it. Append instead. --- crates/taskito-python/src/py_queue/worker.rs | 7 ++++++- sdks/python/taskito/mixins/decorators.py | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 55675546..91e14ef2 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -277,6 +277,11 @@ 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. // @@ -511,7 +516,7 @@ impl PyQueue { let pool_arc: Arc = Arc::new(crate::native_async::NativeAsyncPool::new( num_workers, - async_concurrency.max(1) as usize, + async_concurrency, registry_arc.clone(), filters_arc.clone(), async_executor, diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 4d468f23..b26340b0 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -378,7 +378,6 @@ def task( codecs: list[str] | None = None, max_retry_delay: int | None = None, max_concurrent: int | None = None, - max_in_flight_per_task: int | None = None, idempotent: bool = False, compensates: TaskWrapper | str | None = None, batch: bool | dict[str, Any] | None = None, @@ -386,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.