From 94d4cc1c585c06ef860c014ee4c2957635196818 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:22:54 +0530 Subject: [PATCH 01/13] refactor(python): extract the task lifecycle both paths share MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The blocking wrapper and the async executor each implement the lifecycle a task runs through, and they have drifted: the native copy silently lost the queue hooks, the saga compensation context, soft_timeout and per-item batch handling. Patching the gaps into the executor would make a second copy that drifts again. Move the body to taskito/task_lifecycle.py and have the wrapper drive it. Pure move — the callers keep only what genuinely differs, which is where arguments come from and where the result goes. current_job needs no special handling: its context lookup already reads the contextvar before the thread-local. soft_timeout moves into a _task_soft_timeouts dict, mirroring the eleven other per-task options. It was the only one reachable solely through the wrapper's closure, which is why the executor could never honour it. _clear_context stays with the wrapper: it is edge state, and in the shared body it would null out the thread-local of whichever thread the executor loop runs on. No behaviour change — 1217 tests unchanged. --- sdks/python/taskito/app.py | 4 + sdks/python/taskito/mixins/decorators.py | 236 ++------------- sdks/python/taskito/task_lifecycle.py | 275 ++++++++++++++++++ .../observability/test_task_lifecycle_logs.py | 2 +- 4 files changed, 307 insertions(+), 210 deletions(-) create mode 100644 sdks/python/taskito/task_lifecycle.py diff --git a/sdks/python/taskito/app.py b/sdks/python/taskito/app.py index 2fd518b1..255bd5c2 100644 --- a/sdks/python/taskito/app.py +++ b/sdks/python/taskito/app.py @@ -278,6 +278,10 @@ def __init__( self._task_idempotent: dict[str, bool] = {} self._task_compensates: dict[str, str] = {} self._task_batch_configs: dict[str, BatchConfig] = {} + # Per-task cooperative timeout, mirrored here like every other per-task + # option so the lifecycle can reach it through the queue rather than a + # closure only the blocking wrapper had. + self._task_soft_timeouts: dict[str, float] = {} self._batch_accumulator: BatchAccumulator | None = None self._global_middleware: list[TaskMiddleware] = middleware or [] self._task_middleware: dict[str, list[TaskMiddleware]] = {} diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 4dce800a..19b880cb 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -16,26 +16,14 @@ from taskito._taskito import PyTaskConfig from taskito.async_support.helpers import run_maybe_async from taskito.batching.config import BatchConfig -from taskito.batching.item_result import ( - BatchPartialFailureError, - BatchResultTypeError, - is_batch_item_result_list, -) from taskito.codecs import CodecSerializer -from taskito.context import _clear_context, current_job +from taskito.context import _clear_context from taskito.dashboard.middleware_store import MiddlewareDisableStore -from taskito.events import EventType -from taskito.exceptions import TaskCancelledError from taskito.inject import Inject, _InjectAlias -from taskito.interception.reconstruct import reconstruct_args from taskito.middleware import middleware_key from taskito.predicates.core import coerce_predicate -from taskito.proxies import cleanup_proxies, reconstruct_proxies from taskito.task import TaskWrapper -from taskito.workflows.saga.context import ( - _reset_compensation_context, - _set_compensation_context, -) +from taskito.task_lifecycle import run_lifecycle if TYPE_CHECKING: from taskito.codecs import PayloadCodec @@ -51,29 +39,11 @@ logger = logging.getLogger("taskito") -# Cap the result repr length in the "succeeded" log so a task returning a -# large structure can't blow out the worker output. Matches Celery's -# `CELERYD_TASK_LOG_FORMAT` truncation in spirit. -_MAX_RESULT_REPR = 80 - # How long a cached middleware chain stays valid without a version bump. Bounds # the worst-case lag for an out-of-process dashboard disable change. _MW_CHAIN_TTL = 1.0 -def _safe_result_repr(value: Any) -> str: - """Render a task return value for the success log, bounded and crash-proof.""" - if value is None: - return "None" - try: - text = repr(value) - except Exception: - return "" - if len(text) > _MAX_RESULT_REPR: - return text[: _MAX_RESULT_REPR - 1] + "…" - return text - - def _resolve_module_name(module_name: str) -> str: """Resolve __main__ to the actual module name.""" if module_name != "__main__": @@ -104,6 +74,7 @@ class QueueDecoratorMixin: _task_idempotent: dict[str, bool] _task_compensates: dict[str, str] _task_batch_configs: dict[str, Any] + _task_soft_timeouts: dict[str, float] _task_middleware: dict[str, list[TaskMiddleware]] _task_retry_filters: dict[str, dict[str, list[type[Exception]]]] _task_inject_map: dict[str, list[str]] @@ -168,191 +139,36 @@ def _get_middleware_chain(self, task_name: str) -> list[TaskMiddleware]: self._mw_chain_cache[task_name] = (chain, version, time.monotonic()) return chain - def _wrap_task( - self, fn: Callable, task_name: str, soft_timeout: float | None = None - ) -> Callable: - """Wrap a task function with hooks, middleware, and job context.""" - hooks = self._hooks + def _wrap_task(self, fn: Callable, task_name: str) -> Callable: + """The blocking entry point Rust calls: drive the shared lifecycle to completion. + + Rust hands over deserialized args and expects a plain value back (or an + exception), so the coroutine is driven here rather than awaited. It runs + on a `spawn_blocking` thread with no event loop of its own, which is what + makes that safe — `run_maybe_async` raises if one is already running. + """ queue_ref = self + is_async = inspect.iscoroutinefunction(fn) @functools.wraps(fn) def wrapper(*args: Any, **kwargs: Any) -> Any: - job_id = current_job.id - logger.info("Task %s[%s] received", task_name, job_id) - started_at = time.perf_counter() - - # Worker-dispatch predicate gate. Evaluated on the raw - # deserialized payload (before arg/proxy reconstruction) so - # re-enqueue on defer can round-trip cleanly. - if task_name in queue_ref._task_predicates: - queue_ref._apply_dispatch_predicate( - task_name=task_name, - args=args, - kwargs=kwargs, - job_id=current_job.id, - queue_name=current_job.queue_name, - retry_count=current_job.retry_count, - ) - - # Reconstruct intercepted arguments (CONVERT markers → original types) - redirects: dict[str, str] = {} - if queue_ref._interceptor is not None: - args, kwargs, redirects = reconstruct_args(args, kwargs) - - # Reconstruct proxy markers (PROXY → live objects) - proxy_cleanup: list[Any] = [] - if queue_ref._proxy_registry is not None and not queue_ref._test_mode_active: - args, kwargs, proxy_cleanup = reconstruct_proxies( - args, - kwargs, - queue_ref._proxy_registry, - signing_secret=queue_ref._recipe_signing_key, - max_timeout=queue_ref._max_reconstruction_timeout, - metrics=queue_ref._proxy_metrics, - ) - - # Inject resources from runtime - release_callbacks: list[Any] = [] - runtime = queue_ref._resource_runtime - if runtime is not None: - # From explicit inject=["db"] on task decorator - for res_name in queue_ref._task_inject_map.get(task_name, []): - if res_name not in kwargs: - instance, release = runtime.acquire_for_task(res_name) - kwargs[res_name] = instance - if release is not None: - release_callbacks.append(release) - # From interception REDIRECT markers - for kwarg_name, resource_name in redirects.items(): - instance, release = runtime.acquire_for_task(resource_name) - kwargs[kwarg_name] = instance - if release is not None: - release_callbacks.append(release) - - middleware_chain = queue_ref._get_middleware_chain(task_name) - - # Set soft timeout on context if configured - if soft_timeout is not None: - current_job._set_soft_timeout(soft_timeout) - - # Run middleware before hooks (skipping middlewares whose - # predicate filter excludes this job) - completed_mw: list[Any] = [] - for mw in middleware_chain: - if not mw._should_apply(current_job): - continue - try: - mw.before(current_job) - completed_mw.append(mw) - except Exception: - logger.exception("middleware before() error") - - for hook in hooks["before_task"]: - hook(task_name, args, kwargs) - - # Saga compensation context: if this job was dispatched by the - # saga orchestrator, look up the in-memory CompensationContext - # and push it onto a contextvar so the compensator body can - # call ``current_compensation_context()`` to introspect the - # forward execution. Pop in the ``finally`` below. - comp_ctx_token = None - tracker = getattr(queue_ref, "_workflow_tracker", None) - saga = getattr(tracker, "_saga", None) if tracker is not None else None - if saga is not None and saga.is_compensation_job(job_id): - comp_ctx = saga.take_compensation_context(job_id) - if comp_ctx is not None: - comp_ctx_token = _set_compensation_context(comp_ctx) - - error = None - result = None try: - ret = run_maybe_async(fn(*args, **kwargs)) - # Per-item batch result: when the task is batched with - # per_item_results=True, enforce the return type contract - # and surface partial failures via BatchPartialFailureError - # so the existing retry/DLQ machinery applies. - batch_cfg = queue_ref._task_batch_configs.get(task_name) - if batch_cfg is not None and getattr(batch_cfg, "per_item_results", False): - if not is_batch_item_result_list(ret): - raise BatchResultTypeError( - f"task {task_name!r} declares per_item_results=True but " - f"returned {type(ret).__name__} instead of " - f"list[BatchItemResult]" - ) - failed = [item for item in ret if item.status == "failure"] - if failed: - # Pass the result through so the worker can still - # store it (per-item outcomes are stored verbatim); - # the existing on_failure path emits JOB_FAILED. - result = ret - raise BatchPartialFailureError(failed_items=failed) - result = ret - except Exception as exc: - error = exc - elapsed = time.perf_counter() - started_at - if isinstance(exc, TaskCancelledError): - logger.info( - "Task %s[%s] cancelled in %.3fs: %s", + # The mixin is a Queue once composed; mypy only sees the mixin. + return run_maybe_async( + run_lifecycle( + queue_ref, # type: ignore[arg-type] task_name, - job_id, - elapsed, - exc, + fn, + args, + kwargs, + is_async=is_async, ) - else: - logger.error( - "Task %s[%s] raised %s in %.3fs: %r", - task_name, - job_id, - type(exc).__name__, - elapsed, - exc, - ) - for hook in hooks["on_failure"]: - hook(task_name, args, kwargs, exc) - raise - else: - elapsed = time.perf_counter() - started_at - logger.info( - "Task %s[%s] succeeded in %.3fs: %s", - task_name, - job_id, - elapsed, - _safe_result_repr(result), ) - for hook in hooks["on_success"]: - hook(task_name, args, kwargs, result) - return result finally: - # Pop the saga compensation context (no-op outside saga flow). - if comp_ctx_token is not None: - _reset_compensation_context(comp_ctx_token) - # Release task/request-scoped resources - for release_fn in release_callbacks: - try: - release_fn() - except Exception: - logger.exception("resource release error") - # Clean up reconstructed proxies (LIFO order) - cleanup_proxies(proxy_cleanup, metrics=queue_ref._proxy_metrics) - for hook in hooks["after_task"]: - hook(task_name, args, kwargs, result, error) - # Run middleware after hooks (only those whose before() succeeded) - for mw in completed_mw: - try: - mw.after(current_job, result, error) - except Exception: - logger.exception("middleware after() error") - # Emit job lifecycle events - event_payload: dict[str, Any] = { - "task_name": task_name, - "job_id": current_job.id, - "queue": current_job.queue_name, - } - if error is not None: - event_payload["error"] = str(error) - queue_ref._emit_event(EventType.JOB_FAILED, event_payload) - else: - queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) + # Rust clears the context too, but the classic pool relies on this + # one. It is edge state, not lifecycle state, so it stays out of + # the shared body — there it would null out the thread-local of + # whichever thread the executor loop happens to run on. _clear_context() return wrapper @@ -600,7 +416,9 @@ def decorator(fn: Callable) -> TaskWrapper: self._task_inject_map[task_name] = final_inject # Wrap the function with hooks, middleware, and context - wrapped = self._wrap_task(fn, task_name, soft_timeout) + if soft_timeout is not None: + self._task_soft_timeouts[task_name] = soft_timeout + wrapped = self._wrap_task(fn, task_name) # NOTE: `_taskito_is_async` is deliberately NOT set on `wrapped`. # The pool reads that flag off this registry entry to choose native diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py new file mode 100644 index 00000000..7ecbe3af --- /dev/null +++ b/sdks/python/taskito/task_lifecycle.py @@ -0,0 +1,275 @@ +"""The lifecycle one task runs through, shared by both dispatch paths. + +A task reaches the worker one of two ways: the blocking path, where Rust calls +the wrapped function on a pool thread, or native async dispatch, where the +executor awaits the coroutine on its own event loop. Everything between "we have +the arguments" and "we have a result" is identical, and lives here. + +It lives here rather than in either caller because the two used to reimplement it +separately, and drifted: the native copy silently lost the queue hooks, the saga +compensation context, ``soft_timeout`` and per-item batch handling. One body +means a new lifecycle concern cannot reach one path and miss the other. + +The callers keep only what genuinely differs — how arguments arrive, and how the +result is delivered: + +=================== ========================== ============================= + blocking native async +=================== ========================== ============================= +payload/result Rust calls the queue's the executor calls them + (de)serializers directly +delivery return/raise to Rust ``PyResultSender`` +job context ``_set_context`` (thread) ``set_async_context`` (ctxvar) +=================== ========================== ============================= + +The context split needs no handling here: ``JobContext._require_context`` reads +the contextvar first and falls back to the thread-local, so ``current_job`` +resolves on either path. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +from taskito.async_support.helpers import run_maybe_async +from taskito.batching.item_result import ( + BatchPartialFailureError, + BatchResultTypeError, + is_batch_item_result_list, +) +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 +from taskito.workflows.saga.context import ( + _reset_compensation_context, + _set_compensation_context, +) + +if TYPE_CHECKING: + from taskito.app import Queue + +logger = logging.getLogger("taskito") + +_MAX_RESULT_REPR = 80 + + +def _safe_result_repr(value: Any) -> str: + """Render a task return value for the success log, bounded and crash-proof.""" + if value is None: + return "None" + try: + text = repr(value) + except Exception: + return "" + if len(text) > _MAX_RESULT_REPR: + return text[: _MAX_RESULT_REPR - 1] + "…" + return text + + +async def run_lifecycle( + queue_ref: Queue, + task_name: str, + fn: Callable, + args: Any, + kwargs: Any, + *, + is_async: bool, +) -> Any: + """Run one task end to end: gates, reconstruction, injection, middleware, hooks. + + Returns the task's raw return value, or raises what it raised — delivery + belongs to the caller, which knows whether it is answering Rust or a result + sender. ``args``/``kwargs`` arrive already deserialized. + + ``is_async`` selects how the body is invoked: awaited directly when the task + is a coroutine function (native dispatch), or driven to completion by + ``run_maybe_async`` when it is not. A blocking caller must therefore drive + the whole coroutine with ``run_maybe_async``; a native one awaits it. + """ + job_id = current_job.id + logger.info("Task %s[%s] received", task_name, job_id) + started_at = time.perf_counter() + + # Worker-dispatch predicate gate. Evaluated on the raw deserialized payload + # (before arg/proxy reconstruction) so re-enqueue on defer can round-trip + # cleanly. + if task_name in queue_ref._task_predicates: + queue_ref._apply_dispatch_predicate( + task_name=task_name, + args=args, + kwargs=kwargs, + job_id=current_job.id, + queue_name=current_job.queue_name, + retry_count=current_job.retry_count, + ) + + # Reconstruct intercepted arguments (CONVERT markers → original types) + redirects: dict[str, str] = {} + if queue_ref._interceptor is not None: + args, kwargs, redirects = reconstruct_args(args, kwargs) + + # Reconstruct proxy markers (PROXY → live objects) + proxy_cleanup: list[Any] = [] + if queue_ref._proxy_registry is not None and not queue_ref._test_mode_active: + args, kwargs, proxy_cleanup = reconstruct_proxies( + args, + kwargs, + queue_ref._proxy_registry, + signing_secret=queue_ref._recipe_signing_key, + max_timeout=queue_ref._max_reconstruction_timeout, + metrics=queue_ref._proxy_metrics, + ) + + # Inject resources from runtime + release_callbacks: list[Any] = [] + runtime = queue_ref._resource_runtime + if runtime is not None: + # From explicit inject=["db"] on task decorator + for res_name in queue_ref._task_inject_map.get(task_name, []): + if res_name not in kwargs: + instance, release = runtime.acquire_for_task(res_name) + kwargs[res_name] = instance + if release is not None: + release_callbacks.append(release) + # From interception REDIRECT markers + for kwarg_name, resource_name in redirects.items(): + instance, release = runtime.acquire_for_task(resource_name) + kwargs[kwarg_name] = instance + if release is not None: + release_callbacks.append(release) + + middleware_chain = queue_ref._get_middleware_chain(task_name) + hooks = queue_ref._hooks + + # Set soft timeout on context if configured + soft_timeout = queue_ref._task_soft_timeouts.get(task_name) + if soft_timeout is not None: + current_job._set_soft_timeout(soft_timeout) + + # Run middleware before hooks (skipping middlewares whose predicate filter + # excludes this job) + completed_mw: list[Any] = [] + for mw in middleware_chain: + if not mw._should_apply(current_job): + continue + try: + mw.before(current_job) + completed_mw.append(mw) + except Exception: + logger.exception("middleware before() error") + + for hook in hooks["before_task"]: + hook(task_name, args, kwargs) + + # Saga compensation context: if this job was dispatched by the saga + # orchestrator, look up the in-memory CompensationContext and push it onto a + # contextvar so the compensator body can call + # ``current_compensation_context()`` to introspect the forward execution. + # Pop in the ``finally`` below. + comp_ctx_token = None + tracker = getattr(queue_ref, "_workflow_tracker", None) + saga = getattr(tracker, "_saga", None) if tracker is not None else None + if saga is not None and saga.is_compensation_job(job_id): + comp_ctx = saga.take_compensation_context(job_id) + if comp_ctx is not None: + comp_ctx_token = _set_compensation_context(comp_ctx) + + error = None + result = None + try: + called = fn(*args, **kwargs) + # Native dispatch awaits the coroutine on the executor's loop; the + # blocking path has no loop of its own, so run_maybe_async drives it. + ret = await called if is_async else run_maybe_async(called) + # Per-item batch result: when the task is batched with + # per_item_results=True, enforce the return type contract and surface + # partial failures via BatchPartialFailureError so the existing + # retry/DLQ machinery applies. + batch_cfg = queue_ref._task_batch_configs.get(task_name) + if batch_cfg is not None and getattr(batch_cfg, "per_item_results", False): + if not is_batch_item_result_list(ret): + raise BatchResultTypeError( + f"task {task_name!r} declares per_item_results=True but " + f"returned {type(ret).__name__} instead of " + f"list[BatchItemResult]" + ) + failed = [item for item in ret if item.status == "failure"] + if failed: + # Pass the result through so the worker can still store it + # (per-item outcomes are stored verbatim); the existing + # on_failure path emits JOB_FAILED. + result = ret + raise BatchPartialFailureError(failed_items=failed) + result = ret + except Exception as exc: + error = exc + elapsed = time.perf_counter() - started_at + if isinstance(exc, TaskCancelledError): + logger.info( + "Task %s[%s] cancelled in %.3fs: %s", + task_name, + job_id, + elapsed, + exc, + ) + else: + logger.error( + "Task %s[%s] raised %s in %.3fs: %r", + task_name, + job_id, + type(exc).__name__, + elapsed, + exc, + ) + for hook in hooks["on_failure"]: + hook(task_name, args, kwargs, exc) + raise + else: + elapsed = time.perf_counter() - started_at + logger.info( + "Task %s[%s] succeeded in %.3fs: %s", + task_name, + job_id, + elapsed, + _safe_result_repr(result), + ) + for hook in hooks["on_success"]: + hook(task_name, args, kwargs, result) + return result + finally: + # Pop the saga compensation context (no-op outside saga flow). + if comp_ctx_token is not None: + _reset_compensation_context(comp_ctx_token) + # Release task/request-scoped resources + for release_fn in release_callbacks: + try: + release_fn() + except Exception: + logger.exception("resource release error") + # Clean up reconstructed proxies (LIFO order) + cleanup_proxies(proxy_cleanup, metrics=queue_ref._proxy_metrics) + for hook in hooks["after_task"]: + hook(task_name, args, kwargs, result, error) + # Run middleware after hooks (only those whose before() succeeded) + for mw in completed_mw: + try: + mw.after(current_job, result, error) + except Exception: + logger.exception("middleware after() error") + # Emit job lifecycle events. The Rust outcome loop skips Success on the + # grounds that this body emits it, so it must stay here for both paths. + event_payload: dict[str, Any] = { + "task_name": task_name, + "job_id": current_job.id, + "queue": current_job.queue_name, + } + if error is not None: + event_payload["error"] = str(error) + queue_ref._emit_event(EventType.JOB_FAILED, event_payload) + else: + queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) diff --git a/sdks/python/tests/observability/test_task_lifecycle_logs.py b/sdks/python/tests/observability/test_task_lifecycle_logs.py index 95b2aa23..04e52aea 100644 --- a/sdks/python/tests/observability/test_task_lifecycle_logs.py +++ b/sdks/python/tests/observability/test_task_lifecycle_logs.py @@ -9,7 +9,7 @@ from taskito import Queue from taskito.exceptions import TaskCancelledError -from taskito.mixins.decorators import _MAX_RESULT_REPR, _safe_result_repr +from taskito.task_lifecycle import _MAX_RESULT_REPR, _safe_result_repr _RECEIVED = re.compile(r"^Task (\S+)\[(\S+)\] received$") _SUCCEEDED = re.compile(r"^Task (\S+)\[(\S+)\] succeeded in \d+\.\d{3}s: (.+)$") From 1d5e45bd2044b61c6a1097b923455d4d4ad04f4e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:37:30 +0530 Subject: [PATCH 02/13] test(python): cover lifecycle parity across sync and async Drives each lifecycle guarantee through a real worker, parametrised over def and async def. The async cases route sync while native dispatch is dormant; they are the net for activating it. --- .../tests/worker/test_lifecycle_parity.py | 350 ++++++++++++++++++ 1 file changed, 350 insertions(+) create mode 100644 sdks/python/tests/worker/test_lifecycle_parity.py diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py new file mode 100644 index 00000000..b6b1ab71 --- /dev/null +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -0,0 +1,350 @@ +"""The lifecycle must behave identically whether a task is sync or async. + +Every test here runs a real ``@queue.task()`` through a real worker, and is +parametrised over ``def`` vs ``async def``. That combination is the point: the +existing async tests drive ``AsyncTaskExecutor`` directly with a mocked queue, so +they cannot see whether a hook fired or a compensation context was set, and the +existing hook/saga/batch tests use a real queue but only sync tasks. Neither +shape catches a gap between the two dispatch paths. + +Each test asserts on behaviour a caller was promised, not on which path ran, so +the same assertions hold before and after native async dispatch is activated. +That is also why none of them skip when the build lacks ``native-async``: the +``is_async=True`` cases then route through the blocking path, which is worth +asserting in its own right. Proving *which* path is live is a separate job, and +belongs to the canary that asserts coroutines run on the executor thread — +these tests deliberately cannot tell, and must not. + +Tasks are declared before the worker starts on purpose: the scheduler registers +task policy once at worker boot, so a task declared afterwards silently falls +back to ``RetryPolicy::default()`` and ``max_retries=0`` is ignored. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Generator +from pathlib import Path +from typing import Any + +import pytest + +from taskito import BatchItemResult, MaxRetriesExceededError, Queue +from taskito.context import current_job +from taskito.exceptions import SoftTimeoutError +from taskito.middleware import TaskMiddleware +from taskito.workflows import Workflow, WorkflowState +from taskito.workflows.saga import current_compensation_context + +# Parametrise the *task*, not the assertion: `async`ness is a property of the +# function object, so each case defines its own body under an `if is_async`. +BOTH_PATHS = pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) + +_TIMEOUT = 15.0 + + +@pytest.fixture +def noretry_queue(tmp_path: Path) -> Generator[Queue]: + """Queue whose jobs DLQ on first failure, so failures stay observable.""" + q = Queue(db_path=str(tmp_path / "parity.db"), workers=2, default_retry=0) + try: + yield q + finally: + q.close() + + +def _start_worker(queue: Queue) -> threading.Thread: + """Start a worker. Call only after every task in the test is declared.""" + thread = threading.Thread(target=queue.run_worker, daemon=True) + thread.start() + return thread + + +def _stop_worker(queue: Queue, thread: threading.Thread) -> None: + queue._inner.request_shutdown() + thread.join(timeout=10) + + +def _wait_for(predicate: Any, message: str) -> None: + """Poll until ``predicate()`` is truthy, or fail with ``message``.""" + deadline = time.monotonic() + _TIMEOUT + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.02) + raise AssertionError(message) + + +@BOTH_PATHS +def test_queue_hooks_fire(queue: Queue, is_async: bool) -> None: + """before_task/on_success/after_task fire for either kind of task.""" + events: list[tuple[Any, ...]] = [] + + @queue.before_task + def on_before(task_name: str, args: tuple, kwargs: dict) -> None: + events.append(("before", task_name)) + + @queue.on_success + def on_ok(task_name: str, args: tuple, kwargs: dict, result: Any) -> None: + events.append(("success", task_name, result)) + + @queue.after_task + def on_after(task_name: str, args: tuple, kwargs: dict, result: Any, error: Any) -> None: + events.append(("after", task_name, result, error)) + + if is_async: + + @queue.task(name="hooked") + async def hooked(a: int, b: int) -> int: + return a + b + + else: + + @queue.task(name="hooked") + def hooked(a: int, b: int) -> int: + return a + b + + job = hooked.delay(1, 2) + thread = _start_worker(queue) + try: + assert job.result(timeout=_TIMEOUT) == 3 + finally: + _stop_worker(queue, thread) + + assert ("before", "hooked") in events + assert ("success", "hooked", 3) in events + assert ("after", "hooked", 3, None) in events + + +@BOTH_PATHS +def test_on_failure_hook_sees_the_exception(noretry_queue: Queue, is_async: bool) -> None: + """on_failure receives the live exception for either kind of task.""" + failures: list[BaseException] = [] + + @noretry_queue.on_failure + def on_fail(task_name: str, args: tuple, kwargs: dict, error: BaseException) -> None: + failures.append(error) + + if is_async: + + @noretry_queue.task(name="boom", max_retries=0) + async def boom() -> None: + raise ValueError("expected failure") + + else: + + @noretry_queue.task(name="boom", max_retries=0) + def boom() -> None: + raise ValueError("expected failure") + + boom.delay() + thread = _start_worker(noretry_queue) + try: + _wait_for(lambda: failures, "on_failure hook never fired") + finally: + _stop_worker(noretry_queue, thread) + + assert isinstance(failures[0], ValueError) + assert "expected failure" in str(failures[0]) + + +@BOTH_PATHS +def test_soft_timeout_makes_check_timeout_raise(noretry_queue: Queue, is_async: bool) -> None: + """check_timeout() raises once soft_timeout has elapsed, for either kind. + + soft_timeout reaches the task body through per-task config the lifecycle + reads, so a body that never sees it would return "no timeout" instead. + """ + outcomes: list[str] = [] + + if is_async: + + @noretry_queue.task(name="slow", soft_timeout=0.05, max_retries=0) + async def slow() -> str: + time.sleep(0.2) + try: + current_job.check_timeout() + except SoftTimeoutError: + outcomes.append("raised") + raise + outcomes.append("no timeout") + return "no timeout" + + else: + + @noretry_queue.task(name="slow", soft_timeout=0.05, max_retries=0) + def slow() -> str: + time.sleep(0.2) + try: + current_job.check_timeout() + except SoftTimeoutError: + outcomes.append("raised") + raise + outcomes.append("no timeout") + return "no timeout" + + slow.delay() + thread = _start_worker(noretry_queue) + try: + _wait_for(lambda: outcomes, "task body never ran") + finally: + _stop_worker(noretry_queue, thread) + + assert outcomes == ["raised"], "check_timeout() should raise once soft_timeout elapsed" + + +@BOTH_PATHS +def test_batch_per_item_failure_is_not_a_success(noretry_queue: Queue, is_async: bool) -> None: + """A per-item failure fails the batch for either kind of task. + + Without per-item handling the task returns a list and looks like a plain + success, so the failed item's caller would receive that list as its result + instead of an error. + """ + batch = {"max_size": 2, "max_wait_ms": 60_000, "per_item_results": True} + + if is_async: + + @noretry_queue.task(name="items", batch=batch, max_retries=0) + async def items(values: list[int]) -> list[BatchItemResult]: + return [ + BatchItemResult.failure(item_index=0, error="permanent"), + BatchItemResult.success(item_index=1, result=values[1]), + ] + + else: + + @noretry_queue.task(name="items", batch=batch, max_retries=0) + def items(values: list[int]) -> list[BatchItemResult]: + return [ + BatchItemResult.failure(item_index=0, error="permanent"), + BatchItemResult.success(item_index=1, result=values[1]), + ] + + thread = _start_worker(noretry_queue) + try: + first = items.delay(7) + items.delay(8) + with pytest.raises(MaxRetriesExceededError, match="permanent"): + first.result(timeout=_TIMEOUT) + finally: + _stop_worker(noretry_queue, thread) + + +@BOTH_PATHS +def test_saga_compensation_context_is_visible( + queue: Queue, workflow_worker: Any, is_async: bool +) -> None: + """A compensator body can read its context for either kind of task. + + Parametrised on the *compensator*, since that is the body that reads the + context. Its forward step is followed by a failing step to trigger it. + """ + captured: dict[str, Any] = {} + + if is_async: + + @queue.task(name="refund", max_retries=0) + async def refund(forward_args: tuple, forward_kwargs: dict, forward_result: Any) -> None: + ctx = current_compensation_context() + captured["ctx"] = ctx + if ctx is not None: + captured["node_name"] = ctx.workflow_node_name + captured["forward_result"] = ctx.forward_result + + else: + + @queue.task(name="refund", max_retries=0) + def refund(forward_args: tuple, forward_kwargs: dict, forward_result: Any) -> None: + ctx = current_compensation_context() + captured["ctx"] = ctx + if ctx is not None: + captured["node_name"] = ctx.workflow_node_name + captured["forward_result"] = ctx.forward_result + + @queue.task(name="charge", max_retries=0, compensates=refund) + def charge(amount: int) -> str: + return f"charge-{amount}" + + @queue.task(name="fail_later", max_retries=0) + def fail_later() -> None: + raise RuntimeError("trigger compensation") + + wf = Workflow(name="parity_saga") + wf.step("charge", charge, args=(250,)) + wf.step("fail", fail_later, after="charge") + + with workflow_worker(): + run = queue.submit_workflow(wf) + final = run.wait(timeout=30) + assert final.state in {WorkflowState.COMPENSATED, WorkflowState.COMPENSATION_FAILED} + + assert captured.get("ctx") is not None, "compensator saw no compensation context" + assert captured["node_name"] == "charge" + assert captured["forward_result"] == "charge-250" + + +@BOTH_PATHS +def test_lifecycle_is_logged(queue: Queue, is_async: bool, caplog: Any) -> None: + """The received and succeeded lines are emitted for either kind of task.""" + if is_async: + + @queue.task(name="logged") + async def logged() -> str: + return "ok" + + else: + + @queue.task(name="logged") + def logged() -> str: + return "ok" + + with caplog.at_level("INFO", logger="taskito"): + job = logged.delay() + thread = _start_worker(queue) + try: + assert job.result(timeout=_TIMEOUT) == "ok" + finally: + _stop_worker(queue, thread) + + assert "Task logged" in caplog.text + assert "received" in caplog.text + assert "succeeded" in caplog.text + + +@BOTH_PATHS +def test_middleware_after_receives_the_error(tmp_path: Path, is_async: bool) -> None: + """Middleware .after() receives the exception, not None, for either kind.""" + seen: list[Any] = [] + ran = threading.Event() + + class Recorder(TaskMiddleware): + def after(self, job: Any, result: Any, error: Any) -> None: + seen.append(error) + ran.set() + + q = Queue(db_path=str(tmp_path / "mw.db"), workers=2, default_retry=0, middleware=[Recorder()]) + try: + if is_async: + + @q.task(name="failing", max_retries=0) + async def failing() -> None: + raise RuntimeError("expected failure") + + else: + + @q.task(name="failing", max_retries=0) + def failing() -> None: + raise RuntimeError("expected failure") + + failing.delay() + thread = _start_worker(q) + try: + assert ran.wait(timeout=_TIMEOUT), "middleware after() never fired" + finally: + _stop_worker(q, thread) + + assert isinstance(seen[0], RuntimeError), f"after() should see the error, got {seen[0]!r}" + finally: + q.close() From c2b44e2c68cec780b0eb98c388731bfb439fa63e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:38:17 +0530 Subject: [PATCH 03/13] fix(python): don't report a torn-down task as completed `except Exception` misses a BaseException, so the cleanup saw no error and emitted JOB_COMPLETED for a task whose body never returned. Latent today; asyncio.CancelledError is a BaseException and is routine once async tasks dispatch natively. --- sdks/python/taskito/task_lifecycle.py | 46 +++++++++++---- .../tests/worker/test_lifecycle_parity.py | 56 +++++++++++++++++++ 2 files changed, 91 insertions(+), 11 deletions(-) diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py index 7ecbe3af..af2a9be5 100644 --- a/sdks/python/taskito/task_lifecycle.py +++ b/sdks/python/taskito/task_lifecycle.py @@ -179,7 +179,10 @@ async def run_lifecycle( if comp_ctx is not None: comp_ctx_token = _set_compensation_context(comp_ctx) - error = None + # `torn_down` tracks a BaseException separately from a task failure: both must + # keep the cleanup below from reporting success, but only a failure is one. + error: BaseException | None = None + torn_down = False result = None try: called = fn(*args, **kwargs) @@ -229,6 +232,23 @@ async def run_lifecycle( for hook in hooks["on_failure"]: hook(task_name, args, kwargs, exc) raise + except BaseException as exc: + # asyncio.CancelledError is a BaseException, so the clause above cannot see + # it — and native dispatch raises one on every cancellation and loop + # shutdown. Record it, or the cleanup below reads `error is None` and calls + # a task that never returned a success. The failure hooks stay out of it: + # being torn down is not the task failing, and the disposition event comes + # from the outcome loop (JOB_CANCELLED), not from here. + error = exc + torn_down = True + logger.info( + "Task %s[%s] torn down after %.3fs by %s", + task_name, + job_id, + time.perf_counter() - started_at, + type(exc).__name__, + ) + raise else: elapsed = time.perf_counter() - started_at logger.info( @@ -263,13 +283,17 @@ async def run_lifecycle( logger.exception("middleware after() error") # Emit job lifecycle events. The Rust outcome loop skips Success on the # grounds that this body emits it, so it must stay here for both paths. - event_payload: dict[str, Any] = { - "task_name": task_name, - "job_id": current_job.id, - "queue": current_job.queue_name, - } - if error is not None: - event_payload["error"] = str(error) - queue_ref._emit_event(EventType.JOB_FAILED, event_payload) - else: - queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) + # A torn-down job is the exception: it has no outcome to report yet, and + # the outcome loop emits its own (JOB_CANCELLED, or JOB_RETRYING/JOB_DEAD + # once the failure is classified). + if not torn_down: + event_payload: dict[str, Any] = { + "task_name": task_name, + "job_id": current_job.id, + "queue": current_job.queue_name, + } + if error is not None: + event_payload["error"] = str(error) + queue_ref._emit_event(EventType.JOB_FAILED, event_payload) + else: + queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index b6b1ab71..66af3e52 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -32,6 +32,7 @@ from taskito import BatchItemResult, MaxRetriesExceededError, Queue from taskito.context import current_job +from taskito.events import EventType from taskito.exceptions import SoftTimeoutError from taskito.middleware import TaskMiddleware from taskito.workflows import Workflow, WorkflowState @@ -348,3 +349,58 @@ def failing() -> None: assert isinstance(seen[0], RuntimeError), f"after() should see the error, got {seen[0]!r}" finally: q.close() + + +@BOTH_PATHS +def test_base_exception_is_not_reported_as_completed(tmp_path: Path, is_async: bool) -> None: + """A task killed by a BaseException must not emit JOB_COMPLETED. + + ``except Exception`` does not catch a BaseException, so without an explicit + clause the lifecycle's cleanup sees no error and calls the job a success — + it emits JOB_COMPLETED and hands ``error=None`` to after_task/middleware for + a task whose body never returned. Latent on the blocking path, where only an + exotic BaseException reaches it, but ``asyncio.CancelledError`` is a + BaseException and is routine on the native path: every cancelled coroutine + and every loop shutdown raises one. + """ + completed: list[dict] = [] + after_task_errors: list[Any] = [] + ran = threading.Event() + + class Killed(BaseException): + """Not an Exception, so `except Exception` cannot see it.""" + + q = Queue(db_path=str(tmp_path / "base_exc.db"), workers=2, default_retry=0) + try: + q.on_event(EventType.JOB_COMPLETED, lambda _t, payload: completed.append(payload)) + + @q.after_task + def on_after(task_name: str, args: tuple, kwargs: dict, result: Any, error: Any) -> None: + after_task_errors.append(error) + ran.set() + + if is_async: + + @q.task(name="killed", max_retries=0) + async def killed() -> None: + raise Killed("torn down") + + else: + + @q.task(name="killed", max_retries=0) + def killed() -> None: + raise Killed("torn down") + + killed.delay() + thread = _start_worker(q) + try: + assert ran.wait(timeout=_TIMEOUT), "after_task never fired" + finally: + _stop_worker(q, thread) + + assert completed == [], f"a torn-down task must not emit JOB_COMPLETED, got {completed}" + assert isinstance(after_task_errors[0], Killed), ( + f"after_task should see the BaseException, got {after_task_errors[0]!r}" + ) + finally: + q.close() From f063ff413d57f4d0cd46ea06fb276bfe731627d3 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 02:08:21 +0530 Subject: [PATCH 04/13] refactor(python): run native dispatch through the shared lifecycle The executor reimplemented the lifecycle and had drifted, silently missing the queue hooks, saga context, soft_timeout and per-item batch handling. It now awaits run_lifecycle and keeps only what differs: decoding the payload and answering the result sender. --- sdks/python/taskito/async_support/executor.py | 143 ++----------- sdks/python/tests/worker/test_native_async.py | 189 ++++++------------ 2 files changed, 82 insertions(+), 250 deletions(-) diff --git a/sdks/python/taskito/async_support/executor.py b/sdks/python/taskito/async_support/executor.py index 45b51ce9..f2410e17 100644 --- a/sdks/python/taskito/async_support/executor.py +++ b/sdks/python/taskito/async_support/executor.py @@ -10,12 +10,9 @@ 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 from taskito.task_errors import encode_task_error +from taskito.task_lifecycle import run_lifecycle if TYPE_CHECKING: from taskito.app import Queue @@ -30,9 +27,10 @@ class AsyncTaskExecutor: """Runs async tasks natively on a dedicated event loop. - Receives jobs from Rust via :meth:`submit_job`, executes the async coroutine - with full lifecycle support (interception, proxies, resource injection, - middleware), and reports results back via a ``PyResultSender``. + Receives jobs from Rust via :meth:`submit_job`, awaits the coroutine through + :func:`~taskito.task_lifecycle.run_lifecycle` — the same body the blocking + path runs, so both dispatch paths honour one lifecycle — and reports the + outcome back via a ``PyResultSender``. """ def __init__( @@ -143,21 +141,19 @@ async def _run_job( max_retries: int, queue_name: str, ) -> None: - """Run one job: gate on the concurrency backstop, execute, report.""" + """Run one job: gate on the concurrency backstop, execute, report. + + Everything between the arguments and the result belongs to + :func:`run_lifecycle`; what stays here is what native dispatch does + differently — decoding the payload, and answering the result sender. + """ 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) - release_callbacks: list[Any] = [] - proxy_cleanup: list[Any] = [] - 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 @@ -165,65 +161,17 @@ async def _run_job( # a hardcoded cloudpickle.loads would ignore @task(serializer=...). args, kwargs = queue._deserialize_payload(task_name, payload_bytes) - # Worker-dispatch predicate gate (raw args, pre-reconstruction). - if task_name in queue._task_predicates: - queue._apply_dispatch_predicate( - task_name=task_name, - args=args, - kwargs=kwargs, - job_id=job_id, - queue_name=queue_name, - retry_count=retry_count, - ) - - # Reconstruct intercepted arguments - redirects: dict[str, str] = {} - if queue._interceptor is not None: - args, kwargs, redirects = reconstruct_args(args, kwargs) - - # Reconstruct proxy markers - if queue._proxy_registry is not None and not queue._test_mode_active: - args, kwargs, proxy_cleanup = reconstruct_proxies( - args, - kwargs, - queue._proxy_registry, - signing_secret=queue._recipe_signing_key, - max_timeout=queue._max_reconstruction_timeout, - metrics=queue._proxy_metrics, - ) + # Awaited, not run_maybe_async'd: this is the executor's own loop, + # and the coroutine belongs on it. + result = await run_lifecycle( + queue, + task_name, + self._registry[task_name]._taskito_async_fn, + args, + kwargs, + is_async=True, + ) - # Inject resources - runtime = queue._resource_runtime - if runtime is not None: - for res_name in queue._task_inject_map.get(task_name, []): - if res_name not in kwargs: - instance, release = runtime.acquire_for_task(res_name) - kwargs[res_name] = instance - if release is not None: - release_callbacks.append(release) - for kwarg_name, resource_name in redirects.items(): - instance, release = runtime.acquire_for_task(resource_name) - kwargs[kwarg_name] = instance - if release is not None: - release_callbacks.append(release) - - # Middleware before hooks (skipping filtered middlewares) - middleware_chain = queue._get_middleware_chain(task_name) - for mw in middleware_chain: - if not mw._should_apply(current_job): - continue - try: - mw.before(current_job) - completed_mw.append(mw) - except Exception: - logger.exception("middleware before() error") - - # Execute the async function - wrapper = self._registry[task_name] - fn = wrapper._taskito_async_fn - result = await fn(*args, **kwargs) - - # Serialize and report success result_bytes = ( queue._serialize_result(task_name, result) if result is not None else None ) @@ -236,8 +184,6 @@ async def _run_job( ) except TaskCancelledError: - error = None # Don't treat cancellation as an error for middleware - cancelled = True if not reported: wall_ns = time.monotonic_ns() - start_ns reported = True @@ -249,7 +195,6 @@ async def _run_job( # 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 @@ -260,7 +205,6 @@ async def _run_job( raise except Exception as exc: - error = exc if not reported: wall_ns = time.monotonic_ns() - start_ns reported = True @@ -279,53 +223,8 @@ async def _run_job( ) finally: - # Release task/request-scoped resources - for release_fn in release_callbacks: - try: - release_fn() - except Exception: - logger.exception("resource release error") - - # Clean up reconstructed proxies - cleanup_proxies(proxy_cleanup, metrics=self._queue_ref._proxy_metrics) - - # Middleware after hooks (only those whose before() succeeded) - for mw in completed_mw: - try: - mw.after(current_job, result, error) - except Exception: - logger.exception("middleware after() error") - - # Emit job lifecycle events. The Rust outcome loop deliberately - # skips Success — the blocking wrapper emits it — and native - # dispatch bypasses that wrapper, so it has to emit here too or - # nothing downstream (workflow tracker, subscribers) ever learns - # the job finished. Cancellation is the exception — that outcome - # does have a Rust-side event. - if not cancelled: - self._emit_lifecycle_event(job_id, task_name, queue_name, error) - clear_async_context(token) - def _emit_lifecycle_event( - self, job_id: str, task_name: str, queue_name: str, error: Exception | None - ) -> None: - """Emit JOB_COMPLETED/JOB_FAILED, mirroring the blocking task wrapper.""" - payload: dict[str, Any] = { - "task_name": task_name, - "job_id": job_id, - "queue": queue_name, - } - if error is not None: - payload["error"] = str(error) - try: - self._queue_ref._emit_event( - EventType.JOB_FAILED if error is not None else EventType.JOB_COMPLETED, - payload, - ) - except Exception: - logger.exception("job lifecycle event error") - def _check_retry(self, task_name: str, exc: Exception) -> bool: """Check retry filters to decide if this exception should be retried.""" filters = self._queue_ref._task_retry_filters.get(task_name) diff --git a/sdks/python/tests/worker/test_native_async.py b/sdks/python/tests/worker/test_native_async.py index 299ce572..6571ebb8 100644 --- a/sdks/python/tests/worker/test_native_async.py +++ b/sdks/python/tests/worker/test_native_async.py @@ -8,6 +8,8 @@ from typing import Any from unittest.mock import MagicMock +import cloudpickle + from taskito import Queue, TaskCancelledError, current_job from taskito.async_support.context import ( clear_async_context, @@ -18,6 +20,38 @@ PollUntil = Any # the conftest fixture's runtime type + +def _fake_queue(**overrides: Any) -> MagicMock: + """A queue stub covering every attribute the shared lifecycle reads. + + MagicMock invents any attribute asked of it, and the lifecycle branches on + ``is not None`` and on truthiness — so an unset attribute does not raise, it + silently takes the wrong branch. An absent ``_task_batch_configs`` is the + sharp one: the auto-made mock reads as a per-item batch config, and every + task then fails the return-type contract. Listing them here once means a new + lifecycle read is one edit rather than one per test. + """ + queue = MagicMock() + queue._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) + queue._serialize_result.side_effect = lambda _name, result: cloudpickle.dumps(result) + queue._interceptor = None + queue._proxy_registry = None + queue._proxy_metrics = None + queue._test_mode_active = False + queue._resource_runtime = None + queue._task_inject_map = {} + queue._task_retry_filters = {} + queue._task_predicates = {} + queue._task_batch_configs = {} + queue._task_soft_timeouts = {} + queue._workflow_tracker = None + queue._hooks = {"before_task": [], "after_task": [], "on_success": [], "on_failure": []} + queue._get_middleware_chain.return_value = [] + for name, value in overrides.items(): + setattr(queue, name, value) + return queue + + # ── Async detection ────────────────────────────────────────────── @@ -135,8 +169,6 @@ def test_async_executor_lifecycle() -> None: def test_async_executor_submit_and_execute(poll_until: PollUntil) -> None: """Basic async task produces correct result via executor.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -150,17 +182,7 @@ class FakeWrapper: registry: dict[str, Any] = {"test_mod.my_task": FakeWrapper()} - 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 + queue_ref = _fake_queue() executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -180,8 +202,6 @@ class FakeWrapper: def test_async_exception_reported(poll_until: PollUntil) -> None: """Exception in async task → failure result with traceback.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -194,17 +214,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.failing_task": FakeWrapper()} - 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 + queue_ref = _fake_queue() executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -223,8 +233,6 @@ class FakeWrapper: def test_async_cancellation(poll_until: PollUntil) -> None: """TaskCancelledError → cancelled result.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -237,17 +245,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.cancelling_task": FakeWrapper()} - 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 + queue_ref = _fake_queue() executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -265,8 +263,6 @@ class FakeWrapper: def test_async_retry_filter(poll_until: PollUntil) -> None: """Failed async task respects retry_on filter.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -279,14 +275,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.flaky_task": FakeWrapper()} - 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 = _fake_queue() # Only retry on ValueError, not TypeError queue_ref._task_retry_filters = { "mod.flaky_task": {"retry_on": [ValueError], "dont_retry_on": []}, @@ -308,8 +297,6 @@ class FakeWrapper: def test_async_concurrency_limit(poll_until: PollUntil) -> None: """Semaphore bounds concurrent async tasks.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -331,17 +318,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.slow_task": FakeWrapper()} - 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 + queue_ref = _fake_queue() # Set concurrency to 2 executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=2) @@ -364,8 +341,6 @@ class FakeWrapper: def test_async_middleware_hooks(poll_until: PollUntil) -> None: """Middleware before/after called for async tasks.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor before_called: list[str] = [] @@ -388,17 +363,8 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.simple_task": FakeWrapper()} - 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 = _fake_queue() queue_ref._get_middleware_chain.return_value = [TestMiddleware()] - queue_ref._proxy_metrics = None executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -414,8 +380,6 @@ class FakeWrapper: def test_async_task_with_injection(poll_until: PollUntil) -> None: """inject=["db"] works for async tasks via executor.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -430,21 +394,14 @@ class FakeWrapper: fake_db = "fake-conn" - 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._task_inject_map = {"mod.db_task": ["db"]} - queue_ref._task_retry_filters = {} - queue_ref._get_middleware_chain.return_value = [] - queue_ref._proxy_metrics = None - # Mock resource runtime runtime = MagicMock() runtime.acquire_for_task.return_value = (fake_db, None) - queue_ref._resource_runtime = runtime + + queue_ref = _fake_queue( + _task_inject_map={"mod.db_task": ["db"]}, + _resource_runtime=runtime, + ) executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -461,8 +418,6 @@ class FakeWrapper: def test_async_context_available_inside_task(poll_until: PollUntil) -> None: """current_job.id works inside an async task via contextvars.""" - import cloudpickle - from taskito.async_support.executor import AsyncTaskExecutor sender = MagicMock() @@ -477,17 +432,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.ctx_task": FakeWrapper()} - 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 + queue_ref = _fake_queue() executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) executor.start() @@ -531,19 +476,11 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.add": FakeWrapper()} - queue_ref = MagicMock() - # Route deserialization through a NON-cloudpickle serializer; a hardcoded + queue_ref = _fake_queue() + # Route serialization through a NON-cloudpickle serializer; a hardcoded # cloudpickle.loads would raise on this JSON payload. queue_ref._deserialize_payload.side_effect = lambda _name, data: serializer.loads(data) queue_ref._serialize_result.side_effect = lambda _name, result: serializer.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, registry, queue_ref, max_concurrency=10) executor.start() @@ -574,24 +511,12 @@ def release(self) -> None: 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 + queue_ref = _fake_queue() executor = AsyncTaskExecutor(sender, {task_name: FakeWrapper()}, queue_ref, max_concurrency=10) executor.start() @@ -599,8 +524,6 @@ class FakeWrapper: def _payload() -> bytes: - import cloudpickle - payload: bytes = cloudpickle.dumps(((), {})) return payload @@ -778,7 +701,15 @@ async def boom_task() -> None: 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.""" + """Cancellation has its own Rust-side event; emitting COMPLETED would be a lie. + + Both paths run one lifecycle, so both report a cancel the same way: + TaskCancelledError is an ordinary Exception, so it emits JOB_FAILED here and + the outcome loop adds JOB_CANCELLED. That pairing is the blocking path's + long-standing behaviour — what must never appear is JOB_COMPLETED. + """ + from taskito.events import EventType + sender = MagicMock() sender.try_report_cancelled.return_value = True @@ -791,4 +722,6 @@ async def cancelling_task() -> None: 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" + emitted = [call.args[0] for call in queue_ref._emit_event.call_args_list] + assert EventType.JOB_COMPLETED not in emitted, "a cancelled job did not complete" + assert EventType.JOB_FAILED in emitted, "cancel should report like the blocking path" From 6fb650d65da1512f349fc2aff977c44ace298d54 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 02:08:46 +0530 Subject: [PATCH 05/13] fix(python): don't retain a failed task's exception MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lifecycle held the exception in a plain local and passed it as a log argument, so both the frame and any retained LogRecord kept it — and through its traceback, every frame that raised. Latent on the blocking path; on native dispatch those frames own the result sender the worker's drain loop waits on. --- sdks/python/taskito/task_lifecycle.py | 21 ++++++-- .../tests/worker/test_lifecycle_parity.py | 50 +++++++++++++++++++ 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py index af2a9be5..fe7230e3 100644 --- a/sdks/python/taskito/task_lifecycle.py +++ b/sdks/python/taskito/task_lifecycle.py @@ -212,22 +212,27 @@ async def run_lifecycle( except Exception as exc: error = exc elapsed = time.perf_counter() - started_at + # Format the exception into the message rather than passing it as an + # argument: a LogRecord keeps its args, and any handler that retains + # records (a buffering handler, an error reporter, a test's log capture) + # would then hold the exception — and through its traceback, every frame + # that raised, for as long as the record lives. if isinstance(exc, TaskCancelledError): logger.info( "Task %s[%s] cancelled in %.3fs: %s", task_name, job_id, elapsed, - exc, + str(exc), ) else: logger.error( - "Task %s[%s] raised %s in %.3fs: %r", + "Task %s[%s] raised %s in %.3fs: %s", task_name, job_id, type(exc).__name__, elapsed, - exc, + repr(exc), ) for hook in hooks["on_failure"]: hook(task_name, args, kwargs, exc) @@ -297,3 +302,13 @@ async def run_lifecycle( queue_ref._emit_event(EventType.JOB_FAILED, event_payload) else: queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) + # Drop the exception before this frame outlives it. Its traceback points + # back at this very frame, so holding it in an ordinary local is a + # reference cycle — Python auto-clears an `except ... as` name for exactly + # this reason, but `error` is a plain local and keeps it. Until the cyclic + # collector runs, everything reachable from this frame stays alive, which + # on native dispatch includes the async result sender: the worker's drain + # loop ends when that sender drops, so a leaked one stalls every shutdown + # after a failed task until the drain times out. Must stay last — it is + # only safe once nothing above still reads `error`. + error = None diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 66af3e52..6b66ef4e 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -404,3 +404,53 @@ def killed() -> None: ) finally: q.close() + + +@BOTH_PATHS +def test_failed_task_does_not_stall_shutdown(tmp_path: Path, is_async: bool) -> None: + """A failed task must not keep the worker from stopping promptly. + + The lifecycle holds the exception in a local while it runs the cleanup, and + the exception's traceback points back at that frame — a cycle, so the frame + and everything reachable from it survives until the cyclic collector runs. + On native dispatch the async result sender is reachable from that frame, and + the worker's drain loop only ends once the sender drops and the result + channel disconnects. Leak it and every shutdown after a failed task waits out + the full drain timeout instead of stopping. + """ + q = Queue(db_path=str(tmp_path / "stall.db"), workers=2, default_retry=0) + try: + if is_async: + + @q.task(name="failing", max_retries=0) + async def failing() -> None: + raise ValueError("expected failure") + + else: + + @q.task(name="failing", max_retries=0) + def failing() -> None: + raise ValueError("expected failure") + + job = failing.delay() + thread = _start_worker(q) + + # Wait for the job to reach a terminal state, so the failure is genuinely + # behind us and the only thing left to measure is the shutdown itself. + # `status` is the last-fetched value, so it needs a refresh to advance. + def is_terminal() -> bool: + job.refresh() + return job.status in ("failed", "dead") + + _wait_for(is_terminal, "job never reached a terminal state") + + started = time.monotonic() + q._inner.request_shutdown() + thread.join(timeout=30) + elapsed = time.monotonic() - started + + assert not thread.is_alive(), "worker did not stop after a failed task" + # Prompt is ~0.2s; the leak shows up as the multi-second drain timeout. + assert elapsed < 5, f"shutdown after a failed task took {elapsed:.1f}s — sender leaked?" + finally: + q.close() From 7dd99aca0ddc3af214843671221305bcf9cd7ffc Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 02:09:00 +0530 Subject: [PATCH 06/13] docs(python): record what still blocks native async dispatch The lifecycle gaps the old note cites are closed; shutdown is what is left. --- sdks/python/taskito/mixins/decorators.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 19b880cb..52995f92 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -420,13 +420,21 @@ def decorator(fn: Callable) -> TaskWrapper: self._task_soft_timeouts[task_name] = soft_timeout wrapped = self._wrap_task(fn, task_name) - # 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. + # 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. self._task_registry[task_name] = wrapped cb_threshold = None From 4e32084b8f8d57339365afd17152672899a381b2 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:09:19 +0530 Subject: [PATCH 07/13] fix(python): await what a task returns, not what it declares A plain `def` handing back a coroutine worked before the lifecycle became one: run_maybe_async now nests inside the lifecycle's own loop and raises. Deciding on the returned value covers both kinds and drops the is_async parameter. --- sdks/python/taskito/async_support/executor.py | 1 - sdks/python/taskito/mixins/decorators.py | 2 -- sdks/python/taskito/task_lifecycle.py | 21 ++++++++-------- .../tests/worker/test_lifecycle_parity.py | 24 +++++++++++++++++++ 4 files changed, 35 insertions(+), 13 deletions(-) diff --git a/sdks/python/taskito/async_support/executor.py b/sdks/python/taskito/async_support/executor.py index f2410e17..77287342 100644 --- a/sdks/python/taskito/async_support/executor.py +++ b/sdks/python/taskito/async_support/executor.py @@ -169,7 +169,6 @@ async def _run_job( self._registry[task_name]._taskito_async_fn, args, kwargs, - is_async=True, ) result_bytes = ( diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 52995f92..7d52ffd2 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -148,7 +148,6 @@ def _wrap_task(self, fn: Callable, task_name: str) -> Callable: makes that safe — `run_maybe_async` raises if one is already running. """ queue_ref = self - is_async = inspect.iscoroutinefunction(fn) @functools.wraps(fn) def wrapper(*args: Any, **kwargs: Any) -> Any: @@ -161,7 +160,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: fn, args, kwargs, - is_async=is_async, ) ) finally: diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py index fe7230e3..fc06bc77 100644 --- a/sdks/python/taskito/task_lifecycle.py +++ b/sdks/python/taskito/task_lifecycle.py @@ -29,12 +29,12 @@ from __future__ import annotations +import inspect import logging import time from collections.abc import Callable from typing import TYPE_CHECKING, Any -from taskito.async_support.helpers import run_maybe_async from taskito.batching.item_result import ( BatchPartialFailureError, BatchResultTypeError, @@ -77,8 +77,6 @@ async def run_lifecycle( fn: Callable, args: Any, kwargs: Any, - *, - is_async: bool, ) -> Any: """Run one task end to end: gates, reconstruction, injection, middleware, hooks. @@ -86,10 +84,9 @@ async def run_lifecycle( belongs to the caller, which knows whether it is answering Rust or a result sender. ``args``/``kwargs`` arrive already deserialized. - ``is_async`` selects how the body is invoked: awaited directly when the task - is a coroutine function (native dispatch), or driven to completion by - ``run_maybe_async`` when it is not. A blocking caller must therefore drive - the whole coroutine with ``run_maybe_async``; a native one awaits it. + This is a coroutine, so a blocking caller must drive it (``run_maybe_async``) + while a native one awaits it. Either way it ends up on an event loop, which + is why the task body is awaited here rather than driven. """ job_id = current_job.id logger.info("Task %s[%s] received", task_name, job_id) @@ -186,9 +183,13 @@ async def run_lifecycle( result = None try: called = fn(*args, **kwargs) - # Native dispatch awaits the coroutine on the executor's loop; the - # blocking path has no loop of its own, so run_maybe_async drives it. - ret = await called if is_async else run_maybe_async(called) + # Await whatever is awaitable, rather than trusting the task's declared + # kind: a plain `def` may hand back a coroutine, and that has always + # worked. Driving it with run_maybe_async instead would raise — this body + # already runs inside an event loop on both paths (``asyncio.run`` on the + # blocking one, the executor's own loop on the native one), and + # run_maybe_async refuses to nest. + ret = await called if inspect.isawaitable(called) else called # Per-item batch result: when the task is batched with # per_item_results=True, enforce the return type contract and surface # partial failures via BatchPartialFailureError so the existing diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 6b66ef4e..7e27d5f3 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -77,6 +77,30 @@ def _wait_for(predicate: Any, message: str) -> None: raise AssertionError(message) +def test_sync_task_returning_a_coroutine_is_awaited(queue: Queue) -> None: + """A plain ``def`` may hand back a coroutine, and its value is the result. + + Not parametrised: the whole point is a task the lifecycle would not classify + as async. The body must be awaited on what it returns, not on how it was + declared — driving it instead would raise, since the lifecycle already runs + on an event loop. + """ + + async def inner(x: int) -> int: + return x * 2 + + @queue.task(name="returns_coro") + def returns_coro(x: int) -> Any: + return inner(x) + + job = returns_coro.delay(21) + thread = _start_worker(queue) + try: + assert job.result(timeout=_TIMEOUT) == 42 + finally: + _stop_worker(queue, thread) + + @BOTH_PATHS def test_queue_hooks_fire(queue: Queue, is_async: bool) -> None: """before_task/on_success/after_task fire for either kind of task.""" From 08580e473e9ccc0828f334256e267d79848fe0fb Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:09:48 +0530 Subject: [PATCH 08/13] fix(python): make the lifecycle's cleanup unmissable Resources and proxies are acquired before the body is entered, so a raising before_task hook stranded them. The teardown could also skip clearing the exception if a hook raised, which was the point of clearing it. --- sdks/python/taskito/task_lifecycle.py | 211 ++++++++++-------- .../tests/worker/test_lifecycle_parity.py | 35 +++ 2 files changed, 156 insertions(+), 90 deletions(-) diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py index fc06bc77..fc1d13bb 100644 --- a/sdks/python/taskito/task_lifecycle.py +++ b/sdks/python/taskito/task_lifecycle.py @@ -58,6 +58,26 @@ _MAX_RESULT_REPR = 80 +def _release_acquired( + release_callbacks: list[Any], proxy_cleanup: list[Any], queue_ref: Queue +) -> None: + """Hand back everything acquired for one task. Never raises. + + Called from the task's teardown, and again if setup fails partway — by then + resources may already be held and proxies reconstructed, and nothing else + would give them back. + """ + for release_fn in release_callbacks: + try: + release_fn() + except Exception: + logger.exception("resource release error") + try: + cleanup_proxies(proxy_cleanup, metrics=queue_ref._proxy_metrics) + except Exception: + logger.exception("proxy cleanup error") + + def _safe_result_repr(value: Any) -> str: """Render a task return value for the success log, bounded and crash-proof.""" if value is None: @@ -122,59 +142,72 @@ async def run_lifecycle( metrics=queue_ref._proxy_metrics, ) - # Inject resources from runtime + # From here on the task holds things — resources, proxies — that only the + # teardown below gives back, and the teardown only runs once the body is + # entered. Anything that raises in between (acquiring a resource, a + # before_task hook) would strand them, so unwind explicitly on the way out. release_callbacks: list[Any] = [] - runtime = queue_ref._resource_runtime - if runtime is not None: - # From explicit inject=["db"] on task decorator - for res_name in queue_ref._task_inject_map.get(task_name, []): - if res_name not in kwargs: - instance, release = runtime.acquire_for_task(res_name) - kwargs[res_name] = instance + completed_mw: list[Any] = [] + comp_ctx_token = None + try: + # Inject resources from runtime + runtime = queue_ref._resource_runtime + if runtime is not None: + # From explicit inject=["db"] on task decorator + for res_name in queue_ref._task_inject_map.get(task_name, []): + if res_name not in kwargs: + instance, release = runtime.acquire_for_task(res_name) + kwargs[res_name] = instance + if release is not None: + release_callbacks.append(release) + # From interception REDIRECT markers + for kwarg_name, resource_name in redirects.items(): + instance, release = runtime.acquire_for_task(resource_name) + kwargs[kwarg_name] = instance if release is not None: release_callbacks.append(release) - # From interception REDIRECT markers - for kwarg_name, resource_name in redirects.items(): - instance, release = runtime.acquire_for_task(resource_name) - kwargs[kwarg_name] = instance - if release is not None: - release_callbacks.append(release) - middleware_chain = queue_ref._get_middleware_chain(task_name) - hooks = queue_ref._hooks + middleware_chain = queue_ref._get_middleware_chain(task_name) + hooks = queue_ref._hooks - # Set soft timeout on context if configured - soft_timeout = queue_ref._task_soft_timeouts.get(task_name) - if soft_timeout is not None: - current_job._set_soft_timeout(soft_timeout) + # Set soft timeout on context if configured + soft_timeout = queue_ref._task_soft_timeouts.get(task_name) + if soft_timeout is not None: + current_job._set_soft_timeout(soft_timeout) - # Run middleware before hooks (skipping middlewares whose predicate filter - # excludes this job) - completed_mw: list[Any] = [] - for mw in middleware_chain: - if not mw._should_apply(current_job): - continue - try: - mw.before(current_job) - completed_mw.append(mw) - except Exception: - logger.exception("middleware before() error") + # Run middleware before hooks (skipping middlewares whose predicate filter + # excludes this job) + for mw in middleware_chain: + if not mw._should_apply(current_job): + continue + try: + mw.before(current_job) + completed_mw.append(mw) + except Exception: + logger.exception("middleware before() error") - for hook in hooks["before_task"]: - hook(task_name, args, kwargs) + for hook in hooks["before_task"]: + hook(task_name, args, kwargs) - # Saga compensation context: if this job was dispatched by the saga - # orchestrator, look up the in-memory CompensationContext and push it onto a - # contextvar so the compensator body can call - # ``current_compensation_context()`` to introspect the forward execution. - # Pop in the ``finally`` below. - comp_ctx_token = None - tracker = getattr(queue_ref, "_workflow_tracker", None) - saga = getattr(tracker, "_saga", None) if tracker is not None else None - if saga is not None and saga.is_compensation_job(job_id): - comp_ctx = saga.take_compensation_context(job_id) - if comp_ctx is not None: - comp_ctx_token = _set_compensation_context(comp_ctx) + # Saga compensation context: if this job was dispatched by the saga + # orchestrator, look up the in-memory CompensationContext and push it onto a + # contextvar so the compensator body can call + # ``current_compensation_context()`` to introspect the forward execution. + # Pop in the ``finally`` below. + tracker = getattr(queue_ref, "_workflow_tracker", None) + saga = getattr(tracker, "_saga", None) if tracker is not None else None + if saga is not None and saga.is_compensation_job(job_id): + comp_ctx = saga.take_compensation_context(job_id) + if comp_ctx is not None: + comp_ctx_token = _set_compensation_context(comp_ctx) + except BaseException: + # Setup failed, so the body never runs and its teardown never will. Give + # back only what was acquired: after_task and the lifecycle events belong + # to a task that started, and this one did not. + if comp_ctx_token is not None: + _reset_compensation_context(comp_ctx_token) + _release_acquired(release_callbacks, proxy_cleanup, queue_ref) + raise # `torn_down` tracks a BaseException separately from a task failure: both must # keep the cleanup below from reporting success, but only a failure is one. @@ -268,48 +301,46 @@ async def run_lifecycle( hook(task_name, args, kwargs, result) return result finally: - # Pop the saga compensation context (no-op outside saga flow). - if comp_ctx_token is not None: - _reset_compensation_context(comp_ctx_token) - # Release task/request-scoped resources - for release_fn in release_callbacks: - try: - release_fn() - except Exception: - logger.exception("resource release error") - # Clean up reconstructed proxies (LIFO order) - cleanup_proxies(proxy_cleanup, metrics=queue_ref._proxy_metrics) - for hook in hooks["after_task"]: - hook(task_name, args, kwargs, result, error) - # Run middleware after hooks (only those whose before() succeeded) - for mw in completed_mw: - try: - mw.after(current_job, result, error) - except Exception: - logger.exception("middleware after() error") - # Emit job lifecycle events. The Rust outcome loop skips Success on the - # grounds that this body emits it, so it must stay here for both paths. - # A torn-down job is the exception: it has no outcome to report yet, and - # the outcome loop emits its own (JOB_CANCELLED, or JOB_RETRYING/JOB_DEAD - # once the failure is classified). - if not torn_down: - event_payload: dict[str, Any] = { - "task_name": task_name, - "job_id": current_job.id, - "queue": current_job.queue_name, - } - if error is not None: - event_payload["error"] = str(error) - queue_ref._emit_event(EventType.JOB_FAILED, event_payload) - else: - queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) - # Drop the exception before this frame outlives it. Its traceback points - # back at this very frame, so holding it in an ordinary local is a - # reference cycle — Python auto-clears an `except ... as` name for exactly - # this reason, but `error` is a plain local and keeps it. Until the cyclic - # collector runs, everything reachable from this frame stays alive, which - # on native dispatch includes the async result sender: the worker's drain - # loop ends when that sender drops, so a leaked one stalls every shutdown - # after a failed task until the drain times out. Must stay last — it is - # only safe once nothing above still reads `error`. - error = None + # The clearing of `error` below is not optional, so nothing here may skip + # it — a raising after_task hook or event subscriber must not be able to + # leave the exception behind. + try: + # Pop the saga compensation context (no-op outside saga flow). + if comp_ctx_token is not None: + _reset_compensation_context(comp_ctx_token) + _release_acquired(release_callbacks, proxy_cleanup, queue_ref) + for hook in hooks["after_task"]: + hook(task_name, args, kwargs, result, error) + # Run middleware after hooks (only those whose before() succeeded) + for mw in completed_mw: + try: + mw.after(current_job, result, error) + except Exception: + logger.exception("middleware after() error") + # Emit job lifecycle events. The Rust outcome loop skips Success on + # the grounds that this body emits it, so it must stay here for both + # paths. A torn-down job is the exception: it has no outcome to report + # yet, and the outcome loop emits its own (JOB_CANCELLED, or + # JOB_RETRYING/JOB_DEAD once the failure is classified). + if not torn_down: + event_payload: dict[str, Any] = { + "task_name": task_name, + "job_id": current_job.id, + "queue": current_job.queue_name, + } + if error is not None: + event_payload["error"] = str(error) + queue_ref._emit_event(EventType.JOB_FAILED, event_payload) + else: + queue_ref._emit_event(EventType.JOB_COMPLETED, event_payload) + finally: + # Drop the exception before this frame outlives it. Its traceback + # points back at this very frame, so holding it in an ordinary local + # is a reference cycle — Python auto-clears an `except ... as` name + # for exactly this reason, but `error` is a plain local and keeps it. + # Until the cyclic collector runs, everything reachable from this + # frame stays alive, which on native dispatch includes the async + # result sender: the worker's drain loop ends when that sender drops, + # so a leaked one stalls every shutdown after a failed task until the + # drain times out. + error = None diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 7e27d5f3..991870a3 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -77,6 +77,41 @@ def _wait_for(predicate: Any, message: str) -> None: raise AssertionError(message) +def test_setup_failure_releases_what_it_acquired(tmp_path: Path) -> None: + """A task that dies during setup must still hand back its resources. + + Resources are injected before the before_task hooks run, so a raising hook + leaves them held — and the teardown that would release them only runs once + the body is entered, which it never is. + """ + released: list[str] = [] + + q = Queue(db_path=str(tmp_path / "setup.db"), workers=2, default_retry=0) + try: + # Request scope, because its release calls `teardown` outright — a + # task-scoped resource goes back to a pool, where nothing observes it. + @q.worker_resource(name="db", scope="request", teardown=lambda conn: released.append(conn)) + def make_db() -> str: + return "conn" + + @q.before_task + def boom(task_name: str, args: tuple, kwargs: dict) -> None: + raise RuntimeError("hook blew up during setup") + + @q.task(name="needs_db", max_retries=0, inject=["db"]) + def needs_db(db: Any = None) -> str: + return f"got-{db}" + + needs_db.delay() + thread = _start_worker(q) + try: + _wait_for(lambda: released, "the request-scoped resource was never released") + finally: + _stop_worker(q, thread) + finally: + q.close() + + def test_sync_task_returning_a_coroutine_is_awaited(queue: Queue) -> None: """A plain ``def`` may hand back a coroutine, and its value is the result. From cbf7fdeb453f0528c97d9155ebe0c9928dc0a798 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:10:17 +0530 Subject: [PATCH 09/13] fix(python): clear soft_timeout when a task is re-registered Registering a name again replaces the task, but left the old timeout in the map for the new one to inherit. --- sdks/python/taskito/mixins/decorators.py | 8 ++++++-- sdks/python/tests/worker/test_lifecycle_parity.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/sdks/python/taskito/mixins/decorators.py b/sdks/python/taskito/mixins/decorators.py index 7d52ffd2..f4f46bcd 100644 --- a/sdks/python/taskito/mixins/decorators.py +++ b/sdks/python/taskito/mixins/decorators.py @@ -413,8 +413,12 @@ def decorator(fn: Callable) -> TaskWrapper: if final_inject: self._task_inject_map[task_name] = final_inject - # Wrap the function with hooks, middleware, and context - if soft_timeout is not None: + # Wrap the function with hooks, middleware, and context. Re-registering + # a name replaces the task, so an absent soft_timeout has to clear any + # earlier one rather than leave it for the new task to inherit. + if soft_timeout is None: + self._task_soft_timeouts.pop(task_name, None) + else: self._task_soft_timeouts[task_name] = soft_timeout wrapped = self._wrap_task(fn, task_name) diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 991870a3..2d4d9858 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -77,6 +77,20 @@ def _wait_for(predicate: Any, message: str) -> None: raise AssertionError(message) +def test_re_registering_a_task_clears_its_soft_timeout(queue: Queue) -> None: + """Replacing a task must not leave it wearing the old task's soft_timeout.""" + + @queue.task(name="dup", soft_timeout=0.05) + def first() -> None: ... + + assert queue._task_soft_timeouts["dup"] == 0.05 + + @queue.task(name="dup") + def second() -> None: ... + + assert "dup" not in queue._task_soft_timeouts + + def test_setup_failure_releases_what_it_acquired(tmp_path: Path) -> None: """A task that dies during setup must still hand back its resources. From 343e947a6c054c2a3d931b1a7b12f947cadabe8c Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:11:01 +0530 Subject: [PATCH 10/13] fix(python): report a torn-down job from native dispatch A BaseException that is not a cancellation slipped past every clause, so the job reported nothing and sat Running until the reaper called it a timeout. The queue stub is now concrete: a bare MagicMock invents whatever the lifecycle asks for, which is how a missing attribute takes the wrong branch in silence. --- sdks/python/taskito/async_support/executor.py | 26 ++++ sdks/python/tests/worker/test_native_async.py | 120 ++++++++++++++---- 2 files changed, 123 insertions(+), 23 deletions(-) diff --git a/sdks/python/taskito/async_support/executor.py b/sdks/python/taskito/async_support/executor.py index 77287342..60acc553 100644 --- a/sdks/python/taskito/async_support/executor.py +++ b/sdks/python/taskito/async_support/executor.py @@ -221,6 +221,32 @@ async def _run_job( ) ) + except BaseException as exc: + # Neither a cancellation nor an ordinary failure — KeyboardInterrupt, + # SystemExit, or any other BaseException the clause above cannot see. + # Unreported, the job sits Running until the reaper mislabels it a + # timeout, so report it as the blocking path would: a failure, with + # the retry policy left to decide. Hand off without awaiting — the + # loop may be going down — then re-raise, since swallowing these + # breaks interpreter and loop teardown. + if not reported: + wall_ns = time.monotonic_ns() - start_ns + reported = True + error_msg = encode_task_error(exc) + self._hand_off_now( + lambda: self._sender.try_report_failure( + job_id, + task_name, + error_msg, + retry_count, + max_retries, + wall_ns, + True, + ), + job_id, + ) + raise + finally: clear_async_context(token) diff --git a/sdks/python/tests/worker/test_native_async.py b/sdks/python/tests/worker/test_native_async.py index 6571ebb8..4770a285 100644 --- a/sdks/python/tests/worker/test_native_async.py +++ b/sdks/python/tests/worker/test_native_async.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock import cloudpickle +import pytest from taskito import Queue, TaskCancelledError, current_job from taskito.async_support.context import ( @@ -21,32 +22,76 @@ PollUntil = Any # the conftest fixture's runtime type -def _fake_queue(**overrides: Any) -> MagicMock: +class _FakeQueue: """A queue stub covering every attribute the shared lifecycle reads. - MagicMock invents any attribute asked of it, and the lifecycle branches on - ``is not None`` and on truthiness — so an unset attribute does not raise, it - silently takes the wrong branch. An absent ``_task_batch_configs`` is the - sharp one: the auto-made mock reads as a per-item batch config, and every - task then fails the return-type contract. Listing them here once means a new - lifecycle read is one edit rather than one per test. + Deliberately not a bare ``MagicMock``: that invents whatever it is asked for, + and the lifecycle branches on ``is not None`` and on truthiness, so an unset + attribute would not raise — it would quietly take the wrong branch. An absent + ``_task_batch_configs`` is the sharp one, since the auto-made mock reads as a + per-item batch config and every task then fails the return-type contract. + + ``__slots__`` makes the trade explicit: a lifecycle read this class does not + declare raises here, and so does a typo'd override, instead of turning into a + silently wrong test. """ - queue = MagicMock() - queue._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) - queue._serialize_result.side_effect = lambda _name, result: cloudpickle.dumps(result) - queue._interceptor = None - queue._proxy_registry = None - queue._proxy_metrics = None - queue._test_mode_active = False - queue._resource_runtime = None - queue._task_inject_map = {} - queue._task_retry_filters = {} - queue._task_predicates = {} - queue._task_batch_configs = {} - queue._task_soft_timeouts = {} - queue._workflow_tracker = None - queue._hooks = {"before_task": [], "after_task": [], "on_success": [], "on_failure": []} - queue._get_middleware_chain.return_value = [] + + __slots__ = ( + "_apply_dispatch_predicate", + "_deserialize_payload", + "_emit_event", + "_get_middleware_chain", + "_hooks", + "_interceptor", + "_max_reconstruction_timeout", + "_proxy_metrics", + "_proxy_registry", + "_recipe_signing_key", + "_resource_runtime", + "_serialize_result", + "_task_batch_configs", + "_task_inject_map", + "_task_predicates", + "_task_retry_filters", + "_task_soft_timeouts", + "_test_mode_active", + "_workflow_tracker", + ) + + def __init__(self) -> None: + self._deserialize_payload = MagicMock( + side_effect=lambda _name, data: cloudpickle.loads(data) + ) + self._serialize_result = MagicMock( + side_effect=lambda _name, result: cloudpickle.dumps(result) + ) + self._get_middleware_chain = MagicMock(return_value=[]) + self._apply_dispatch_predicate = MagicMock() + self._emit_event = MagicMock() + self._interceptor = None + self._proxy_registry = None + self._proxy_metrics = None + self._recipe_signing_key = None + self._max_reconstruction_timeout = None + self._test_mode_active = False + self._resource_runtime = None + self._workflow_tracker = None + self._task_inject_map: dict[str, Any] = {} + self._task_retry_filters: dict[str, Any] = {} + self._task_predicates: dict[str, Any] = {} + self._task_batch_configs: dict[str, Any] = {} + self._task_soft_timeouts: dict[str, Any] = {} + self._hooks: dict[str, list[Any]] = { + "before_task": [], + "after_task": [], + "on_success": [], + "on_failure": [], + } + + +def _fake_queue(**overrides: Any) -> Any: + """A :class:`_FakeQueue` with per-test overrides applied.""" + queue = _FakeQueue() for name, value in overrides.items(): setattr(queue, name, value) return queue @@ -528,6 +573,35 @@ def _payload() -> bytes: return payload +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_base_exception_reports_failure(poll_until: PollUntil) -> None: + """A non-cancellation BaseException must still reach a terminal report. + + KeyboardInterrupt and friends slip past `except Exception` just as + CancelledError does, but they are not cancellations — unreported, the job + sits Running until the reaper mislabels it a timeout. + + The interrupt is re-raised after reporting, which is the contract (swallowing + it would break loop teardown), so it surfaces on the executor thread and + pytest notices — expected here, hence the filter. + """ + sender = MagicMock() + sender.try_report_failure.return_value = True + + async def interrupted_task() -> None: + raise KeyboardInterrupt + + executor = _backpressure_executor(sender, interrupted_task, "mod.interrupted_task") + executor.submit_job("job-k", "mod.interrupted_task", _payload(), 0, 3, "default") + poll_until( + lambda: sender.try_report_failure.called, + message="KeyboardInterrupt was not reported", + ) + executor.stop() + + assert not sender.try_report_cancelled.called, "an interrupt is not a cancellation" + + def test_cancelled_coroutine_reports_cancelled(poll_until: PollUntil) -> None: """asyncio.CancelledError is a BaseException, so `except Exception` misses it. From 3741c959759ac11bccabebc5e8c32eb39491bb43 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:11:29 +0530 Subject: [PATCH 11/13] test(python): stop the parity tests hiding and causing stalls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stop helper returned on a live worker, which is how the shutdown stall in this series nearly escaped: assertions had already passed, so only the clock showed it. The hooks also held live exceptions — the stall's own trigger. --- .../tests/worker/test_lifecycle_parity.py | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 2d4d9858..986652c9 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -63,8 +63,15 @@ def _start_worker(queue: Queue) -> threading.Thread: def _stop_worker(queue: Queue, thread: threading.Thread) -> None: + """Stop the worker, and fail if it does not actually stop. + + Returning quietly on a live thread is how a shutdown regression hides: the + assertions have already passed by this point, so the test still goes green + and only the wall clock shows anything is wrong. + """ queue._inner.request_shutdown() thread.join(timeout=10) + assert not thread.is_alive(), "worker did not stop within 10s" def _wait_for(predicate: Any, message: str) -> None: @@ -194,11 +201,15 @@ def hooked(a: int, b: int) -> int: @BOTH_PATHS def test_on_failure_hook_sees_the_exception(noretry_queue: Queue, is_async: bool) -> None: """on_failure receives the live exception for either kind of task.""" - failures: list[BaseException] = [] + # Record what the hook saw, not the exception itself: holding a live + # exception here would pin its traceback — the documented condition that + # stalls worker shutdown once dispatch goes native — and turn every failure + # test into a slow one for no added assurance. + failures: list[tuple[type, str]] = [] @noretry_queue.on_failure def on_fail(task_name: str, args: tuple, kwargs: dict, error: BaseException) -> None: - failures.append(error) + failures.append((type(error), str(error))) if is_async: @@ -219,8 +230,8 @@ def boom() -> None: finally: _stop_worker(noretry_queue, thread) - assert isinstance(failures[0], ValueError) - assert "expected failure" in str(failures[0]) + assert failures[0][0] is ValueError + assert "expected failure" in failures[0][1] @BOTH_PATHS @@ -390,12 +401,13 @@ def logged() -> str: @BOTH_PATHS def test_middleware_after_receives_the_error(tmp_path: Path, is_async: bool) -> None: """Middleware .after() receives the exception, not None, for either kind.""" - seen: list[Any] = [] + seen: list[tuple[type, str] | None] = [] ran = threading.Event() class Recorder(TaskMiddleware): def after(self, job: Any, result: Any, error: Any) -> None: - seen.append(error) + # Type and message only — see the note in the on_failure test. + seen.append(None if error is None else (type(error), str(error))) ran.set() q = Queue(db_path=str(tmp_path / "mw.db"), workers=2, default_retry=0, middleware=[Recorder()]) @@ -419,7 +431,8 @@ def failing() -> None: finally: _stop_worker(q, thread) - assert isinstance(seen[0], RuntimeError), f"after() should see the error, got {seen[0]!r}" + assert seen[0] is not None, "after() should see the error, not None" + assert seen[0][0] is RuntimeError, f"after() should see the error, got {seen[0]!r}" finally: q.close() @@ -437,7 +450,7 @@ def test_base_exception_is_not_reported_as_completed(tmp_path: Path, is_async: b and every loop shutdown raises one. """ completed: list[dict] = [] - after_task_errors: list[Any] = [] + after_task_errors: list[tuple[type, str] | None] = [] ran = threading.Event() class Killed(BaseException): @@ -449,7 +462,8 @@ class Killed(BaseException): @q.after_task def on_after(task_name: str, args: tuple, kwargs: dict, result: Any, error: Any) -> None: - after_task_errors.append(error) + # Type and message only — see the note in the on_failure test. + after_task_errors.append(None if error is None else (type(error), str(error))) ran.set() if is_async: @@ -472,7 +486,8 @@ def killed() -> None: _stop_worker(q, thread) assert completed == [], f"a torn-down task must not emit JOB_COMPLETED, got {completed}" - assert isinstance(after_task_errors[0], Killed), ( + assert after_task_errors[0] is not None, "after_task saw no error for a torn-down task" + assert after_task_errors[0][0] is Killed, ( f"after_task should see the BaseException, got {after_task_errors[0]!r}" ) finally: From 68272d9528a09e0355d4dc72ee7363b37edabe96 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:30:23 +0530 Subject: [PATCH 12/13] fix(python): pair middleware after() when setup fails `completed_mw` exists so every before() gets an after(); the setup unwind released resources but skipped it, leaving middleware-owned state behind. --- sdks/python/taskito/task_lifecycle.py | 16 +++++++++--- .../tests/worker/test_lifecycle_parity.py | 25 ++++++++++++++----- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/sdks/python/taskito/task_lifecycle.py b/sdks/python/taskito/task_lifecycle.py index fc1d13bb..d4b90f6e 100644 --- a/sdks/python/taskito/task_lifecycle.py +++ b/sdks/python/taskito/task_lifecycle.py @@ -200,12 +200,20 @@ async def run_lifecycle( comp_ctx = saga.take_compensation_context(job_id) if comp_ctx is not None: comp_ctx_token = _set_compensation_context(comp_ctx) - except BaseException: - # Setup failed, so the body never runs and its teardown never will. Give - # back only what was acquired: after_task and the lifecycle events belong - # to a task that started, and this one did not. + except BaseException as exc: + # Setup failed, so the body never runs and its teardown never will. Undo + # what setup did, and no more: after_task and the lifecycle events belong + # to a task that started, and this one did not. Middleware is the + # exception — `completed_mw` exists to pair every before() with an + # after(), and a middleware that set something up is owed the chance to + # take it down whether or not the task it prepared for ever ran. if comp_ctx_token is not None: _reset_compensation_context(comp_ctx_token) + for mw in completed_mw: + try: + mw.after(current_job, None, exc) + except Exception: + logger.exception("middleware after() error") _release_acquired(release_callbacks, proxy_cleanup, queue_ref) raise diff --git a/sdks/python/tests/worker/test_lifecycle_parity.py b/sdks/python/tests/worker/test_lifecycle_parity.py index 986652c9..7d40f72d 100644 --- a/sdks/python/tests/worker/test_lifecycle_parity.py +++ b/sdks/python/tests/worker/test_lifecycle_parity.py @@ -98,16 +98,23 @@ def second() -> None: ... assert "dup" not in queue._task_soft_timeouts -def test_setup_failure_releases_what_it_acquired(tmp_path: Path) -> None: - """A task that dies during setup must still hand back its resources. +def test_setup_failure_undoes_what_setup_did(tmp_path: Path) -> None: + """A task that dies during setup must still be torn down. - Resources are injected before the before_task hooks run, so a raising hook - leaves them held — and the teardown that would release them only runs once - the body is entered, which it never is. + Resources are injected and middleware prepared before the before_task hooks + run, so a raising hook leaves them held — and the teardown that would undo + them only runs once the body is entered, which it never is. """ released: list[str] = [] + torn_down: list[tuple[type, str]] = [] - q = Queue(db_path=str(tmp_path / "setup.db"), workers=2, default_retry=0) + class Recorder(TaskMiddleware): + def after(self, job: Any, result: Any, error: Any) -> None: + torn_down.append((type(error), str(error))) + + q = Queue( + db_path=str(tmp_path / "setup.db"), workers=2, default_retry=0, middleware=[Recorder()] + ) try: # Request scope, because its release calls `teardown` outright — a # task-scoped resource goes back to a pool, where nothing observes it. @@ -115,6 +122,7 @@ def test_setup_failure_releases_what_it_acquired(tmp_path: Path) -> None: def make_db() -> str: return "conn" + # Raises after the resource is injected and Recorder.before() has run. @q.before_task def boom(task_name: str, args: tuple, kwargs: dict) -> None: raise RuntimeError("hook blew up during setup") @@ -127,8 +135,13 @@ def needs_db(db: Any = None) -> str: thread = _start_worker(q) try: _wait_for(lambda: released, "the request-scoped resource was never released") + _wait_for(lambda: torn_down, "middleware that ran before() never got its after()") finally: _stop_worker(q, thread) + + assert torn_down[0][0] is RuntimeError, ( + f"after() should see the error, got {torn_down[0]!r}" + ) finally: q.close() From 06cafb000c8b4194b0cb57820baa9914e2a6cd0b Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 17 Jul 2026 03:30:39 +0530 Subject: [PATCH 13/13] test(python): assert the executor thread stops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `stop()` only logs when its join times out, so the test could pass on a leaked thread — the same way the worker stop helper hid a stall. --- sdks/python/tests/worker/test_native_async.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sdks/python/tests/worker/test_native_async.py b/sdks/python/tests/worker/test_native_async.py index 4770a285..25ad210e 100644 --- a/sdks/python/tests/worker/test_native_async.py +++ b/sdks/python/tests/worker/test_native_async.py @@ -600,6 +600,9 @@ async def interrupted_task() -> None: executor.stop() assert not sender.try_report_cancelled.called, "an interrupt is not a cancellation" + # `stop()` only logs when its join times out, and the thread-exception warning + # is filtered here — without this the test could go green on a leaked thread. + assert not executor._thread.is_alive(), "the executor thread outlived stop()" def test_cancelled_coroutine_reports_cancelled(poll_until: PollUntil) -> None: