perf(scaling): tier 3 async backpressure (S16, S17, S19)#433
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe scheduler adds per-task in-flight limits and task-aware accounting. Python task configuration exposes the limit, while native async execution gains separate concurrency budgets, permit handoff, non-blocking result reporting, backoff, and lifecycle events. ChangesScheduler and async execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PyQueue
participant NativeAsyncPool
participant AsyncTaskExecutor
participant PyResultSender
participant Scheduler
PyQueue->>NativeAsyncPool: Submit async job
NativeAsyncPool->>AsyncTaskExecutor: Pass PyJobPermit
AsyncTaskExecutor->>PyResultSender: Try to report result
PyResultSender-->>AsyncTaskExecutor: Accept or signal full channel
AsyncTaskExecutor->>Scheduler: Release permit and emit lifecycle event
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
crates/taskito-core/src/scheduler/mod.rs (1)
313-348: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winPer-task bulkhead (
max_in_flight_per_task) silently becomes a no-op unless the pool-widemax_in_flightcap is also set.
track_in_flight(Lines 313-320) andrelease_in_flight(Lines 333-348) only touch theInFlightmap whenself.config.max_in_flight.is_some().task_in_flight(Lines 322-328) reads from that very same map. So if aTaskConfig.max_in_flight_per_taskis configured whileSchedulerConfig.max_in_flightisNone(the default), the map is never populated,task_in_flight()always returns0, and the new per-task gate inpoller.rs::check_post_claim_concurrency(Lines 421-430) can never reject a job — the bulkhead this PR (S19) introduces is silently disabled. The Python binding happens to always passmax_in_flight: Some(..)(py_queue/worker.rs), which masks this for the shipped SDK, but it's a latent trap for any othertaskito-coreconsumer or future binding.🐛 Proposed fix: track/release unconditionally, gate only the wake logic on `max_in_flight`
- /// 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, 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, task_name); - } - } + /// Record a job as in flight once it is handed to the worker channel. + /// Always tracked: the per-task bulkhead (`max_in_flight_per_task`) needs + /// this even when the pool-wide `max_in_flight` cap is unset. + fn track_in_flight(&self, job_id: &str, task_name: &str) { + self.in_flight + .lock() + .unwrap_or_else(|p| p.into_inner()) + .insert(job_id, task_name); + }fn release_in_flight(&self, job_id: &str) { - let Some(max) = self.config.max_in_flight else { - return; - }; let mut in_flight = self.in_flight.lock().unwrap_or_else(|p| p.into_inner()); - let was_full = in_flight.len() >= max; + let was_full = match self.config.max_in_flight { + Some(max) => in_flight.len() >= max, + None => false, + }; let freed = in_flight.remove(job_id); drop(in_flight); if freed && was_full { self.dispatch_wake.notify_one(); } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/taskito-core/src/scheduler/mod.rs` around lines 313 - 348, Make the InFlight map track and release jobs whenever per-task or pool-wide concurrency tracking is needed, rather than requiring max_in_flight to be configured. Update track_in_flight and release_in_flight to operate unconditionally, while retaining the max_in_flight check only for deciding whether release_in_flight should wake the poller; ensure task_in_flight can observe per-task counts when only max_in_flight_per_task is configured.crates/taskito-python/src/py_queue/worker.rs (1)
465-520: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClamp
async_concurrencybefore creatingAsyncTaskExecutor
AsyncTaskExecutor.start()buildsasyncio.Semaphore(self._max_concurrency)directly, soasync_concurrency=0leaves async jobs waiting forever even though the Rust side now allows dispatch.Suggested fix
let executor = cls.call1(( sender_obj, registry_arc.clone_ref(py), queue_ref, - async_concurrency, + async_concurrency.max(1), ))?;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/taskito-python/src/py_queue/worker.rs` around lines 465 - 520, Clamp async_concurrency to at least 1 before the native-async block creates AsyncTaskExecutor, so the value passed to its constructor and start path is always positive. Update the surrounding worker initialization flow without changing prefork or non-native-async behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@crates/taskito-core/src/scheduler/mod.rs`:
- Around line 313-348: Make the InFlight map track and release jobs whenever
per-task or pool-wide concurrency tracking is needed, rather than requiring
max_in_flight to be configured. Update track_in_flight and release_in_flight to
operate unconditionally, while retaining the max_in_flight check only for
deciding whether release_in_flight should wake the poller; ensure task_in_flight
can observe per-task counts when only max_in_flight_per_task is configured.
In `@crates/taskito-python/src/py_queue/worker.rs`:
- Around line 465-520: Clamp async_concurrency to at least 1 before the
native-async block creates AsyncTaskExecutor, so the value passed to its
constructor and start path is always positive. Update the surrounding worker
initialization flow without changing prefork or non-native-async behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0fa8908e-b9e1-4075-babc-a5796fbc2561
📒 Files selected for processing (15)
crates/taskito-core/src/scheduler/mod.rscrates/taskito-core/src/scheduler/poller.rscrates/taskito-java/src/worker.rscrates/taskito-node/src/convert/task_config.rscrates/taskito-python/src/lib.rscrates/taskito-python/src/native_async/mod.rscrates/taskito-python/src/native_async/permit.rscrates/taskito-python/src/native_async/pool.rscrates/taskito-python/src/native_async/result_sender.rscrates/taskito-python/src/py_config.rscrates/taskito-python/src/py_queue/worker.rssdks/python/taskito/_taskito.pyisdks/python/taskito/async_support/executor.pysdks/python/taskito/mixins/decorators.pysdks/python/tests/worker/test_native_async.py
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.
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.
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.
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.
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.
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.
cbb85c3 to
35412aa
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/taskito-python/src/py_queue/worker.rs`:
- Line 514: Normalize async_concurrency before passing it to the Python
AsyncTaskExecutor, using the same minimum-of-one behavior already applied for
NativeAsyncPool::new. Update the executor argument near the
async_concurrency.max(1) expression so negative and zero values become 1, while
positive values remain unchanged.
In `@sdks/python/taskito/mixins/decorators.py`:
- Line 381: Move the max_in_flight_per_task parameter to the end of the
task(...) signature, after all existing parameters, preserving the positional
order of idempotent, compensates, batch, and later arguments.
In `@sdks/python/tests/worker/test_in_flight_bound.py`:
- Around line 47-50: Update the peak-running assertion in the in-flight bound
test to compare against the actual running limit reported by queue.stats(),
rather than the hard-coded 8 cap. Preserve the completed-job assertion and use
the stats-exposed limit so the test reflects the configured two-worker bound.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 957efe8a-6281-41a6-bbd3-27f7f3eeb524
📒 Files selected for processing (16)
crates/taskito-core/src/scheduler/mod.rscrates/taskito-core/src/scheduler/poller.rscrates/taskito-java/src/worker.rscrates/taskito-node/src/convert/task_config.rscrates/taskito-python/src/lib.rscrates/taskito-python/src/native_async/mod.rscrates/taskito-python/src/native_async/permit.rscrates/taskito-python/src/native_async/pool.rscrates/taskito-python/src/native_async/result_sender.rscrates/taskito-python/src/py_config.rscrates/taskito-python/src/py_queue/worker.rssdks/python/taskito/_taskito.pyisdks/python/taskito/async_support/executor.pysdks/python/taskito/mixins/decorators.pysdks/python/tests/worker/test_in_flight_bound.pysdks/python/tests/worker/test_native_async.py
🚧 Files skipped from review as they are similar to previous changes (12)
- crates/taskito-python/src/native_async/mod.rs
- crates/taskito-python/src/lib.rs
- crates/taskito-python/src/native_async/permit.rs
- crates/taskito-python/src/py_config.rs
- crates/taskito-python/src/native_async/pool.rs
- crates/taskito-node/src/convert/task_config.rs
- sdks/python/taskito/_taskito.pyi
- crates/taskito-core/src/scheduler/poller.rs
- crates/taskito-python/src/native_async/result_sender.rs
- crates/taskito-java/src/worker.rs
- sdks/python/taskito/async_support/executor.py
- crates/taskito-core/src/scheduler/mod.rs
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.
Tier 3 of the scaling catalog: bound async dispatch, stop the result path from
freezing the event loop, and keep one slow task from taking a whole worker.
What's here
S17 — non-blocking result report.
PyResultSender's blocking send parked thecaller while holding the GIL. The result drain loop re-acquires the GIL every
iteration, so a full channel wedged both threads permanently. Results now hand off
with
try_sendand back off withawait, which keeps the loop responsive and neverdrops a result. This had to land before the sizing change below — raising the
in-flight cap first would have made the deadlock reachable rather than latent.
S16 — permit before submit. The async branch submitted to the Python executor
without taking a permit, so the dispatch loop drained the job channel unbounded and
piled up work that was
Runningbut not yet executing. A permit is now taken beforesubmit and owned by the coroutine frame, so it comes back on every exit path —
including the ones that report nothing (an escaping
CancelledError, a loop stoppedmid-flight). A
job_id -> permitmap would leak a slot permanently on each of those.The in-flight cap becomes
num_workers + async_concurrency. Sync and async dispatchhave separate budgets but share one in-flight set, so the cap is their sum; a
max()would let either starve the other. Both channels are sized to it so the cheap
pre-claim gate binds before the channel path, which claims and then rolls back.
S19 — per-task bulkhead.
@queue.task(max_in_flight_per_task=N)caps a task'sshare of one worker's slots. It lives in the scheduler, not the pool: the pool cannot
reject (
WorkerDispatcher::runhas onlyjob_rx/result_tx, and aFailure{should_retry}would burn a retry), and a blocking per-task acquire on thedispatch loop would starve the neighbours the bulkhead exists to protect. Complements
max_concurrent, which stays the cluster-wide cap and costs a database read.Bugs found along the way
asyncio.CancelledErrorreported nothing — it's aBaseException, so it escapedexcept Exception. The job satRunninguntil thereaper mislabelled it a timeout.
JOB_COMPLETEDfrom native dispatch. The Rust outcome loop skipsSuccesson the grounds that the blocking wrapper emits it, so a workflow with anasync step would hang in
runningforever.Each has a regression test verified to fail without its fix.
Note: this hardens a dormant path
Native async dispatch has never actually run. The decorator registers
wrappedinthe task registry but sets
_taskito_is_asyncon theTaskWrapperit returns, so thepool's lookup always misses and every async task takes the sync branch. That single
mismatch is why
async_concurrencylooked like dead config and why the S17 deadlockwas never hit in the wild.
Activating it is two lines and demonstrably works (coroutines move to the executor
thread; concurrency goes 2 -> 20 at
workers=2, async_concurrency=20) — but thenative path is a second, independent implementation of the task lifecycle and still
lacks the blocking wrapper's queue hooks, saga compensation context,
soft_timeoutand per-item batch handling. Four of those fail silently. Activation is therefore its
own PR;
decorators.pycarries a note so the mismatch doesn't read as an oversight.Verification
cargo check --workspacex {default, postgres, redis, native-async} — cleancargo test --workspace— 185 passed; clippy and fmt cleanuv run python -m pytest tests/— 1214 passed, 4 skippedruff check+mypyovertaskito/andtests/— cleanSummary by CodeRabbit
try_*methods returningbool.try_*APIs.