From 88fa10f01bbc4e7f6ea540ba85ae9cf03af9418e Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Mon, 13 Jul 2026 23:31:12 +0200 Subject: [PATCH] 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 --- docs/guide/cli.md | 6 +- taskiq/cli/worker/args.py | 5 +- taskiq/receiver/receiver.py | 136 ++++--- tests/receiver/receiver_lifecycle_support.py | 179 +++++++++ tests/receiver/test_receiver_lifecycle.py | 374 +++++++++++++++++++ 5 files changed, 653 insertions(+), 47 deletions(-) create mode 100644 tests/receiver/receiver_lifecycle_support.py create mode 100644 tests/receiver/test_receiver_lifecycle.py diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 8b8a7ab4..92779472 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -143,10 +143,10 @@ kill -HUP
### Graceful and force shutdowns -If you send `SIGINT` or `SIGKILL` to the main process by pressing Ctrl+C or using the `kill` command, it will initiate the shutdown process. +If you send `SIGINT` or `SIGTERM` to the main process by pressing Ctrl+C 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 diff --git a/taskiq/cli/worker/args.py b/taskiq/cli/worker/args.py index df3eb957..95716768 100644 --- a/taskiq/cli/worker/args.py +++ b/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", diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 871d13d6..4043d6c1 100644 --- a/taskiq/receiver/receiver.py +++ b/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: """ diff --git a/tests/receiver/receiver_lifecycle_support.py b/tests/receiver/receiver_lifecycle_support.py new file mode 100644 index 00000000..d77eb613 --- /dev/null +++ b/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() diff --git a/tests/receiver/test_receiver_lifecycle.py b/tests/receiver/test_receiver_lifecycle.py new file mode 100644 index 00000000..e2aa4cf8 --- /dev/null +++ b/tests/receiver/test_receiver_lifecycle.py @@ -0,0 +1,374 @@ +import asyncio +import logging + +import anyio +import pytest + +from taskiq.abc.broker import AckableMessage +from taskiq.receiver.receiver import QUEUE_DONE, Receiver +from tests.receiver.receiver_lifecycle_support import ( + BlockingRemoveMiddleware, + ControlledReceiver, + FailingRemoveMiddleware, + FailingSecondRemoveMiddleware, + ReceiverLifecycleError, + ReceiverQueue, + ShieldCallCounter, + assert_exact_capacity, + start_callback, + wait_for_signals, +) +from tests.utils import AsyncQueueBroker + +RUNNER_PROBE_TASK_NAME = "receiver-drain-cancellation-probe" + + +async def test_finite_timeout_cancels_and_drains_callback() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=0.01, + max_async_tasks=None, + ) + queue, runner_task = await start_callback(receiver) + + try: + await queue.put(QUEUE_DONE) + await asyncio.wait_for(queue.shutdown_received.wait(), timeout=1) + await wait_for_signals(receiver.cleanup_started_callbacks) + assert not runner_task.done() + + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.finished_callbacks.qsize() == 1 + assert receiver.callback_task is not None + assert receiver.callback_task.cancelled() + finally: + await receiver.settle(runner_task) + + +async def test_finite_timeout_drains_every_active_callback() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=0.01, + max_async_tasks=None, + ) + queue = ReceiverQueue() + await queue.put(b"first") + await queue.put(b"second") + runner_task = asyncio.create_task(receiver.runner(queue)) + + try: + await wait_for_signals(receiver.started_callbacks, count=2) + await queue.put(QUEUE_DONE) + await asyncio.wait_for(queue.shutdown_received.wait(), timeout=1) + await wait_for_signals(receiver.cleanup_started_callbacks, count=2) + assert not runner_task.done() + + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + assert len(receiver.callback_tasks) == 2 + assert all(task.cancelled() for task in receiver.callback_tasks) + assert receiver.finished_callbacks.qsize() == 2 + finally: + await receiver.settle(runner_task) + + +async def test_outer_cancellation_drains_callback_before_runner_stops() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + _, runner_task = await start_callback(receiver) + + try: + runner_task.cancel() + await wait_for_signals(receiver.cleanup_started_callbacks) + assert not runner_task.done() + + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.finished_callbacks.qsize() == 1 + finally: + await receiver.settle(runner_task) + + +async def test_cancellation_during_graceful_wait_drains_callback() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=None, + max_async_tasks=None, + ) + queue, runner_task = await start_callback(receiver) + + try: + await queue.put(QUEUE_DONE) + await asyncio.wait_for(queue.shutdown_received.wait(), timeout=1) + assert not runner_task.done() + + runner_task.cancel() + await wait_for_signals(receiver.cleanup_started_callbacks) + receiver.release_cleanup.set() + + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.finished_callbacks.qsize() == 1 + finally: + await receiver.settle(runner_task) + + +async def test_repeated_cancellation_does_not_interrupt_callback_cleanup() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + _, runner_task = await start_callback(receiver) + + try: + runner_task.cancel() + await wait_for_signals(receiver.cleanup_started_callbacks) + runner_task.cancel() + receiver.release_cleanup.set() + + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.finished_callbacks.qsize() == 1 + finally: + await receiver.settle(runner_task) + + +async def test_listen_cancellation_preserves_multicheckpoint_task_cleanup() -> None: + broker = AsyncQueueBroker() + callback_started = asyncio.Event() + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + cleanup_finished = asyncio.Event() + callback_task: asyncio.Task[None] | None = None + + @broker.task + async def blocked_task() -> None: + nonlocal callback_task + callback_task = asyncio.current_task() + callback_started.set() + try: + await asyncio.Event().wait() + finally: + cleanup_started.set() + await release_cleanup.wait() + for _ in range(20): + await asyncio.sleep(0) + cleanup_finished.set() + + await blocked_task.kiq() + receiver = Receiver( + broker, + max_async_tasks=1, + run_startup=False, + wait_tasks_timeout=None, + ) + listen_task = asyncio.create_task(receiver.listen(asyncio.Event())) + + try: + await asyncio.wait_for(callback_started.wait(), timeout=1) + listen_task.cancel() + await asyncio.wait_for(cleanup_started.wait(), timeout=1) + listen_task.cancel() + release_cleanup.set() + await asyncio.gather(listen_task, return_exceptions=True) + + assert cleanup_finished.is_set() + finally: + release_cleanup.set() + listen_task.cancel() + if callback_task is not None and not callback_task.done(): + callback_task.cancel() + await asyncio.gather( + *(task for task in (callback_task, listen_task) if task is not None), + return_exceptions=True, + ) + + +async def test_level_cancellation_does_not_busy_spin_during_callback_drain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + queue = ReceiverQueue() + await queue.put(b"payload") + shield_counter = ShieldCallCounter(RUNNER_PROBE_TASK_NAME) + monkeypatch.setattr(asyncio, "shield", shield_counter) + + async with anyio.create_task_group() as task_group: + task_group.start_soon( + receiver.runner, + queue, + name=RUNNER_PROBE_TASK_NAME, + ) + await wait_for_signals(receiver.started_callbacks) + task_group.cancel_scope.cancel() + + with anyio.CancelScope(shield=True): + try: + await wait_for_signals(receiver.cleanup_started_callbacks) + for _ in range(20): + await asyncio.sleep(0) + finally: + receiver.release_cleanup.set() + + assert receiver.finished_callbacks.qsize() == 1 + assert shield_counter.runner_calls <= 2 + + +async def test_level_cancellation_interrupts_unbounded_graceful_wait() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=None, + max_async_tasks=None, + ) + queue = ReceiverQueue() + await queue.put(b"payload") + + async with anyio.create_task_group() as task_group: + task_group.start_soon(receiver.runner, queue) + await wait_for_signals(receiver.started_callbacks) + await queue.put(QUEUE_DONE) + await asyncio.wait_for(queue.shutdown_received.wait(), timeout=1) + task_group.cancel_scope.cancel() + + with anyio.CancelScope(shield=True): + try: + await wait_for_signals(receiver.cleanup_started_callbacks) + finally: + receiver.release_callback.set() + receiver.release_cleanup.set() + + assert receiver.callback_task is not None + assert receiver.callback_task.cancelled() + + +async def test_unbounded_wait_preserves_graceful_callback_completion() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=None, + max_async_tasks=None, + ) + queue, runner_task = await start_callback(receiver) + + try: + await queue.put(QUEUE_DONE) + await asyncio.wait_for(queue.shutdown_received.wait(), timeout=1) + assert not runner_task.done() + + receiver.release_callback.set() + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.finished_callbacks.qsize() == 1 + assert receiver.callback_task is not None + assert not receiver.callback_task.cancelled() + finally: + await receiver.settle(runner_task) + + +async def test_runner_cancellation_before_message_releases_capacity() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + queue = ReceiverQueue() + runner_task = asyncio.create_task(receiver.runner(queue)) + + try: + await asyncio.wait_for(queue.get_started.wait(), timeout=1) + assert receiver.sem is not None + assert receiver.sem.locked() + + runner_task.cancel() + await asyncio.gather(runner_task, return_exceptions=True) + + await assert_exact_capacity(receiver) + finally: + await receiver.settle(runner_task) + + +async def test_prefetch_remove_failure_releases_capacity() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + receiver.broker.with_middlewares(FailingRemoveMiddleware()) + queue: asyncio.Queue[bytes | AckableMessage] = asyncio.Queue() + await queue.put(b"payload") + + with pytest.raises(ReceiverLifecycleError, match="prefetch remove failed"): + await receiver.runner(queue) + + await assert_exact_capacity(receiver) + + +async def test_prefetch_remove_cancellation_releases_capacity() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + middleware = BlockingRemoveMiddleware() + receiver.broker.with_middlewares(middleware) + queue: asyncio.Queue[bytes | AckableMessage] = asyncio.Queue() + await queue.put(b"payload") + runner_task = asyncio.create_task(receiver.runner(queue)) + + try: + await asyncio.wait_for(middleware.started.wait(), timeout=1) + runner_task.cancel() + await asyncio.gather(runner_task, return_exceptions=True) + + assert receiver.started_callbacks.empty() + await assert_exact_capacity(receiver) + finally: + await receiver.settle(runner_task) + + +async def test_runner_failure_drains_active_callback_and_releases_capacity() -> None: + receiver = ControlledReceiver( + wait_tasks_timeout=None, + max_async_tasks=2, + ) + receiver.broker.with_middlewares(FailingSecondRemoveMiddleware()) + queue = ReceiverQueue() + await queue.put(b"first") + runner_task = asyncio.create_task(receiver.runner(queue)) + + try: + await wait_for_signals(receiver.started_callbacks) + await queue.put(b"second") + await wait_for_signals(receiver.cleanup_started_callbacks) + assert not runner_task.done() + + receiver.release_cleanup.set() + with pytest.raises( + ReceiverLifecycleError, + match="second prefetch remove failed", + ): + await asyncio.wait_for(runner_task, timeout=1) + + assert receiver.callback_task is not None + assert receiver.callback_task.cancelled() + assert receiver.finished_callbacks.qsize() == 1 + await assert_exact_capacity(receiver, slots=2) + finally: + await receiver.settle(runner_task) + + +async def test_callback_completion_releases_exactly_one_capacity_slot() -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None) + queue, runner_task = await start_callback(receiver) + + try: + await queue.put(QUEUE_DONE) + receiver.release_callback.set() + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + await assert_exact_capacity(receiver) + finally: + await receiver.settle(runner_task) + + +async def test_callback_failure_is_logged_and_releases_capacity( + caplog: pytest.LogCaptureFixture, +) -> None: + receiver = ControlledReceiver(wait_tasks_timeout=None, fail=True) + queue, runner_task = await start_callback(receiver) + caplog.set_level(logging.ERROR, logger="taskiq.receiver.receiver") + + try: + await queue.put(QUEUE_DONE) + receiver.release_callback.set() + receiver.release_cleanup.set() + await asyncio.wait_for(runner_task, timeout=1) + + await assert_exact_capacity(receiver) + assert "Receiver callback failed outside task execution handling" in caplog.text + finally: + await receiver.settle(runner_task)