From 4321e09d24ed6008c2ed8fbc42107030f670281d Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Wed, 15 Jul 2026 18:29:42 +0200 Subject: [PATCH] fix(receiver): harden prefetch and listener lifecycle Track delivery and prefetch capacity ownership across shutdown paths. Close pending listener reads, preserve late deliveries, and restore fail-fast transport error propagation. Keep middleware accounting balanced, report discarded deliveries, and retain deferred lifecycle diagnostics when task group failures overlap. --- docs/guide/cli.md | 2 +- taskiq/api/receiver.py | 3 +- taskiq/cli/worker/args.py | 5 +- taskiq/receiver/receiver.py | 440 +++++++++--- tests/receiver/receiver_listener_support.py | 138 ++++ tests/receiver/test_receiver_listener.py | 720 ++++++++++++++++++++ 6 files changed, 1224 insertions(+), 84 deletions(-) create mode 100644 tests/receiver/receiver_listener_support.py create mode 100644 tests/receiver/test_receiver_listener.py diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 8b8a7ab4..0c8c4d81 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -162,7 +162,7 @@ The number of signals before a hard kill can be configured with the `--hardkill- * `--log-format` is used to set a log format (default `%(asctime)s][%(name)s][%(levelname)-7s][%(processName)s] %(message)s`). * `--max-async-tasks` - maximum number of simultaneously running async tasks. * `--max-async-tasks-jitter` – Randomly varies the max async task limit between --max-async-tasks and a jittered value, helping prevent simultaneous worker restarts. -* `--max-prefetch` - number of tasks to be prefetched before execution. (Useful for systems with high message rates, but brokers should support acknowledgements). +* `--max-prefetch` - maximum number of deliveries allowed to wait for execution beyond the available async execution capacity. With a finite `--max-async-tasks` limit, a worker admits at most its effective async execution limit plus `max_prefetch` running or waiting deliveries. The default `0` disables additional buffering, and the value must be non-negative. This option is useful for systems with high message rates, but brokers should support acknowledgements. * `--max-threadpool-threads` - number of threads for sync function execution. * `--no-propagate-errors` - if this parameter is enabled, exceptions won't be thrown in generator dependencies. * `--receiver` - python path to custom receiver class. diff --git a/taskiq/api/receiver.py b/taskiq/api/receiver.py index 03b84ca3..3c940b2c 100644 --- a/taskiq/api/receiver.py +++ b/taskiq/api/receiver.py @@ -45,7 +45,8 @@ async def run_receiver_task( :param validate_params: whether to validate params or not. :param max_async_tasks: maximum number of simultaneous async tasks. :param max_async_tasks_jitter: random jitter to add to max_async_tasks. - :param max_prefetch: maximum number of tasks to prefetch. + :param max_prefetch: maximum number of deliveries allowed to wait beyond + the available async execution capacity. Must be non-negative. :param propagate_exceptions: whether to propagate exceptions in generators or not. :param run_startup: whether to run startup function or not. :param ack_time: acknowledge type to use. diff --git a/taskiq/cli/worker/args.py b/taskiq/cli/worker/args.py index df3eb957..e2dca0c4 100644 --- a/taskiq/cli/worker/args.py +++ b/taskiq/cli/worker/args.py @@ -224,7 +224,10 @@ def from_cli( type=int, dest="max_prefetch", default=0, - help="Maximum prefetched tasks per worker process. ", + help=( + "Maximum deliveries waiting beyond async execution capacity. " + "Must be non-negative. " + ), ) parser.add_argument( "--no-configure-logging", diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index 871d13d6..afc458d1 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -4,8 +4,10 @@ import inspect import random import sys -from collections.abc import Callable +from collections.abc import AsyncGenerator, Callable from concurrent.futures import Executor, ProcessPoolExecutor +from dataclasses import dataclass +from enum import Enum, auto from logging import getLogger from time import time from typing import Any, get_type_hints @@ -26,7 +28,37 @@ logger = getLogger(__name__) PY_VERSION = sys.version_info -QUEUE_DONE = b"-1" + + +class _QueueSignal(Enum): + """Control signals exchanged by the Receiver queue.""" + + DONE = auto() + + +@dataclass(frozen=True, slots=True) +class _PrefetchedMessage: + """A delivery and its admission-capacity ownership.""" + + data: bytes | AckableMessage + owns_delivery_slot: bool + + +@dataclass(slots=True) +class _PrefetchState: + """State owned by one broker listener.""" + + iterator: AsyncGenerator[bytes | AckableMessage, None] + current_message: asyncio.Task[bytes | AckableMessage] | None = None + owns_delivery_slot: bool = False + + +@dataclass(frozen=True, slots=True) +class _StartedCallback: + """A callback task and the delivery capacity transferred to it.""" + + task: asyncio.Task[Any] + owns_delivery_slot: bool def _execute_sync_task_in_executor( @@ -78,9 +110,14 @@ 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._listen_error: BaseException | None = None + if max_prefetch < 0: + raise ValueError("max_prefetch cannot be negative.") for task in self.broker.get_all_tasks().values(): self._prepare_task(task.task_name, task.original_func) + self.sem: asyncio.Semaphore | None = None + delivery_capacity = max_prefetch + 1 if max_async_tasks is not None and max_async_tasks > 0: # Apply jitter to prevent all workers from hitting the limit simultaneously actual_limit = max_async_tasks @@ -91,12 +128,13 @@ def __init__( max_async_tasks_jitter, ) self.sem = asyncio.Semaphore(actual_limit) + delivery_capacity = actual_limit + max_prefetch else: logger.warning( "Setting unlimited number of async tasks " "can result in undefined behavior", ) - self.sem_prefetch = asyncio.Semaphore(max_prefetch) + self.sem_prefetch = asyncio.Semaphore(delivery_capacity) self.is_process_pool = isinstance(executor, ProcessPoolExecutor) async def callback( # noqa: C901, PLR0912 @@ -386,19 +424,37 @@ async def listen(self, finish_event: asyncio.Event) -> None: # pragma: no cover """ if self.run_startup: await self.broker.startup() + self._listen_error = None logger.info("Listening started.") - queue: asyncio.Queue[bytes | AckableMessage] = asyncio.Queue() + queue: asyncio.Queue[_PrefetchedMessage | _QueueSignal] = asyncio.Queue() + + try: + try: + async with anyio.create_task_group() as gr: + gr.start_soon(self.prefetcher, queue, finish_event) + gr.start_soon(self.runner, queue) + finally: + with anyio.CancelScope(shield=True): + await self._discard_queued_messages(queue) + except BaseException: + if self._listen_error is not None: + error = self._listen_error + logger.error( + "A Receiver listener lifecycle error was recorded before " + "the task group failed.", + exc_info=(type(error), error, error.__traceback__), + ) + raise - async with anyio.create_task_group() as gr: - gr.start_soon(self.prefetcher, queue, finish_event) - gr.start_soon(self.runner, queue) + if self._listen_error is not None: + raise self._listen_error if self.on_exit is not None: self.on_exit(self) async def prefetcher( self, - queue: "asyncio.Queue[bytes | AckableMessage]", + queue: "asyncio.Queue[_PrefetchedMessage | _QueueSignal]", finish_event: asyncio.Event, ) -> None: """ @@ -407,54 +463,246 @@ async def prefetcher( :param queue: queue for prefetched data. :param finish_event: event to indicate that we need to stop prefetching. """ - fetched_tasks: int = 0 - iterator = self.broker.listen() - current_message: asyncio.Task[bytes | AckableMessage] = asyncio.create_task( - iterator.__anext__(), # type: ignore - ) + try: + state = _PrefetchState(iterator=self.broker.listen()) + except BaseException as exc: + self._record_listen_error(exc) + queue.put_nowait(_QueueSignal.DONE) + return - while True: - if finish_event.is_set(): - break - try: - await self.sem_prefetch.acquire() - if ( - self.max_tasks_to_execute - and fetched_tasks >= self.max_tasks_to_execute - ): - logger.info("Max number of tasks executed.") + fetched_tasks = 0 + finish_waiter = asyncio.create_task(finish_event.wait()) + try: + while not self._should_stop_prefetch(finish_event, fetched_tasks): + try: + message = await self._get_prefetched_message( + state, + finish_event, + finish_waiter, + ) + except StopAsyncIteration: break - # Here we wait for the message to be fetched, - # but we make it with timeout so it can be interrupted - done, _ = await asyncio.wait({current_message}, timeout=0.3) - # If the message is not fetched, we release the semaphore - # and continue the loop. So it will check if finished event was set. - if not done: - self.sem_prefetch.release() + + if message is None: continue - # We're done, so now we need to check - # whether task has returned an error. - message = current_message.result() - current_message = asyncio.create_task(iterator.__anext__()) # type: ignore + fetched_tasks += 1 - await queue.put(message) - # Custom hooks for OTel and any future instrumentations - for middleware in reversed(self.broker.middlewares): - if hasattr(middleware, "on_prefetch_queue_add"): - await maybe_awaitable( - middleware.on_prefetch_queue_add(), # type: ignore - ) + queue.put_nowait( + _PrefetchedMessage( + data=message, + owns_delivery_slot=True, + ), + ) + state.owns_delivery_slot = False + try: + await self._notify_prefetch_hook("on_prefetch_queue_add") + except asyncio.CancelledError: + raise + except BaseException as exc: + self._record_listen_error(exc) + break + finally: + logger.info("Stopping prefetching messages...") + with anyio.CancelScope(shield=True): + try: + await self._enqueue_late_prefetched_message(queue, state) + finally: + queue.put_nowait(_QueueSignal.DONE) + finish_waiter.cancel() + await asyncio.gather(finish_waiter, return_exceptions=True) + + def _should_stop_prefetch( + self, + finish_event: asyncio.Event, + fetched_tasks: int, + ) -> bool: + """Return whether this Receiver should stop requesting deliveries.""" + if finish_event.is_set(): + return True + if self.max_tasks_to_execute and fetched_tasks >= self.max_tasks_to_execute: + logger.info("Max number of tasks executed.") + return True + return False + + async def _get_prefetched_message( + self, + state: _PrefetchState, + finish_event: asyncio.Event, + finish_waiter: asyncio.Task[bool], + ) -> bytes | AckableMessage | None: + """Acquire capacity and wait for one delivery or the finish signal.""" + if not await self._acquire_delivery_slot( + state, + finish_event, + finish_waiter, + ): + return None + + try: + state.current_message = asyncio.create_task(anext(state.iterator)) + current_message = state.current_message + done, _ = await asyncio.wait( + {current_message, finish_waiter}, + return_when=asyncio.FIRST_COMPLETED, + ) + if current_message in done: + state.current_message = None + return current_message.result() + + return None + except BaseException: + self._release_delivery_slot(state) + raise + + async def _acquire_delivery_slot( + self, + state: _PrefetchState, + finish_event: asyncio.Event, + finish_waiter: asyncio.Task[bool], + ) -> bool: + """Acquire one delivery-admission slot unless shutdown wins the race.""" + acquire_task = asyncio.create_task(self.sem_prefetch.acquire()) + try: + done, _ = await asyncio.wait( + {acquire_task, finish_waiter}, + return_when=asyncio.FIRST_COMPLETED, + ) + except BaseException: + await self._settle_delivery_acquire(acquire_task) + raise + + if finish_waiter in done or finish_event.is_set(): + await self._settle_delivery_acquire(acquire_task) + return False + + acquire_task.result() + state.owns_delivery_slot = True + return True + + async def _settle_delivery_acquire(self, acquire_task: asyncio.Task[bool]) -> None: + """Cancel a capacity waiter and return a concurrently acquired slot.""" + acquire_task.cancel() + with anyio.CancelScope(shield=True): + result = await asyncio.gather(acquire_task, return_exceptions=True) + if result and result[0] is True: + self.sem_prefetch.release() + + def _release_delivery_slot(self, state: _PrefetchState) -> None: + """Release the slot currently owned by the prefetch state.""" + if not state.owns_delivery_slot: + return + state.owns_delivery_slot = False + self.sem_prefetch.release() + + async def _enqueue_late_prefetched_message( + self, + queue: "asyncio.Queue[_PrefetchedMessage | _QueueSignal]", + state: _PrefetchState, + ) -> None: + """Close listener state and retain a delivery completed during stop.""" + late_delivery = await self._close_prefetch_state(state) + if late_delivery is None: + return + + queue.put_nowait(late_delivery) + try: + await self._notify_prefetch_hook("on_prefetch_queue_add") + except BaseException as exc: + self._record_listen_error(exc) + + async def _close_prefetch_state( + self, + state: _PrefetchState, + ) -> _PrefetchedMessage | None: + """Close the pending read and iterator, retaining their first error.""" + late_message: bytes | AckableMessage | None = None + current_message = state.current_message + state.current_message = None + if current_message is not None: + current_message.cancel() + try: + late_message = await current_message except (asyncio.CancelledError, StopAsyncIteration): + pass + except BaseException as exc: + self._record_listen_error(exc) + + try: + await state.iterator.aclose() + except (asyncio.CancelledError, StopAsyncIteration): + pass + except BaseException as exc: + self._record_listen_error(exc) + + if late_message is None: + self._release_delivery_slot(state) + return None + + late_delivery = _PrefetchedMessage( + data=late_message, + owns_delivery_slot=state.owns_delivery_slot, + ) + state.owns_delivery_slot = False + return late_delivery + + async def _notify_prefetch_hook(self, hook_name: str) -> None: + """Run all prefetch hooks and preserve the first failure.""" + first_error: BaseException | None = None + for middleware in reversed(self.broker.middlewares): + hook = getattr(middleware, hook_name, None) + if hook is not None: + try: + await maybe_awaitable(hook()) + except BaseException as exc: + if first_error is None: + first_error = exc + else: + logger.error( + "Additional error while running prefetch hook %s.", + hook_name, + exc_info=(type(exc), exc, exc.__traceback__), + ) + + if first_error is not None: + raise first_error + + async def _discard_queued_messages( + self, + queue: "asyncio.Queue[_PrefetchedMessage | _QueueSignal]", + ) -> None: + """Release capacity and instrumentation for abandoned queue entries.""" + discarded_messages = 0 + while True: + try: + queued_message = queue.get_nowait() + except asyncio.QueueEmpty: break - # We don't want to fetch new messages if we are shutting down. - logger.info("Stopping prefetching messages...") - current_message.cancel() - await queue.put(QUEUE_DONE) - self.sem_prefetch.release() + if isinstance(queued_message, _PrefetchedMessage): + discarded_messages += 1 + if queued_message.owns_delivery_slot: + self.sem_prefetch.release() + + if discarded_messages: + logger.warning( + "Discarding %d prefetched deliveries during Receiver cleanup.", + discarded_messages, + ) + + # Restore all capacity before cleanup hooks introduce suspension points. + first_error: BaseException | None = None + for _ in range(discarded_messages): + try: + await self._notify_prefetch_hook("on_prefetch_queue_remove") + except BaseException as exc: + if first_error is None: + first_error = exc + + if first_error is not None: + self._record_listen_error(first_error) async def runner( self, - queue: "asyncio.Queue[bytes | AckableMessage]", + queue: "asyncio.Queue[_PrefetchedMessage | _QueueSignal]", ) -> None: """ Run tasks. @@ -463,29 +711,10 @@ async def runner( """ 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: - # Waits for semaphore to be released. - if self.sem is not None: - await self.sem.acquire() - - self.sem_prefetch.release() - message = await queue.get() - if message is QUEUE_DONE: + 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( @@ -494,18 +723,8 @@ def task_cb(task: "asyncio.Task[Any]") -> None: await asyncio.wait(tasks, timeout=self.wait_tasks_timeout) logger.info("No more tasks to wait for. Shutting down.") 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 - ) - - task = asyncio.create_task( - self.callback(message=message, raise_err=False), - ) - tasks.add(task) + started_callback = await self._start_callback(queued_message) + tasks.add(started_callback.task) # We want the task to remove itself from the set when it's done. # @@ -513,12 +732,71 @@ def task_cb(task: "asyncio.Task[Any]") -> None: # 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) + 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: break logger.info("The runner is stopped.") + async def _start_callback( + self, + message: _PrefetchedMessage, + ) -> _StartedCallback: + """Transfer execution and delivery capacity to a callback task.""" + owns_delivery_slot = message.owns_delivery_slot + owns_execution_slot = False + try: + await self._notify_prefetch_hook("on_prefetch_queue_remove") + if self.sem is not None: + await self.sem.acquire() + owns_execution_slot = True + + if self.sem is None and owns_delivery_slot: + self.sem_prefetch.release() + owns_delivery_slot = False + return _StartedCallback( + task=asyncio.create_task( + self.callback(message=message.data, raise_err=False), + ), + owns_delivery_slot=owns_delivery_slot, + ) + except BaseException: + if owns_delivery_slot: + self.sem_prefetch.release() + if owns_execution_slot and self.sem is not None: + self.sem.release() + raise + + def _on_callback_done( + self, + task: asyncio.Task[Any], + *, + active_tasks: set[asyncio.Task[Any]], + owns_delivery_slot: bool, + ) -> None: + """Release capacity transferred to a completed callback task.""" + active_tasks.discard(task) + if self.sem is not None: + self.sem.release() + if owns_delivery_slot: + self.sem_prefetch.release() + + def _record_listen_error(self, error: BaseException) -> None: + """Preserve the first listener error and report cleanup failures.""" + if self._listen_error is None: + self._listen_error = error + return + logger.error( + "Additional Receiver listener lifecycle error.", + exc_info=(type(error), error, error.__traceback__), + ) + def _prepare_task(self, name: str, handler: Callable[..., Any]) -> None: """ Prepare task for execution. diff --git a/tests/receiver/receiver_listener_support.py b/tests/receiver/receiver_listener_support.py new file mode 100644 index 00000000..6d6e5404 --- /dev/null +++ b/tests/receiver/receiver_listener_support.py @@ -0,0 +1,138 @@ +import asyncio +from collections.abc import AsyncGenerator, Callable, Sequence +from typing import Literal, cast + +from taskiq.abc.broker import AckableMessage, AsyncBroker +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.message import BrokerMessage, TaskiqMessage + + +class ReceiverLifecycleError(RuntimeError): + """Marker error for listener and middleware lifecycle failures.""" + + +class PrefetchCounterMiddleware(TaskiqMiddleware): + """Track the observable number of messages in the prefetch queue.""" + + def __init__(self) -> None: + super().__init__() + self.queued_messages = 0 + + def on_prefetch_queue_add(self) -> None: + """Record one queued delivery.""" + self.queued_messages += 1 + + def on_prefetch_queue_remove(self) -> None: + """Record one removed or discarded delivery.""" + self.queued_messages -= 1 + + +class ObservedSemaphore(asyncio.Semaphore): + """Semaphore that exposes when a Receiver starts waiting for capacity.""" + + def __init__(self, value: int) -> None: + super().__init__(value) + self.acquire_started = asyncio.Event() + self.acquire_attempts: asyncio.Queue[int] = asyncio.Queue() + self._acquire_count = 0 + + async def acquire(self) -> Literal[True]: + """Record the wait before delegating to asyncio.Semaphore.""" + self._acquire_count += 1 + self.acquire_attempts.put_nowait(self._acquire_count) + self.acquire_started.set() + return await super().acquire() + + +class ControlledBroker(AsyncBroker): + """Queue-backed broker with deterministic read and close barriers.""" + + def __init__(self) -> None: + super().__init__() + self.incoming: asyncio.Queue[bytes | BaseException] = asyncio.Queue() + self.read_started: asyncio.Queue[int] = asyncio.Queue() + self.closed = asyncio.Event() + self.listen_calls = 0 + self._read_count = 0 + + async def kick(self, message: BrokerMessage) -> None: + """Put one encoded message into the controlled listener queue.""" + await self.incoming.put(message.message) + + def listen(self) -> AsyncGenerator[bytes | AckableMessage, None]: + """Create a new controlled listener.""" + self.listen_calls += 1 + return self._listen() + + async def _listen(self) -> AsyncGenerator[bytes | AckableMessage, None]: + try: + while True: + self._read_count += 1 + await self.read_started.put(self._read_count) + item = await self.incoming.get() + if isinstance(item, BaseException): + raise item + yield item + finally: + self.closed.set() + + +class ListenerBroker(AsyncBroker): + """Broker backed by a test-specific async-generator factory.""" + + def __init__( + self, + listener_factory: Callable[ + [], + AsyncGenerator[bytes | AckableMessage, None], + ], + ) -> None: + super().__init__() + self.listener_factory = listener_factory + + async def kick(self, message: BrokerMessage) -> None: + """Ignore sends in listener-only tests.""" + + def listen(self) -> AsyncGenerator[bytes | AckableMessage, None]: + """Return a fresh listener from the configured factory.""" + return self.listener_factory() + + +def encoded_message(broker: AsyncBroker, task_name: str) -> bytes: + """Build one valid transport payload for a registered task.""" + return broker.formatter.dumps( + TaskiqMessage( + task_id="receiver-listener-task", + task_name=task_name, + labels={}, + args=[], + kwargs={}, + ), + ).message + + +def contains_exception(error: BaseException, expected: BaseException) -> bool: + """Return whether an error or a cross-version exception group contains a cause.""" + if error is expected: + return True + nested = cast( + Sequence[BaseException], + getattr(error, "exceptions", ()), + ) + return any(contains_exception(item, expected) for item in nested) + + +async def assert_semaphore_capacity( + semaphore: asyncio.Semaphore, + expected_capacity: int, +) -> None: + """Assert exact available capacity using only semaphore operations.""" + acquired = 0 + try: + for _ in range(expected_capacity): + await asyncio.wait_for(semaphore.acquire(), timeout=1) + acquired += 1 + assert semaphore.locked() + finally: + for _ in range(acquired): + semaphore.release() diff --git a/tests/receiver/test_receiver_listener.py b/tests/receiver/test_receiver_listener.py new file mode 100644 index 00000000..dc498058 --- /dev/null +++ b/tests/receiver/test_receiver_listener.py @@ -0,0 +1,720 @@ +import asyncio +import logging +import unittest.mock +from collections.abc import AsyncGenerator + +import pytest + +from taskiq.abc.broker import AckableMessage +from taskiq.abc.middleware import TaskiqMiddleware +from taskiq.receiver import Receiver +from tests.receiver.receiver_listener_support import ( + ControlledBroker, + ListenerBroker, + ObservedSemaphore, + PrefetchCounterMiddleware, + ReceiverLifecycleError, + assert_semaphore_capacity, + contains_exception, + encoded_message, +) + + +@pytest.mark.parametrize("explicit_zero", [False, True]) +async def test_default_and_zero_prefetch_make_progress(explicit_zero: bool) -> None: + broker = ControlledBroker() + executed = asyncio.Event() + + @broker.task(task_name=f"receiver.prefetch.progress.{explicit_zero}") + async def task() -> None: + executed.set() + + if explicit_zero: + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + ) + else: + receiver = Receiver( + broker, + max_async_tasks=1, + run_startup=False, + ) + finish_event = asyncio.Event() + await task.kiq() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + + try: + await asyncio.wait_for(executed.wait(), timeout=1) + finally: + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + + assert broker.closed.is_set() + + +async def test_zero_prefetch_does_not_read_ahead_of_saturated_execution() -> None: + broker = ControlledBroker() + callback_started = asyncio.Event() + callback_finished = asyncio.Event() + release_callback = asyncio.Event() + + @broker.task(task_name="receiver.prefetch.zero-buffer") + async def task() -> None: + callback_started.set() + try: + await release_callback.wait() + finally: + callback_finished.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + wait_tasks_timeout=0, + ) + delivery_capacity = ObservedSemaphore(1) + receiver.sem_prefetch = delivery_capacity + for _ in range(3): + await task.kiq() + + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + assert await delivery_capacity.acquire_attempts.get() == 1 + assert await broker.read_started.get() == 1 + await asyncio.wait_for(callback_started.wait(), timeout=1) + assert await delivery_capacity.acquire_attempts.get() == 2 + + assert broker.read_started.empty() + assert broker.incoming.qsize() == 2 + assert delivery_capacity.locked() + + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + release_callback.set() + await asyncio.wait_for(callback_finished.wait(), timeout=1) + + await assert_semaphore_capacity(delivery_capacity, 1) + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 1) + + +async def test_zero_prefetch_preserves_bounded_execution_concurrency() -> None: + broker = ControlledBroker() + two_callbacks_started = asyncio.Event() + callbacks_finished = asyncio.Event() + release_callbacks = asyncio.Event() + started = 0 + finished = 0 + + @broker.task(task_name="receiver.prefetch.bounded-concurrency") + async def task() -> None: + nonlocal finished, started + started += 1 + if started == 2: + two_callbacks_started.set() + try: + await release_callbacks.wait() + finally: + finished += 1 + if finished == 2: + callbacks_finished.set() + + receiver = Receiver( + broker, + max_async_tasks=2, + max_prefetch=0, + run_startup=False, + wait_tasks_timeout=0, + ) + delivery_capacity = ObservedSemaphore(2) + receiver.sem_prefetch = delivery_capacity + for _ in range(3): + await task.kiq() + + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + assert [await delivery_capacity.acquire_attempts.get() for _ in range(3)] == [ + 1, + 2, + 3, + ] + assert [await broker.read_started.get() for _ in range(2)] == [1, 2] + await asyncio.wait_for(two_callbacks_started.wait(), timeout=1) + + assert broker.read_started.empty() + assert broker.incoming.qsize() == 1 + assert delivery_capacity.locked() + + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + release_callbacks.set() + await asyncio.wait_for(callbacks_finished.wait(), timeout=1) + + await assert_semaphore_capacity(delivery_capacity, 2) + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 2) + + +async def test_prefetch_is_additional_to_bounded_execution_capacity() -> None: + broker = ControlledBroker() + two_callbacks_started = asyncio.Event() + three_callbacks_finished = asyncio.Event() + release_callbacks = asyncio.Event() + started = 0 + finished = 0 + + @broker.task(task_name="receiver.prefetch.additional-capacity") + async def task() -> None: + nonlocal finished, started + started += 1 + if started == 2: + two_callbacks_started.set() + await release_callbacks.wait() + finished += 1 + if finished == 3: + three_callbacks_finished.set() + + receiver = Receiver( + broker, + max_async_tasks=2, + max_prefetch=1, + run_startup=False, + ) + delivery_capacity = ObservedSemaphore(3) + receiver.sem_prefetch = delivery_capacity + for _ in range(4): + await task.kiq() + + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + assert [await delivery_capacity.acquire_attempts.get() for _ in range(4)] == [ + 1, + 2, + 3, + 4, + ] + assert [await broker.read_started.get() for _ in range(3)] == [1, 2, 3] + await asyncio.wait_for(two_callbacks_started.wait(), timeout=1) + + assert broker.read_started.empty() + assert broker.incoming.qsize() == 1 + assert delivery_capacity.locked() + + finish_event.set() + release_callbacks.set() + await asyncio.wait_for(listen_task, timeout=1) + await asyncio.wait_for(three_callbacks_finished.wait(), timeout=1) + + assert started == 3 + await assert_semaphore_capacity(delivery_capacity, 3) + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 2) + + +async def test_unlimited_execution_keeps_zero_prefetch_handoff_progress() -> None: + broker = ControlledBroker() + two_callbacks_started = asyncio.Event() + callbacks_finished = asyncio.Event() + release_callbacks = asyncio.Event() + finished = 0 + started = 0 + + @broker.task(task_name="receiver.prefetch.unlimited-concurrency") + async def task() -> None: + nonlocal finished, started + started += 1 + if started == 2: + two_callbacks_started.set() + try: + await release_callbacks.wait() + finally: + finished += 1 + if finished == 2: + callbacks_finished.set() + + receiver = Receiver( + broker, + max_prefetch=0, + run_startup=False, + wait_tasks_timeout=0, + ) + await task.kiq() + await task.kiq() + + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + await asyncio.wait_for(two_callbacks_started.wait(), timeout=1) + + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + release_callbacks.set() + await asyncio.wait_for(callbacks_finished.wait(), timeout=1) + + +def test_delivery_capacity_uses_jittered_execution_limit() -> None: + with unittest.mock.patch( + "taskiq.receiver.receiver.random.randint", + return_value=3, + ): + receiver = Receiver( + ControlledBroker(), + max_async_tasks=5, + max_async_tasks_jitter=4, + max_prefetch=2, + run_startup=False, + ) + + assert receiver.sem is not None + assert receiver.sem._value == 8 + assert receiver.sem_prefetch._value == 10 + + +def test_negative_prefetch_is_rejected_before_listener_startup() -> None: + broker = ControlledBroker() + + with pytest.raises(ValueError, match="max_prefetch cannot be negative"): + Receiver(broker, max_prefetch=-1, run_startup=False) + + assert broker.listen_calls == 0 + + +async def test_finish_wakes_prefetcher_blocked_on_capacity() -> None: + broker = ControlledBroker() + receiver = Receiver(broker, max_prefetch=0, run_startup=False) + observed_semaphore = ObservedSemaphore(0) + receiver.sem_prefetch = observed_semaphore + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + + await observed_semaphore.acquire_started.wait() + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + + assert observed_semaphore.locked() + + +async def test_pending_read_is_cancelled_and_iterator_closed() -> None: + broker = ControlledBroker() + receiver = Receiver(broker, max_prefetch=1, run_startup=False) + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + + assert await broker.read_started.get() == 1 + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + + assert broker.closed.is_set() + await assert_semaphore_capacity(receiver.sem_prefetch, 2) + + +async def test_receiver_cancellation_closes_pending_read() -> None: + broker = ControlledBroker() + receiver = Receiver(broker, max_prefetch=0, run_startup=False) + listen_task = asyncio.create_task(receiver.listen(asyncio.Event())) + + assert await broker.read_started.get() == 1 + listen_task.cancel() + with pytest.raises(asyncio.CancelledError): + await listen_task + + assert broker.closed.is_set() + await assert_semaphore_capacity(receiver.sem_prefetch, 1) + + +async def test_ready_delivery_wins_concurrent_finish_signal() -> None: + broker = ControlledBroker() + executed = asyncio.Event() + + @broker.task(task_name="receiver.prefetch.ready-at-finish") + async def task() -> None: + executed.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + ) + finish_event = asyncio.Event() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + assert await broker.read_started.get() == 1 + + broker.incoming.put_nowait(encoded_message(broker, task.task_name)) + finish_event.set() + + await asyncio.wait_for(executed.wait(), timeout=1) + await asyncio.wait_for(listen_task, timeout=1) + + +async def test_late_delivery_is_processed_and_hook_failure_propagates() -> None: + read_started = asyncio.Event() + finish_event = asyncio.Event() + executed = asyncio.Event() + hook_error = ReceiverLifecycleError("late prefetch hook failed") + prefetch_counter = PrefetchCounterMiddleware() + + async def cancellation_delivery() -> AsyncGenerator[bytes | AckableMessage, None]: + read_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + yield encoded_message(broker, task.task_name) + + class FailingAddMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + raise hook_error + + broker = ListenerBroker(cancellation_delivery).with_middlewares( + prefetch_counter, + FailingAddMiddleware(), + ) + + @broker.task(task_name="receiver.prefetch.cancellation-delivery") + async def task() -> None: + executed.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + ) + listen_task = asyncio.create_task(receiver.listen(finish_event)) + await read_started.wait() + + finish_event.set() + with pytest.raises(ReceiverLifecycleError) as exc_info: + await asyncio.wait_for(listen_task, timeout=1) + + assert exc_info.value is hook_error + assert executed.is_set() + assert prefetch_counter.queued_messages == 0 + + +async def test_pending_read_failure_during_shutdown_is_propagated() -> None: + read_started = asyncio.Event() + finish_event = asyncio.Event() + read_error = ReceiverLifecycleError("pending transport read failed") + + async def cancellation_failure() -> AsyncGenerator[bytes | AckableMessage, None]: + read_started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + raise read_error from None + if False: # pragma: no branch + yield b"" + + broker = ListenerBroker(cancellation_failure) + receiver = Receiver(broker, max_prefetch=0, run_startup=False) + listen_task = asyncio.create_task(receiver.listen(finish_event)) + await read_started.wait() + + finish_event.set() + with pytest.raises(ReceiverLifecycleError) as exc_info: + await asyncio.wait_for(listen_task, timeout=1) + + assert exc_info.value is read_error + + +async def test_iterator_close_failure_is_propagated_after_delivery() -> None: + finish_event = asyncio.Event() + executed = asyncio.Event() + close_error = ReceiverLifecycleError("listener close failed") + + async def close_failure() -> AsyncGenerator[bytes | AckableMessage, None]: + try: + yield encoded_message(broker, task.task_name) + await asyncio.Event().wait() + finally: + raise close_error + + class StopAfterPrefetchMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + finish_event.set() + + broker = ListenerBroker(close_failure).with_middlewares( + StopAfterPrefetchMiddleware(), + ) + + @broker.task(task_name="receiver.prefetch.close-failure") + async def task() -> None: + executed.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + ) + + with pytest.raises(ReceiverLifecycleError) as exc_info: + await asyncio.wait_for(receiver.listen(finish_event), timeout=1) + + assert exc_info.value is close_error + assert executed.is_set() + + +async def test_listener_failure_propagates_and_releases_capacity() -> None: + broker = ControlledBroker() + listener_error = ReceiverLifecycleError("listener failed") + receiver = Receiver(broker, max_prefetch=1, run_startup=False) + await broker.incoming.put(listener_error) + + with pytest.raises(BaseException) as exc_info: + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert contains_exception(exc_info.value, listener_error) + assert broker.closed.is_set() + await assert_semaphore_capacity(receiver.sem_prefetch, 2) + + +async def test_listener_failure_interrupts_running_task_wait() -> None: + broker = ControlledBroker() + listener_error = ReceiverLifecycleError("listener failed while task was running") + task_started = asyncio.Event() + task_finished = asyncio.Event() + release_task = asyncio.Event() + + @broker.task(task_name="receiver.prefetch.listener-failure-running-task") + async def task() -> None: + task_started.set() + try: + await release_task.wait() + finally: + task_finished.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=1, + run_startup=False, + ) + await task.kiq() + listen_task = asyncio.create_task(receiver.listen(asyncio.Event())) + await asyncio.wait_for(task_started.wait(), timeout=1) + await broker.incoming.put(listener_error) + + try: + with pytest.raises(BaseException) as exc_info: + await asyncio.wait_for(asyncio.shield(listen_task), timeout=1) + finally: + release_task.set() + await asyncio.wait_for(task_finished.wait(), timeout=1) + await asyncio.gather(listen_task, return_exceptions=True) + + assert contains_exception(exc_info.value, listener_error) + await assert_semaphore_capacity(receiver.sem_prefetch, 2) + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 1) + + +async def test_recorded_listener_error_is_logged_when_runner_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + close_error = ReceiverLifecycleError("listener close failed after runner") + hook_error = ReceiverLifecycleError("prefetch remove hook failed") + + async def close_failure() -> AsyncGenerator[bytes | AckableMessage, None]: + try: + yield encoded_message(broker, task.task_name) + await asyncio.Event().wait() + finally: + raise close_error + + class FailingRemoveMiddleware(TaskiqMiddleware): + def on_prefetch_queue_remove(self) -> None: + raise hook_error + + broker = ListenerBroker(close_failure).with_middlewares( + FailingRemoveMiddleware(), + ) + + @broker.task(task_name="receiver.prefetch.listener-and-runner-failure") + async def task() -> None: + pass + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=1, + run_startup=False, + ) + + with ( + caplog.at_level(logging.ERROR, logger="taskiq.receiver.receiver"), + pytest.raises(BaseException) as exc_info, + ): + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert contains_exception(exc_info.value, hook_error) + assert "A Receiver listener lifecycle error was recorded" in caplog.text + assert str(close_error) in caplog.text + + +async def test_listener_open_failure_propagates() -> None: + listener_error = ReceiverLifecycleError("listener open failed") + + def opening_failure() -> AsyncGenerator[bytes | AckableMessage, None]: + raise listener_error + + broker = ListenerBroker(opening_failure) + receiver = Receiver(broker, run_startup=False) + + with pytest.raises(ReceiverLifecycleError) as exc_info: + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert exc_info.value is listener_error + + +async def test_prefetch_add_hook_failure_preserves_delivery_and_capacity() -> None: + hook_error = ReceiverLifecycleError("prefetch add hook failed") + executed = asyncio.Event() + prefetch_counter = PrefetchCounterMiddleware() + + class FailingAddMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + raise hook_error + + broker = ControlledBroker().with_middlewares( + prefetch_counter, + FailingAddMiddleware(), + ) + + @broker.task(task_name="receiver.prefetch.add-failure") + async def task() -> None: + executed.set() + + receiver = Receiver(broker, max_prefetch=1, run_startup=False) + await task.kiq() + + with pytest.raises(ReceiverLifecycleError) as exc_info: + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert exc_info.value is hook_error + assert executed.is_set() + assert broker.closed.is_set() + assert prefetch_counter.queued_messages == 0 + await assert_semaphore_capacity(receiver.sem_prefetch, 2) + + +async def test_remove_hook_failure_releases_buffered_capacity( + caplog: pytest.LogCaptureFixture, +) -> None: + three_messages_added = asyncio.Event() + hook_error = ReceiverLifecycleError("prefetch remove hook failed") + task_started = asyncio.Event() + added_messages = 0 + prefetch_counter = PrefetchCounterMiddleware() + + class CoordinatedFailingMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + nonlocal added_messages + added_messages += 1 + if added_messages == 3: + three_messages_added.set() + + async def on_prefetch_queue_remove(self) -> None: + await three_messages_added.wait() + raise hook_error + + broker = ControlledBroker().with_middlewares( + prefetch_counter, + CoordinatedFailingMiddleware(), + ) + + @broker.task(task_name="receiver.prefetch.buffered-remove-failure") + async def task() -> None: + task_started.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=2, + run_startup=False, + ) + for _ in range(3): + await task.kiq() + + with ( + caplog.at_level(logging.WARNING, logger="taskiq.receiver.receiver"), + pytest.raises(BaseException) as exc_info, + ): + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert contains_exception(exc_info.value, hook_error) + assert not task_started.is_set() + assert "Discarding 2 prefetched deliveries during Receiver cleanup" in caplog.text + assert prefetch_counter.queued_messages == 0 + await assert_semaphore_capacity(receiver.sem_prefetch, 3) + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 1) + + +async def test_shutdown_sentinel_bypasses_saturated_execution_capacity() -> None: + broker = ControlledBroker() + callback_started = asyncio.Event() + callback_finished = asyncio.Event() + release_callback = asyncio.Event() + + @broker.task(task_name="receiver.prefetch.saturated-shutdown") + async def task() -> None: + callback_started.set() + try: + await release_callback.wait() + finally: + callback_finished.set() + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=0, + run_startup=False, + wait_tasks_timeout=0, + ) + finish_event = asyncio.Event() + await task.kiq() + listen_task = asyncio.create_task(receiver.listen(finish_event)) + await asyncio.wait_for(callback_started.wait(), timeout=1) + + finish_event.set() + try: + await asyncio.wait_for(asyncio.shield(listen_task), timeout=1) + finally: + release_callback.set() + await asyncio.wait_for(callback_finished.wait(), timeout=1) + if not listen_task.done(): + await asyncio.wait_for(listen_task, timeout=1) + + assert receiver.sem is not None + await assert_semaphore_capacity(receiver.sem, 1) + assert broker.closed.is_set() + + +async def test_task_limit_leaves_the_next_delivery_in_the_broker() -> None: + broker = ControlledBroker() + executions = 0 + + @broker.task(task_name="receiver.prefetch.legacy-limit") + async def task() -> None: + nonlocal executions + executions += 1 + + receiver = Receiver( + broker, + max_async_tasks=1, + max_prefetch=1, + max_tasks_to_execute=1, + run_startup=False, + ) + await task.kiq() + await task.kiq() + + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert executions == 1 + assert broker.incoming.qsize() == 1