fix(receiver): preserve task finalization during shutdown

- Defer shutdown cancellation during dependency teardown and executor calls.
- Preserve cancellation across nested execution without interrupting cleanup.
- Cover thread and process executors, retries, and reentrant finalization.
- Clarify graceful shutdown limits and per-worker hard-kill handling.

Author: GefMar

docs/guide/cli.md

@@ -146,10 +146,10 @@ kill -HUP <main pid>
 If you send `SIGINT` or `SIGTERM` to the main process by pressing <kbd>Ctrl</kbd>+<kbd>C</kbd> or using the `kill` command, it will initiate the shutdown process.
 By default, it will stop fetching new messages immediately after receiving the signal but will wait for the completion of all currently executing tasks.
 
-If you don't want to wait indefinitely for tasks to complete, configure `--wait-tasks-timeout` to set their graceful completion period. Once that period expires, the worker requests cancellation of asynchronous task callbacks that are still running and waits for their cleanup before continuing shutdown. Synchronous functions already running in a thread or process executor cannot be forcibly stopped by asyncio cancellation and may continue until the executor shuts down. This is not an absolute process deadline; repeat the termination signal until the configured hard-kill threshold is reached if the process must stop immediately.
+If you don't want to wait indefinitely for tasks to complete, configure `--wait-tasks-timeout` to set their graceful completion period. Once that period expires, the worker requests cancellation of asynchronous task callbacks that are still running and awaits their completion before continuing shutdown. Synchronous functions already running in a thread or process executor cannot be forcibly stopped by asyncio cancellation. Shutdown cancellation is deferred while awaiting such a call so that task dependencies remain usable until the function returns. This is not an absolute process deadline. Enforce a hard deadline through your process supervisor, with forced termination configured for the worker processes and their executor subprocesses; forced termination does not guarantee task cleanup.
 
 ::: tip Cool tip
-The number of signals before a hard kill can be configured with the `--hardkill-count` CLI argument.
+The `--hardkill-count` CLI argument counts termination signals received by each worker process. Additional signals sent only to the main process are not forwarded after shutdown begins.
 :::
 
 
