fix(receiver): harden callback draining during shutdown

- keep callback tasks owned until cleanup completes
- prevent AnyIO level-cancellation busy loops
- preserve interruptible unbounded graceful waits
- release execution capacity exactly once
- clarify wait-tasks-timeout semantics
- add deterministic lifecycle and cancellation tests

Author: GefMar

docs/guide/cli.md

@@ -143,10 +143,10 @@ kill -HUP <main pid>
 
 ### Graceful and force shutdowns
 
-If you send `SIGINT` or `SIGKILL` to the main process by pressing <kbd>Ctrl</kbd>+<kbd>C</kbd> or using the `kill` command, it will initiate the shutdown process.
+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 too long for tasks to complete each time you shut down the worker, you can either send termination signals three times to the main process to perform a hard kill or configure the `--wait-tasks-timeout` to set a hard time limit for shutting down.
+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.
 
 ::: tip Cool tip
 The number of signals before a hard kill can be configured with the `--hardkill-count` CLI argument.
@@ -171,7 +171,7 @@ The number of signals before a hard kill can be configured with the `--hardkill-
 * `--max-tasks-per-child` - maximum number of tasks to be executed by a single worker process before restart.
 * `--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` - if cannot read new messages from the broker or maximum number of tasks is reached, worker will wait for all current tasks to finish. This parameter sets the maximum amount of time to wait until shutdown.
+* `--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.
 
 ## Scheduler

taskiq/cli/worker/args.py

@@ -259,7 +259,10 @@ def from_cli(
             "--wait-tasks-timeout",
             type=float,
             default=None,
-            help="Maximum time to wait for all current tasks to finish before exiting.",
+            help=(
+                "Grace period for current task callbacks before cancellation is "
+                "requested during shutdown."
+            ),
         )
         parser.add_argument(
             "--hardkill-count",

taskiq/receiver/receiver.py

@@ -656,19 +656,12 @@ async def runner(
         :param queue: queue with prefetched data.
         """
         tasks: set[asyncio.Task[Any]] = set()
-
-        while True:
-            try:
+        graceful_shutdown = False
+        try:
+            while True:
                 queued_message = await queue.get()
                 if queued_message is _QueueSignal.DONE:
-                    # asyncio.wait will throw an error if there is nothing to wait for
-                    if tasks:
-                        logger.info(
-                            "Waiting for %d running tasks to complete...",
-                            len(tasks),
-                        )
-                        await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
-                        logger.info("No more tasks to wait for. Shutting down.")
+                    graceful_shutdown = True
                     break
                 execution_semaphore = self.sem
                 owns_execution_slot = execution_semaphore is not None
@@ -684,23 +677,68 @@ async def runner(
                 )
                 tasks.add(started_callback.task)
 
-                # We want the task to remove itself from the set when it's done.
-                #
-                # Because if we won't save it anywhere,
-                # python's GC can silently cancel task
-                # and this behaviour considered to be a Hisenbug.
-                # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
                 started_callback.task.add_done_callback(
                     functools.partial(
                         self._on_callback_done,
                         active_tasks=tasks,
                         owns_delivery_slot=started_callback.owns_delivery_slot,
                     ),
                 )
+        except asyncio.CancelledError:
+            pass
+        finally:
+            await self._drain_active_tasks(
+                tasks,
+                cancel_immediately=not graceful_shutdown,
+            )
+            logger.info("The runner is stopped.")
 
+    async def _run_owned_callback(self, message: bytes | AckableMessage) -> None:
+        """Run one callback outside repeated listener-scope cancellation."""
+        with anyio.CancelScope(shield=True):
+            await self.callback(message=message, raise_err=False)
+
+    async def _drain_active_tasks(
+        self,
+        tasks: set[asyncio.Task[Any]],
+        *,
+        cancel_immediately: bool,
+    ) -> None:
+        """Wait for callbacks and cancel work beyond the graceful boundary."""
+        tasks = set(tasks)
+        if not tasks:
+            return
+
+        logger.info("Waiting for %d running tasks to complete...", len(tasks))
+        if cancel_immediately:
+            pending = {task for task in tasks if not task.done()}
+        else:
+            try:
+                _, pending = await asyncio.wait(
+                    tasks,
+                    timeout=self.wait_tasks_timeout,
+                )
             except asyncio.CancelledError:
-                break
-        logger.info("The runner is stopped.")
+                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(
+        tasks: set[asyncio.Task[Any]],
+    ) -> None:
+        """Cancel callbacks once and await cleanup despite outer cancellation."""
+        for task in 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()
 
     def _start_callback(
         self,
@@ -716,7 +754,7 @@ def _start_callback(
                 owns_delivery_slot = False
             return _StartedCallback(
                 task=asyncio.create_task(
-                    self.callback(message=message.data, raise_err=False),
+                    self._run_owned_callback(message=message.data),
                 ),
                 owns_delivery_slot=owns_delivery_slot,
             )
@@ -737,12 +775,24 @@ def _on_callback_done(
         active_tasks: set[asyncio.Task[Any]],
         owns_delivery_slot: bool,
     ) -> None:
-        """Release capacity transferred to a completed callback task."""
+        """Release callback capacity and retrieve unexpected failures."""
         active_tasks.discard(task)
         if self.sem is not None:
             self.sem.release()
         if owns_delivery_slot:
             self.sem_prefetch.release()
+        if task.cancelled():
+            return
+        task_exception = task.exception()
+        if task_exception is not None:
+            logger.error(
+                "Receiver callback failed outside task execution handling.",
+                exc_info=(
+                    type(task_exception),
+                    task_exception,
+                    task_exception.__traceback__,
+                ),
+            )
 
     def _record_listen_error(self, error: BaseException) -> None:
         """Preserve the first listener error and report cleanup failures."""

tests/receiver/receiver_lifecycle_support.py

@@ -0,0 +1,139 @@
+import asyncio
+from collections.abc import Awaitable
+from typing import Any
+
+from taskiq.abc.broker import AckableMessage
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.receiver.receiver import Receiver, _PrefetchedMessage, _QueueSignal
+from tests.receiver.receiver_listener_support import (
+    ReceiverLifecycleError,
+    assert_semaphore_capacity,
+)
+
+
+class ShieldCallCounter:
+    """Count shield calls made by one named asyncio task."""
+
+    def __init__(self, task_name: str) -> None:
+        self.task_name = task_name
+        self.runner_calls = 0
+        self._shield = asyncio.shield
+
+    def __call__(self, awaitable: Awaitable[Any]) -> asyncio.Future[Any]:
+        current_task = asyncio.current_task()
+        if current_task is not None and current_task.get_name() == self.task_name:
+            self.runner_calls += 1
+        return self._shield(awaitable)
+
+
+class ReceiverQueue(asyncio.Queue[_PrefetchedMessage | _QueueSignal]):
+    """Expose runner reads and graceful-shutdown sentinel delivery."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.get_started = asyncio.Event()
+        self.shutdown_received = asyncio.Event()
+
+    async def get(self) -> _PrefetchedMessage | _QueueSignal:
+        self.get_started.set()
+        message = await super().get()
+        if message is _QueueSignal.DONE:
+            self.shutdown_received.set()
+        return message
+
+
+class ControlledReceiver(Receiver):
+    """Receiver with deterministic callback and cleanup checkpoints."""
+
+    def __init__(
+        self,
+        *,
+        wait_tasks_timeout: float | None,
+        max_async_tasks: int | None = 1,
+        fail: bool = False,
+    ) -> None:
+        super().__init__(
+            InMemoryBroker(),
+            max_async_tasks=max_async_tasks,
+            run_startup=False,
+            wait_tasks_timeout=wait_tasks_timeout,
+        )
+        self.fail = fail
+        self.started_callbacks: asyncio.Queue[None] = asyncio.Queue()
+        self.release_callback = asyncio.Event()
+        self.cleanup_started_callbacks: asyncio.Queue[None] = asyncio.Queue()
+        self.release_cleanup = asyncio.Event()
+        self.finished_callbacks: asyncio.Queue[None] = asyncio.Queue()
+        self.callback_tasks: list[asyncio.Task[None]] = []
+
+    @property
+    def callback_task(self) -> asyncio.Task[None] | None:
+        """Return the most recently started callback task."""
+        if not self.callback_tasks:
+            return None
+        return self.callback_tasks[-1]
+
+    async def callback(
+        self,
+        message: bytes | AckableMessage,
+        raise_err: bool = False,
+    ) -> None:
+        del message, raise_err
+        callback_task = asyncio.current_task()
+        assert callback_task is not None
+        self.callback_tasks.append(callback_task)
+        self.started_callbacks.put_nowait(None)
+        try:
+            await self.release_callback.wait()
+            if self.fail:
+                raise ReceiverLifecycleError("callback failed")
+        finally:
+            self.cleanup_started_callbacks.put_nowait(None)
+            await self.release_cleanup.wait()
+            for _ in range(20):
+                await asyncio.sleep(0)
+            self.finished_callbacks.put_nowait(None)
+
+    async def settle(self, runner_task: asyncio.Task[None]) -> None:
+        """Release all checkpoints and settle test-owned tasks."""
+        self.release_callback.set()
+        self.release_cleanup.set()
+        callback_tasks = set(self.callback_tasks)
+        for callback_task in callback_tasks:
+            if not callback_task.done():
+                callback_task.cancel()
+        if not runner_task.done():
+            runner_task.cancel()
+        await asyncio.gather(
+            *callback_tasks,
+            runner_task,
+            return_exceptions=True,
+        )
+
+
+async def start_callback(
+    receiver: ControlledReceiver,
+) -> tuple[ReceiverQueue, asyncio.Task[None]]:
+    """Start one controlled callback through the real runner boundary."""
+    queue = ReceiverQueue()
+    await queue.put(_PrefetchedMessage(b"payload", owns_delivery_slot=False))
+    runner_task = asyncio.create_task(receiver.runner(queue))
+    await wait_for_signals(receiver.started_callbacks)
+    return queue, runner_task
+
+
+async def wait_for_signals(
+    signals: asyncio.Queue[None],
+    count: int = 1,
+) -> None:
+    """Wait for an exact number of deterministic lifecycle checkpoints."""
+    await asyncio.wait_for(
+        asyncio.gather(*(signals.get() for _ in range(count))),
+        timeout=1,
+    )
+
+
+async def assert_exact_capacity(receiver: Receiver, slots: int = 1) -> None:
+    """Assert that exactly the expected execution permits were returned."""
+    assert receiver.sem is not None
+    await assert_semaphore_capacity(receiver.sem, slots)