diff --git a/crates/taskito-core/src/scheduler/mod.rs b/crates/taskito-core/src/scheduler/mod.rs index 105cdb0e..d402522b 100644 --- a/crates/taskito-core/src/scheduler/mod.rs +++ b/crates/taskito-core/src/scheduler/mod.rs @@ -324,6 +324,19 @@ impl Scheduler { } } + /// Forget a job that was tracked but never reached the worker channel + /// (dispatch rollback). Unlike [`Self::release_in_flight`] this never wakes + /// the poller — the slot was returned by the same thread that took it, and + /// the send just failed, so an immediate retry would only spin. + fn untrack_in_flight(&self, job_id: &str) { + if self.config.max_in_flight.is_some() { + self.in_flight + .lock() + .unwrap_or_else(|p| p.into_inner()) + .remove(job_id); + } + } + /// 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 @@ -385,6 +398,21 @@ impl Scheduler { self.shutdown.clone() } + /// `true` once every job this scheduler dispatched has had its result + /// handled — a "drain is done" signal for a binding's shutdown loop that + /// does not depend on the result channel disconnecting (which can hinge on + /// foreign-object lifetimes, e.g. a Python-held sender). `false` when + /// dispatch is unbounded: nothing is tracked, so emptiness proves nothing. + pub fn in_flight_settled(&self) -> bool { + self.config.max_in_flight.is_some() + && self + .in_flight + .lock() + .unwrap_or_else(|p| p.into_inner()) + .len() + == 0 + } + pub fn register_queue_config(&mut self, queue_name: String, config: QueueConfig) { self.queue_configs.insert(queue_name, config); } @@ -1181,6 +1209,121 @@ mod tests { ); } + #[test] + fn test_full_channel_rollback_does_not_leak_in_flight_slot() { + // A job that fails the channel hand-off must not keep its in-flight + // slot: a leaked entry shrinks the dispatch cap for the life of the + // worker and keeps the drain-exit signal from ever settling. + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); + let config = SchedulerConfig { + max_in_flight: Some(8), + ..SchedulerConfig::default() + }; + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config, None); + scheduler.storage.enqueue(make_job("full_chan")).unwrap(); + + // Capacity-1 channel pre-filled with a sentinel so the hand-off fails. + let (tx, _rx) = make_channel(1); + let sentinel = scheduler.storage.enqueue(make_job("sentinel")).unwrap(); + tx.try_send(sentinel).expect("pre-fill should succeed"); + + assert!(!scheduler.tick_dispatch(&tx), "hand-off must not succeed"); + assert_eq!( + scheduler.in_flight_remaining(), + Some(8), + "failed hand-off must not occupy an in-flight slot" + ); + } + + #[test] + fn test_closed_channel_rollback_untracks_in_flight() { + // A job tracked before a failed hand-off must be untracked by the + // rollback, or the slot leaks and `in_flight_settled` can never turn + // true again — every later shutdown would wait out the drain timeout. + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); + let config = SchedulerConfig { + max_in_flight: Some(8), + ..SchedulerConfig::default() + }; + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config, None); + scheduler.storage.enqueue(make_job("closed_chan")).unwrap(); + + let (tx, rx) = make_channel(16); + drop(rx); + assert!( + !scheduler.tick_dispatch(&tx), + "closed channel must not count as a dispatch" + ); + + assert_eq!( + scheduler.in_flight_remaining(), + Some(8), + "rolled-back job must not occupy an in-flight slot" + ); + let pending = scheduler + .storage + .list_jobs(Some(JobStatus::Pending as i32), None, None, 10, 0, None) + .unwrap(); + assert_eq!( + pending.len(), + 1, + "job returned to Pending for another worker" + ); + } + + #[test] + fn test_in_flight_settled_tracks_dispatch_lifecycle() { + // Bounded dispatch: settled when idle, unsettled while a job is out, + // settled again once its result is handled. Unbounded dispatch tracks + // nothing, so it must never report settled. + let storage = + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()); + let config = SchedulerConfig { + max_in_flight: Some(4), + ..SchedulerConfig::default() + }; + let scheduler = Scheduler::new(storage, vec!["default".to_string()], config, None); + assert!( + scheduler.in_flight_settled(), + "idle bounded scheduler is settled" + ); + + scheduler.storage.enqueue(make_job("settle")).unwrap(); + let (tx, mut rx) = make_channel(16); + while scheduler.tick_dispatch(&tx) {} + let job = rx.try_recv().expect("job dispatched"); + assert!( + !scheduler.in_flight_settled(), + "a dispatched job keeps the scheduler unsettled" + ); + + scheduler + .handle_result(JobResult::Success { + job_id: job.id.clone(), + result: None, + task_name: "settle".to_string(), + wall_time_ns: 1, + }) + .unwrap(); + assert!( + scheduler.in_flight_settled(), + "handled result settles the scheduler" + ); + + let unbounded = Scheduler::new( + StorageBackend::Sqlite(crate::storage::sqlite::SqliteStorage::in_memory().unwrap()), + vec!["default".to_string()], + SchedulerConfig::default(), + None, + ); + assert!( + !unbounded.in_flight_settled(), + "unbounded dispatch can never prove it settled" + ); + } + /// 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 = @@ -1442,17 +1585,16 @@ mod tests { } #[test] - fn test_try_dispatch_reschedules_on_full_channel() { - // Same regression when the channel is *full* rather than closed — - // the worker pool is behind but still alive. The job must come back - // to Pending so the next tick has a chance to dispatch it once the - // pool drains. + fn test_try_dispatch_skips_when_channel_full() { + // When the channel is *full* rather than closed, the poller must not + // dequeue at all: a claim taken now could only roll back on `try_send`, + // churning storage for nothing. The job stays Pending, untouched, and + // dispatch resumes once the pool drains. let scheduler = test_scheduler(); let job = scheduler.storage.enqueue(make_job("ch_full_task")).unwrap(); - // Capacity-1 channel pre-filled with a sentinel job. The poller's - // `try_send` will see `TrySendError::Full`. - let (tx, _rx) = make_channel(1); + // Capacity-1 channel pre-filled with a sentinel job. + let (tx, mut rx) = make_channel(1); let sentinel = scheduler.storage.enqueue(make_job("sentinel")).unwrap(); tx.try_send(sentinel).expect("pre-fill should succeed"); @@ -1461,13 +1603,23 @@ mod tests { let after = scheduler.storage.get_job(&job.id).unwrap().unwrap(); assert_eq!(after.status, JobStatus::Pending); - assert!(after.scheduled_at > now_millis()); + assert_eq!( + after.scheduled_at, job.scheduled_at, + "a skipped job keeps its schedule — it was never touched" + ); let claims = scheduler .storage .list_claims_by_worker("scheduler") .unwrap(); assert!(!claims.contains(&job.id)); + + // Draining the channel restores dispatch. + let _ = rx.try_recv().expect("sentinel was in the channel"); + assert!( + scheduler.tick_dispatch(&tx), + "freed capacity resumes dispatch" + ); } #[test] diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index f42912f0..1aabaf50 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/crates/taskito-core/src/scheduler/poller.rs @@ -93,6 +93,14 @@ impl Scheduler { return Ok(false); } + // A full channel means every dequeue would claim and then immediately + // roll back on `try_send` — pure storage churn. Wait for the pool to + // drain; `try_send` below still guards the race where it fills between + // this check and the hand-off. + if job_tx.capacity() == 0 { + return Ok(false); + } + let job = match self .storage .dequeue_from(&active_queues, now, self.namespace.as_deref())? @@ -126,8 +134,13 @@ impl Scheduler { // Size the batch to the worker pool's free capacity so we never claim // more jobs than the channel can immediately accept. `try_send` still // guards each hand-off, but pre-sizing avoids needless claim/rollback - // churn. Always claim at least one. - let mut budget = self.config.batch_size.min(job_tx.capacity().max(1)); + // churn — including the full-channel case, where any claim would + // immediately roll back. + let capacity = job_tx.capacity(); + if capacity == 0 { + return Ok(false); + } + let mut budget = self.config.batch_size.min(capacity); // Never claim past the in-flight cap: a drained batch that outran the // workers would mark the surplus `Running` and starve peer schedulers. @@ -289,15 +302,20 @@ impl Scheduler { // must reverse the claim — otherwise the job sits in `Running` until // the stale-reaper times it out, which surfaces as a *timeout* in // metrics and middleware (wrong outcome for a job that never ran). + // + // Track BEFORE the send: once the job is in the channel its result can + // arrive on the drain thread at any moment, and a release that races + // ahead of the insert would leak the slot forever. Tracking first means + // the entry always exists by the time a result can be handled; the + // failure arms untrack what never left. let job_id = job.id.clone(); let task_name = job.task_name.clone(); + self.track_in_flight(&job_id, &task_name); match job_tx.try_send(job) { - Ok(()) => { - self.track_in_flight(&job_id, &task_name); - Ok(true) - } + Ok(()) => Ok(true), Err(TrySendError::Full(job)) => { warn!("worker channel full; rescheduling job {job_id} (worker pool is behind)",); + self.untrack_in_flight(&job_id); counts.release(&job.task_name, &job.queue); self.rollback_claim_and_reschedule( &job_id, @@ -309,6 +327,7 @@ impl Scheduler { warn!( "worker channel closed; rescheduling job {job_id} (worker pool shutting down)", ); + self.untrack_in_flight(&job_id); counts.release(&job.task_name, &job.queue); self.rollback_claim_and_reschedule( &job_id, diff --git a/crates/taskito-python/src/py_queue/worker.rs b/crates/taskito-python/src/py_queue/worker.rs index 6b8fa49a..58f88a01 100644 --- a/crates/taskito-python/src/py_queue/worker.rs +++ b/crates/taskito-python/src/py_queue/worker.rs @@ -282,15 +282,25 @@ impl PyQueue { #[allow(unused_variables)] let async_concurrency = async_concurrency.max(1) as usize; + let use_prefork = pool.as_deref() == Some("prefork"); + // 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. + // With the native pool, sync and async work draw on separate budgets: + // blocking tasks are bounded by `num_workers` threads, coroutines by the + // executor's `async_concurrency` semaphore — so the cap is their sum. + // Everywhere else — non-native builds, and prefork, which builds no + // executor — every task is bounded by `num_workers`; adding + // `async_concurrency` there would claim jobs nothing can execute — they + // would sit Running in the channel waiting for a worker slot. + #[cfg(feature = "native-async")] + let max_in_flight = if use_prefork { + self.num_workers + } else { + self.num_workers + async_concurrency + }; + #[cfg(not(feature = "native-async"))] let max_in_flight = self.num_workers; let scheduler_config = SchedulerConfig { @@ -480,43 +490,22 @@ impl PyQueue { pool.as_deref(), ); - // Create the async executor for native async tasks (if feature enabled) - #[cfg(feature = "native-async")] - let async_executor = { - let sender = crate::native_async::PyResultSender::new(result_tx.clone()); - Python::attach(|py| -> PyResult>> { - let sender_obj = pyo3::Py::new(py, sender)?; - let mod_ = py.import("taskito.async_support.executor")?; - let cls = mod_.getattr("AsyncTaskExecutor")?; - let context_mod = py.import("taskito.context")?; - let queue_ref = context_mod.getattr("_queue_ref")?; - let executor = cls.call1(( - sender_obj, - registry_arc.clone_ref(py), - queue_ref, - async_concurrency, - ))?; - executor.call_method0("start")?; - Ok(Arc::new(executor.unbind())) - })? - }; - // Build the dispatcher up front for the prefork case so we can install // it on the queue before the run loop starts — request_cancel relies on // the install to deliver out-of-band cancel signals to child processes. // // For in-process pools (native-async, classic async) `notify_cancel` is // a no-op — running tasks observe cancellation via the storage flag — - // so we deliberately do NOT install the dispatcher on `self`. - // Installing it would keep an `Arc>` reference (the async - // executor / PyResultSender chain) alive on the parent thread until - // `set_dispatcher(None)` runs after `runtime_handle.join()`, deadlocking - // shutdown: the drain loop waits for the result channel to disconnect, - // which can't happen until PyResultSender drops, which can't happen - // until both `async_executor` Arc refs drop — and the second ref is - // exactly what `self.dispatcher` was holding. + // so we don't install the dispatcher on `self`. let num_workers = self.num_workers; - let use_prefork = pool.as_deref() == Some("prefork"); + + // Held on this thread so shutdown can stop the executor's event loop + // once the drain completes. Safe to keep here: the drain exits when + // in-flight work settles, not when the result channel disconnects, so + // this reference cannot deadlock shutdown. + #[cfg(feature = "native-async")] + let mut async_executor_for_shutdown: Option>> = None; + let dispatcher_for_run: Arc = if use_prefork { let pool_arc: Arc = Arc::new( crate::prefork::PreforkPool::new(num_workers, app_path.unwrap_or_default()), @@ -526,6 +515,26 @@ impl PyQueue { } else { #[cfg(feature = "native-async")] { + // The executor is built only for the pool that uses it — + // prefork must not hold a PyResultSender (a result-channel + // clone) it would never send on. + let sender = crate::native_async::PyResultSender::new(result_tx.clone()); + let async_executor = Python::attach(|py| -> PyResult>> { + let sender_obj = pyo3::Py::new(py, sender)?; + let mod_ = py.import("taskito.async_support.executor")?; + let cls = mod_.getattr("AsyncTaskExecutor")?; + let context_mod = py.import("taskito.context")?; + let queue_ref = context_mod.getattr("_queue_ref")?; + let executor = cls.call1(( + sender_obj, + registry_arc.clone_ref(py), + queue_ref, + async_concurrency, + ))?; + executor.call_method0("start")?; + Ok(Arc::new(executor.unbind())) + })?; + async_executor_for_shutdown = Some(async_executor.clone()); let pool_arc: Arc = Arc::new(crate::native_async::NativeAsyncPool::new( num_workers, @@ -548,6 +557,15 @@ impl PyQueue { #[cfg(feature = "mesh")] let mesh_worker_id = worker_id.clone(); + // Mesh-stolen jobs enter the channel through the bridge, bypassing + // `finish_dispatch` — the scheduler's in-flight tracking never sees + // them, so the settled-work drain exit below is only sound without + // mesh. + #[cfg(feature = "mesh")] + let mesh_enabled = mesh_config.is_some(); + #[cfg(not(feature = "mesh"))] + let mesh_enabled = false; + // Captured for the channel-based (Postgres/Redis) wake-source setup // inside the runtime. Gated to the listener-bearing backends so the // default and SQLite-only builds have no unused binding. @@ -664,6 +682,23 @@ impl PyQueue { let drain_timeout = std::time::Duration::from_secs(drain_timeout_secs.unwrap_or(30)); + // Refcount-independent "drain is done" signal. Waiting for the result + // channel to *disconnect* requires every sender clone to drop, and the + // async executor's PyResultSender is owned by a Python object — a failed + // task's traceback can pin it (any hook that retains the exception + // reaches the frame that owns the sender), stalling shutdown for the + // full drain timeout. Instead: the runtime thread finishing proves the + // scheduler returned (no dispatch in progress) and the pool drained its + // queue (the runtime's drop flushes running blocking tasks), and a + // settled in-flight map proves every dispatched job's result was + // handled — this loop is the only consumer, so nothing can still be + // pending. + let work_settled = || { + !mesh_enabled + && runtime_handle.is_finished() + && scheduler_for_results.in_flight_settled() + }; + loop { // Release GIL for one iteration of result polling let action = py.detach(|| { @@ -721,7 +756,12 @@ impl PyQueue { } } } - PollAction::Continue => continue, + PollAction::Continue => { + if work_settled() { + break; + } + continue; + } PollAction::Done => break, PollAction::Shutdown => unreachable!(), } @@ -746,13 +786,32 @@ impl PyQueue { } } } - PollAction::Continue => continue, + PollAction::Continue => { + // Without this, a runtime that died with all work settled + // would leave the loop polling a channel that can never + // disconnect (the executor still holds a sender). + if work_settled() { + break; + } + continue; + } PollAction::Done => break, } } let _ = runtime_handle.join(); + // Stop the async executor's event loop and drop its result sender. + // Best effort — shutdown must finish even if the interpreter is + // tearing down. Without this the daemon thread and the sender live + // until the executor object is garbage collected. + #[cfg(feature = "native-async")] + if let Some(executor) = async_executor_for_shutdown { + if let Err(e) = executor.call_method0(py, "stop") { + log::warn!("[taskito] async executor stop failed: {e}"); + } + } + // Clear the dispatcher reference so post-shutdown cancel requests // become no-ops instead of forwarding to a torn-down pool. self.set_dispatcher(None); diff --git a/sdks/python/taskito/async_support/executor.py b/sdks/python/taskito/async_support/executor.py index 60acc553..0ec63e1a 100644 --- a/sdks/python/taskito/async_support/executor.py +++ b/sdks/python/taskito/async_support/executor.py @@ -268,10 +268,19 @@ def _check_retry(self, task_name: str, exc: Exception) -> bool: return True def stop(self) -> None: - """Stop the executor's event loop and join the thread.""" - if self._loop is not None and self._loop.is_running(): + """Stop the executor's event loop, join the thread, release the sender.""" + # `not is_closed()` rather than `is_running()`: right after `start()` + # the thread may not have entered `run_forever()` yet, and skipping the + # stop here would leave it running forever once it does. + if self._loop is not None and not self._loop.is_closed(): self._loop.call_soon_threadsafe(self._loop.stop) if self._thread is not None: self._thread.join(timeout=5) if self._thread.is_alive(): logger.warning("Async executor thread did not stop within 5s timeout") + return + # The sender wraps a clone of the worker's result channel, and a pinned + # frame (a retained exception's traceback) can keep this object alive + # long past shutdown — so the release must not wait for GC. Only safe + # once the loop thread is down: no coroutine can report after that. + self._sender = None diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index f4f46bcd..ae1d5843 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -422,21 +422,16 @@ def decorator(fn: Callable) -> TaskWrapper: self._task_soft_timeouts[task_name] = soft_timeout wrapped = self._wrap_task(fn, task_name) - # NOTE: `_taskito_is_async`/`_taskito_async_fn` are deliberately NOT - # set on `wrapped`. The pool reads them off this registry entry to - # choose native dispatch, so setting them here activates that path. - # Both must move together when it is activated — the pool reads the - # first to route, the executor reads the second to find the coroutine. - # - # The lifecycle gaps that used to block activation are closed: both - # paths now run `run_lifecycle`. What still blocks it is shutdown. - # The worker's drain loop ends when the result channel disconnects, - # which needs the executor's `PyResultSender` to drop — but a failed - # task's traceback keeps the frame that owns it alive, so anything - # holding the exception (an on_failure hook that collects errors is - # enough) makes every shutdown wait out the full drain timeout: 30s - # by default, past a typical container's grace period. Decoupling the - # drain from Python object lifetime has to land first. + # Native async dispatch keys off this registry entry: the pool reads + # `_taskito_is_async` to route, the executor reads `_taskito_async_fn` + # to find the coroutine — so both must live on `wrapped`, not only on + # the TaskWrapper returned to the caller. On builds without the + # native pool the flags are inert and async tasks run on the + # blocking path via `run_maybe_async`. + is_async = inspect.iscoroutinefunction(fn) + wrapped._taskito_is_async = is_async # type: ignore[attr-defined] + if is_async: + wrapped._taskito_async_fn = fn # type: ignore[attr-defined] self._task_registry[task_name] = wrapped cb_threshold = None @@ -489,8 +484,7 @@ def decorator(fn: Callable) -> TaskWrapper: # Preserve function metadata functools.update_wrapper(wrapper, fn) - # Mark async status for native async dispatch - is_async = inspect.iscoroutinefunction(fn) + # Mirror the async markers for introspection on the returned handle. wrapper._taskito_is_async = is_async if is_async: wrapper._taskito_async_fn = fn diff --git a/sdks/python/tests/worker/test_in_flight_bound.py b/sdks/python/tests/worker/test_in_flight_bound.py index 2f8a6f0f..50b840cb 100644 --- a/sdks/python/tests/worker/test_in_flight_bound.py +++ b/sdks/python/tests/worker/test_in_flight_bound.py @@ -45,6 +45,15 @@ def slow() -> int: 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" + # The failure mode guarded here is the whole backlog going Running at once + # (an uncapped scheduler claims all 40). The legitimate ceiling is the + # worker's pipeline, every stage of which is a bounded channel or pool: + # 2 executing + 4 job-channel + 1 held by the pool loop + # + 4 result-channel + ~5 in one finalize batch (Running until committed) + # = 16. On native-async builds the in-flight cap is num_workers + + # async_concurrency, so for sync work the channel is the binding limiter, + # not the cap — observed peaks are 7-9, the rest is commit lag headroom. + assert peak_running <= 16, ( + f"worker claimed {peak_running} jobs — pipeline ceiling is 16, " + f"anything near the {job_count}-job backlog is the S01 over-claim" + ) diff --git a/sdks/python/tests/worker/test_native_async.py b/sdks/python/tests/worker/test_native_async.py index 25ad210e..899595ce 100644 --- a/sdks/python/tests/worker/test_native_async.py +++ b/sdks/python/tests/worker/test_native_async.py @@ -212,6 +212,60 @@ def test_async_executor_lifecycle() -> None: executor.stop() +def test_stop_before_loop_runs_still_stops_the_thread() -> None: + """stop() racing start() must not leave the loop thread running forever. + + Rebuilds start()'s state with the thread held before ``run_forever`` so the + loop is provably not running when stop() checks it — the race window that + ``is_running()`` misses and ``is_closed()`` does not. + """ + from taskito.async_support.executor import AsyncTaskExecutor + + executor = AsyncTaskExecutor(MagicMock(), {}, MagicMock(), max_concurrency=10) + executor._loop = asyncio.new_event_loop() + entered = threading.Event() + checked = threading.Event() + + def held_run_forever() -> None: + entered.wait(timeout=10) + assert executor._loop is not None + executor._loop.run_forever() + + class JoinSignalsCheckDone(threading.Thread): + """stop() joins only after its loop check, so join marks the check done.""" + + def join(self, timeout: float | None = None) -> None: + checked.set() + super().join(timeout) + + executor._thread = JoinSignalsCheckDone(target=held_run_forever, daemon=True) + executor._thread.start() + + stopper = threading.Thread(target=executor.stop) + stopper.start() + # Release the thread into run_forever only once stop() has provably made + # its loop check against a not-yet-running loop. + assert checked.wait(timeout=10), "stop() never reached its join" + entered.set() + stopper.join(timeout=10) + + assert not stopper.is_alive(), "stop() never returned" + assert not executor._thread.is_alive(), "loop thread kept running forever" + assert executor._sender is None + + +def test_stop_releases_result_sender() -> None: + """stop() must drop the sender itself — a pinned frame can keep the + executor alive indefinitely, so the release cannot wait for GC.""" + from taskito.async_support.executor import AsyncTaskExecutor + + executor = AsyncTaskExecutor(MagicMock(), {}, MagicMock(), max_concurrency=10) + executor.start() + executor.stop() + assert executor._sender is None + assert executor._thread is not None and not executor._thread.is_alive() + + def test_async_executor_submit_and_execute(poll_until: PollUntil) -> None: """Basic async task produces correct result via executor.""" from taskito.async_support.executor import AsyncTaskExecutor diff --git a/sdks/python/tests/worker/test_shutdown_drain.py b/sdks/python/tests/worker/test_shutdown_drain.py new file mode 100644 index 00000000..7038e001 --- /dev/null +++ b/sdks/python/tests/worker/test_shutdown_drain.py @@ -0,0 +1,77 @@ +"""Worker shutdown must not depend on Python object lifetimes.""" + +import threading +import time +from pathlib import Path +from typing import Any + +import pytest + +from taskito import Queue, _taskito + +PollUntil = Any # the conftest fixture's runtime type + +requires_native_async = pytest.mark.skipif( + not hasattr(_taskito, "PyResultSender"), + reason="wheel built without the native-async feature", +) + + +@requires_native_async +def test_async_task_runs_on_native_executor(tmp_path: Path, poll_until: PollUntil) -> None: + """Async tasks must dispatch to the executor's event loop, not a blocking thread.""" + queue = Queue(db_path=str(tmp_path / "native.db"), workers=2) + seen: list[str] = [] + + @queue.task(name="where_am_i") + async def where_am_i() -> None: + seen.append(threading.current_thread().name) + + where_am_i.delay() + thread = threading.Thread(target=queue.run_worker, daemon=True) + thread.start() + try: + poll_until(lambda: len(seen) == 1, timeout=15, message="task never ran") + finally: + queue.shutdown() + thread.join(timeout=10) + + assert not thread.is_alive(), "worker did not shut down" + assert seen == ["taskito-async-executor"] + + +@requires_native_async +def test_failed_async_task_does_not_stall_shutdown(tmp_path: Path, poll_until: PollUntil) -> None: + """A retained exception must not make shutdown wait out the drain timeout. + + A failed task's traceback reaches the executor frame that owns the async + result sender, so a hook that keeps the exception held the result channel + open — and a drain that waits for channel disconnect then burns the full + 30s, past a typical container's grace period. + """ + queue = Queue(db_path=str(tmp_path / "stall.db"), workers=2) + errors: list[BaseException] = [] + + @queue.on_failure + def collect(task_name: str, args: Any, kwargs: Any, error: BaseException) -> None: + errors.append(error) + + @queue.task(name="boom", max_retries=0) + async def boom() -> None: + raise RuntimeError("kept alive by the hook") + + boom.delay() + thread = threading.Thread(target=queue.run_worker, daemon=True) + thread.start() + try: + poll_until(lambda: len(errors) == 1, timeout=15, message="failure hook never fired") + finally: + queue.shutdown() + start = time.monotonic() + thread.join(timeout=20) + elapsed = time.monotonic() - start + + assert not thread.is_alive(), "worker never shut down" + assert elapsed < 5, ( + f"shutdown took {elapsed:.2f}s — the drain is waiting on the retained exception" + )