fix(receiver): harden callback draining during shutdown

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

@@ -256,7 +256,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

@@ -6,6 +6,7 @@
 import sys
 from collections.abc import Callable
 from concurrent.futures import Executor, ProcessPoolExecutor
+from contextlib import suppress
 from logging import getLogger
 from time import time
 from typing import Any, get_type_hints
@@ -78,6 +79,7 @@ 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._active_tasks: set[asyncio.Task[Any]] = set()
         for task in self.broker.get_all_tasks().values():
             self._prepare_task(task.task_name, task.original_func)
         self.sem: asyncio.Semaphore | None = None
@@ -461,63 +463,111 @@ async def runner(
 
         :param queue: queue with prefetched data.
         """
-        tasks: set[asyncio.Task[Any]] = set()
-
-        def task_cb(task: "asyncio.Task[Any]") -> None:
-            """
-            Callback for tasks.
-
-            This function used to remove task
-            from the list of active tasks and release
-            the semaphore, so other tasks can use it.
-
-            :param task: finished task
-            """
-            tasks.discard(task)
-            if self.sem is not None:
-                self.sem.release()
-
-        while True:
-            try:
+        capacity_acquired = False
+        graceful_shutdown = False
+        try:
+            while True:
                 # Waits for semaphore to be released.
                 if self.sem is not None:
                     await self.sem.acquire()
+                    capacity_acquired = True
 
                 self.sem_prefetch.release()
                 message = await queue.get()
                 if message is QUEUE_DONE:
-                    # asyncio.wait will throw an error if there is nothing to wait for
-                    if tasks:
-                        logger.info(
-                            f"Waiting for {len(tasks)} running tasks to complete...",
-                        )
-                        await asyncio.wait(tasks, timeout=self.wait_tasks_timeout)
-                        logger.info("No more tasks to wait for. Shutting down.")
+                    graceful_shutdown = True
                     break
 
-                # Custom hooks for OTel and any future instrumentations
-                for middleware in reversed(self.broker.middlewares):
-                    if hasattr(middleware, "on_prefetch_queue_remove"):
-                        await maybe_awaitable(
-                            middleware.on_prefetch_queue_remove(),  # type: ignore
-                        )
+                await self._notify_prefetch_queue_remove()
+                self._start_callback(message)
+                capacity_acquired = False
+        except asyncio.CancelledError:
+            pass
+        finally:
+            if capacity_acquired and self.sem is not None:
+                self.sem.release()
+            await self._drain_active_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)
 
-                task = asyncio.create_task(
-                    self.callback(message=message, raise_err=False),
+    async def _notify_prefetch_queue_remove(self) -> None:
+        """Run instrumentation before transferring callback ownership."""
+        for middleware in reversed(self.broker.middlewares):
+            if hasattr(middleware, "on_prefetch_queue_remove"):
+                await maybe_awaitable(
+                    middleware.on_prefetch_queue_remove(),  # type: ignore
                 )
-                tasks.add(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/
-                task.add_done_callback(task_cb)
+    def _start_callback(self, message: bytes | AckableMessage) -> None:
+        """Transfer the current execution slot to an owned callback task."""
+        task = asyncio.create_task(self._run_owned_callback(message))
+        self._active_tasks.add(task)
+        task.add_done_callback(self._on_callback_done)
+
+    def _on_callback_done(self, task: "asyncio.Task[Any]") -> None:
+        """Release callback capacity and retrieve unexpected failures."""
+        self._active_tasks.discard(task)
+        if self.sem is not None:
+            self.sem.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__,
+                ),
+            )
+
+    async def _drain_active_tasks(
+        self,
+        *,
+        cancel_immediately: bool,
+    ) -> None:
+        """Wait for callbacks and cancel work beyond the graceful boundary."""
+        tasks = set(self._active_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 suppress(asyncio.CancelledError):
+                await asyncio.shield(waiter)
+        waiter.result()
 
     def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None:
         """

tests/receiver/receiver_lifecycle_support.py

@@ -0,0 +1,179 @@
+import asyncio
+from collections.abc import Awaitable
+from typing import Any
+
+import pytest
+
+from taskiq.abc.broker import AckableMessage
+from taskiq.abc.middleware import TaskiqMiddleware
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.receiver.receiver import QUEUE_DONE, Receiver
+
+
+class ReceiverLifecycleError(RuntimeError):
+    """Marker error for deterministic runner failures."""
+
+
+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 FailingRemoveMiddleware(TaskiqMiddleware):
+    """Fail before the runner transfers capacity to a callback."""
+
+    def on_prefetch_queue_remove(self) -> None:
+        raise ReceiverLifecycleError("prefetch remove failed")
+
+
+class FailingSecondRemoveMiddleware(TaskiqMiddleware):
+    """Fail while one previously admitted callback is still active."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.calls = 0
+
+    def on_prefetch_queue_remove(self) -> None:
+        self.calls += 1
+        if self.calls == 2:
+            raise ReceiverLifecycleError("second prefetch remove failed")
+
+
+class BlockingRemoveMiddleware(TaskiqMiddleware):
+    """Expose cancellation before callback ownership is transferred."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.started = asyncio.Event()
+
+    async def on_prefetch_queue_remove(self) -> None:
+        self.started.set()
+        await asyncio.Event().wait()
+
+
+class ReceiverQueue(asyncio.Queue[bytes | AckableMessage]):
+    """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) -> bytes | AckableMessage:
+        self.get_started.set()
+        message = await super().get()
+        if message is QUEUE_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) | set(self._active_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(b"payload")
+    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
+    for _ in range(slots):
+        await asyncio.wait_for(receiver.sem.acquire(), timeout=0.1)
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(receiver.sem.acquire(), timeout=0.01)
+    for _ in range(slots):
+        receiver.sem.release()