@@ -172,7 +172,7 @@ The number of signals before a hard kill can be configured with the `--hardkill-
 * `--max-fails` - Maximum number of child process exits.
 * `--shutdown-timeout` - maximum amount of time for graceful broker's shutdown in seconds (default 5).
 * `--wait-tasks-timeout` - graceful completion period for current tasks during shutdown. Cancellation is requested for callbacks still running after the period, and their cleanup is awaited. The default `None` waits without a timeout.
-* `--hardkill-count` - Number of termination signals to the main process before performing a hardkill.
+* `--hardkill-count` - per-worker termination-signal threshold for the worker's hard-kill handling.
 
 ## Scheduler
 

taskiq/cli/worker/args.py

@@ -268,8 +268,7 @@ def from_cli(
             "--hardkill-count",
             type=int,
             default=3,
-            help="Number of termination signals to the main "
-            "process before performing a hardkill.",
+            help="Per-worker termination-signal threshold for hard-kill handling.",
         )
         parser.add_argument(
             "--use-process-pool",

taskiq/receiver/receiver.py

@@ -5,6 +5,7 @@
 import inspect
 import random
 import sys
+from collections import Counter
 from collections.abc import AsyncGenerator, Callable
 from concurrent.futures import Executor, ProcessPoolExecutor
 from dataclasses import dataclass
@@ -111,6 +112,11 @@ def __init__(
         self.known_tasks: set[str] = set()
         self.max_tasks_to_execute = max_tasks_to_execute
         self.wait_tasks_timeout = wait_tasks_timeout
+        self._executing_tasks: Counter[asyncio.Task[Any]] = Counter()
+        self._finishing_tasks: Counter[asyncio.Task[Any]] = Counter()
+        self._sync_tasks: set[asyncio.Task[Any]] = set()
+        # Execution depth receiving cancellation, or None for a deferred request.
+        self._shutdown_cancellations: dict[asyncio.Task[Any], int | None] = {}
         self._listen_error: BaseException | None = None
         if max_prefetch < 0:
             raise ValueError("max_prefetch cannot be negative.")
@@ -312,6 +318,9 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
             if PY_VERSION <= (3, 13)
             else inspect.iscoroutinefunction
         )
+        execution_task = asyncio.current_task()
+        if execution_task is not None:
+            self._executing_tasks[execution_task] += 1
         try:
             # We put kwargs resolving here,
             # to be able to catch any exception (for example ),
@@ -348,6 +357,7 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
                         ctx.run,
                         func,
                     )
+                target_future = self._await_sync_result(target_future)
             timeout = message.labels.get("timeout")
             if timeout is not None:
                 if not is_coroutine:
@@ -374,41 +384,99 @@ async def run_task(  # noqa: C901, PLR0912, PLR0915
         except BaseException as exc:
             found_exception = exc
             logger.exception("Exception found while executing function.")
+        finally:
+            if execution_task is not None:
+                if self._shutdown_cancellations.get(execution_task) == (
+                    self._executing_tasks[execution_task]
+                ):
+                    self._shutdown_cancellations[execution_task] = None
+                self._executing_tasks[execution_task] -= 1
+                if not self._executing_tasks[execution_task]:
+                    del self._executing_tasks[execution_task]
         # Stop the timer.
         execution_time = time() - start_time
-        if dep_ctx:
-            args = (None, None, None)
-            if found_exception and self.propagate_exceptions:
-                args = (  # type: ignore
-                    type(found_exception),
-                    found_exception,
-                    found_exception.__traceback__,
-                )
-            await dep_ctx.close(*args)
-
-        # Assemble result.
-        result: TaskiqResult[Any] = TaskiqResult(
-            is_err=found_exception is not None,
-            log=None,
-            return_value=returned,
-            execution_time=round(execution_time, 2),
-            error=found_exception,
-            labels=message.labels,
-        )
-        # If exception is found we execute middlewares.
-        if found_exception is not None:
-            for middleware in reversed(self.broker.middlewares):
-                if middleware.__class__.on_error != TaskiqMiddleware.on_error:
-                    await maybe_awaitable(
-                        middleware.on_error(
-                            message,
-                            result,
-                            found_exception,
-                        ),
+        if execution_task is not None:
+            self._finishing_tasks[execution_task] += 1
+        try:
+            if dep_ctx:
+                args = (None, None, None)
+                if found_exception and self.propagate_exceptions:
+                    args = (  # type: ignore
+                        type(found_exception),
+                        found_exception,
+                        found_exception.__traceback__,
                     )
-
+                await dep_ctx.close(*args)
+
+            # Assemble result.
+            result: TaskiqResult[Any] = TaskiqResult(
+                is_err=found_exception is not None,
+                log=None,
+                return_value=returned,
+                execution_time=round(execution_time, 2),
+                error=found_exception,
+                labels=message.labels,
+            )
+            # If exception is found we execute middlewares.
+            if found_exception is not None:
+                for middleware in reversed(self.broker.middlewares):
+                    if middleware.__class__.on_error != TaskiqMiddleware.on_error:
+                        await maybe_awaitable(
+                            middleware.on_error(
+                                message,
+                                result,
+                                found_exception,
+                            ),
+                        )
+        finally:
+            if execution_task is not None:
+                self._finishing_tasks[execution_task] -= 1
+                if not self._finishing_tasks[execution_task]:
+                    del self._finishing_tasks[execution_task]
+                    if execution_task in self._shutdown_cancellations:
+                        asyncio.get_running_loop().call_soon(
+                            self._deliver_deferred_cancellation,
+                            execution_task,
+                        )
+
+        if execution_task is not None and self._deliver_deferred_cancellation(
+            execution_task,
+        ):
+            await asyncio.sleep(0)
         return result
 
+    def _deliver_deferred_cancellation(self, task: asyncio.Task[Any]) -> bool:
+        """Deliver pending shutdown only when the task resumes execution."""
+        if (
+            task not in self._shutdown_cancellations
+            or task in self._finishing_tasks
+            or task in self._sync_tasks
+        ):
+            return False
+        if task not in self._executing_tasks or task.done():
+            del self._shutdown_cancellations[task]
+            return False
+        if self._shutdown_cancellations[task] is None:
+            self._shutdown_cancellations[task] = self._executing_tasks[task]
+            return task.cancel()
+        return False
+
+    async def _await_sync_result(self, future: asyncio.Future[Any]) -> Any:
+        """Defer receiver shutdown while the executor still uses dependencies."""
+        task = asyncio.current_task()
+        if task is not None:
+            self._sync_tasks.add(task)
+        try:
+            return await future
+        finally:
+            if task is not None:
+                self._sync_tasks.discard(task)
+                if task in self._shutdown_cancellations:
+                    asyncio.get_running_loop().call_soon(
+                        self._deliver_deferred_cancellation,
+                        task,
+                    )
+
     async def listen(self, finish_event: asyncio.Event) -> None:  # pragma: no cover
         """
         This function iterates over tasks asynchronously.
@@ -722,23 +790,36 @@ async def _drain_active_tasks(
                 pending = {task for task in tasks if not task.done()}
 
         if pending:
-            logger.warning("Cancelling %d running callback tasks.", len(pending))
             with anyio.CancelScope(shield=True):
                 await self._cancel_callback_tasks(pending)
         logger.info("No more tasks to wait for. Shutting down.")
 
-    @staticmethod
     async def _cancel_callback_tasks(
+        self,
         tasks: set[asyncio.Task[Any]],
     ) -> None:
-        """Cancel callbacks once and await cleanup despite outer cancellation."""
-        for task in tasks:
+        """Cancel execution without interrupting task finalization."""
+        cancellable_tasks = tasks.difference(self._finishing_tasks, self._sync_tasks)
+        self._shutdown_cancellations.update(
+            (task, self._executing_tasks[task] if task in cancellable_tasks else None)
+            for task in tasks
+        )
+        if cancellable_tasks:
+            logger.warning(
+                "Cancelling %d running callback tasks.",
+                len(cancellable_tasks),
+            )
+        for task in cancellable_tasks:
             task.cancel()
         waiter = asyncio.gather(*tasks, return_exceptions=True)
-        while not waiter.done():
-            with contextlib.suppress(asyncio.CancelledError):
-                await asyncio.shield(waiter)
-        waiter.result()
+        try:
+            while not waiter.done():
+                with contextlib.suppress(asyncio.CancelledError):
+                    await asyncio.shield(waiter)
+            waiter.result()
+        finally:
+            for task in tasks:
+                self._shutdown_cancellations.pop(task, None)
 
     def _start_callback(
         self,

tests/receiver/receiver_lifecycle_support.py

@@ -1,14 +1,129 @@
 import asyncio
-from collections.abc import Awaitable
+import contextvars
+import os
+import threading
+from collections.abc import AsyncGenerator, Awaitable
+from dataclasses import dataclass
+from io import StringIO
+from pathlib import Path
 from typing import Any
 
+import anyio
+from taskiq_dependencies import Depends
+
 from taskiq.abc.broker import AckableMessage
 from taskiq.brokers.inmemory_broker import InMemoryBroker
 from taskiq.receiver.receiver import Receiver, _PrefetchedMessage, _QueueSignal
+from taskiq.state import TaskiqState
 from tests.receiver.receiver_listener_support import (
     ReceiverLifecycleError,
     assert_semaphore_capacity,
 )
+from tests.utils import AsyncQueueBroker
+
+process_resource_broker = AsyncQueueBroker()
+
+
+@dataclass
+class ProcessResourceProbe:
+    """Share a resource path and explicit checkpoints with a spawned worker."""
+
+    resource: Path
+    started: threading.Event
+    release: threading.Event
+
+
+async def process_resource_dependency(
+    state: TaskiqState = Depends(),
+) -> AsyncGenerator[ProcessResourceProbe, None]:
+    """Keep the resource available until the receiver closes the dependency."""
+    probe: ProcessResourceProbe = state.process_resource
+    try:
+        yield probe
+    finally:
+        probe.resource.unlink()
+
+
+@process_resource_broker.task
+def read_process_resource(
+    probe: ProcessResourceProbe = Depends(process_resource_dependency),
+) -> tuple[int, str]:
+    """Read the resource in an importable, spawn-compatible task function."""
+    probe.started.set()
+    if not probe.release.wait(timeout=15):
+        raise TimeoutError("Process resource probe was not released")
+    return os.getpid(), probe.resource.read_text()
+
+
+class SyncResourceProbe:
+    """Keep a real executor function alive while its dependency is inspected."""
+
+    def __init__(self) -> None:
+        self.loop = asyncio.get_running_loop()
+        self.started = asyncio.Event()
+        self.finished = asyncio.Event()
+        self.release = threading.Event()
+        self.resource = StringIO("completed")
+        self.continuation_started = asyncio.Event()
+        self.continuation_finished = asyncio.Event()
+        self.release_continuation = asyncio.Event()
+
+    async def dependency(self) -> AsyncGenerator[StringIO, None]:
+        try:
+            yield self.resource
+        finally:
+            self.resource.close()
+
+    def run(self, resource: StringIO) -> str:
+        self.loop.call_soon_threadsafe(self.started.set)
+        try:
+            assert self.release.wait(timeout=10)
+            return resource.getvalue()
+        finally:
+            self.loop.call_soon_threadsafe(self.finished.set)
+
+    async def continuation(self) -> str:
+        self.continuation_started.set()
+        try:
+            await self.release_continuation.wait()
+            return self.resource.getvalue()
+        finally:
+            assert not self.resource.closed
+            self.continuation_finished.set()
+
+
+class DependencyCleanupProbe:
+    """Expose teardown checkpoints while enforcing the original task/context."""
+
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+        self.release = asyncio.Event()
+        self.finished = asyncio.Event()
+        self.context = contextvars.ContextVar("dependency-cleanup", default=False)
+
+    async def dependency(self) -> AsyncGenerator[None, None]:
+        owner = asyncio.current_task()
+        token = self.context.set(True)
+        try:
+            yield
+        finally:
+            self.started.set()
+            await self.release.wait()
+            for _ in range(20):
+                await asyncio.sleep(0)
+            assert asyncio.current_task() is owner
+            self.context.reset(token)
+            self.finished.set()
+
+
+async def listen_in_scope(
+    receiver: Receiver,
+    finish_event: asyncio.Event,
+    scope: anyio.CancelScope,
+) -> None:
+    """Expose listener cancellation through a real AnyIO scope."""
+    with scope:
+        await receiver.listen(finish_event)
 
 
 class ShieldCallCounter: