From 4321e09d24ed6008c2ed8fbc42107030f670281d Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Wed, 15 Jul 2026 18:29:42 +0200 Subject: [PATCH 1/2] 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 From b6f07dcee773dde68d389fdba11264bc4b823a2b Mon Sep 17 00:00:00 2001 From: "sergei.romanchuk" Date: Fri, 24 Jul 2026 22:15:21 +0200 Subject: [PATCH 2/2] fix(receiver): tighten prefetch ownership checks Reject invalid prefetch values before receiver retries and retain delivery accounting until execution capacity is acquired. Add deterministic coverage for CLI defaults, cancellation races, and listener lifecycle cleanup. Refs #528 --- taskiq/api/receiver.py | 4 + taskiq/cli/worker/args.py | 2 + taskiq/receiver/receiver.py | 32 +++++-- tests/api/test_receiver_task.py | 33 +++++++ tests/cli/worker/test_args.py | 37 ++++++++ tests/receiver/test_receiver_listener.py | 110 ++++++++++++++++++++++- 6 files changed, 208 insertions(+), 10 deletions(-) create mode 100644 tests/cli/worker/test_args.py diff --git a/taskiq/api/receiver.py b/taskiq/api/receiver.py index 3c940b2c..9417398d 100644 --- a/taskiq/api/receiver.py +++ b/taskiq/api/receiver.py @@ -52,7 +52,11 @@ async def run_receiver_task( :param ack_time: acknowledge type to use. :param use_process_pool: whether to use process pool or threadpool. :raises asyncio.CancelledError: if the task was cancelled. + :raises ValueError: if max_prefetch is negative. """ + if max_prefetch < 0: + raise ValueError("max_prefetch cannot be negative.") + finish_event = asyncio.Event() def on_exit(_: Receiver) -> None: diff --git a/taskiq/cli/worker/args.py b/taskiq/cli/worker/args.py index e2dca0c4..8187113c 100644 --- a/taskiq/cli/worker/args.py +++ b/taskiq/cli/worker/args.py @@ -286,6 +286,8 @@ def from_cli( args, namespace=None if defaults is None else Namespace(**defaults), ) + if namespace.max_prefetch < 0: + parser.error("argument --max-prefetch: max_prefetch cannot be negative.") # If there are any patterns specified, remove default. # This is an argparse limitation. if len(namespace.tasks_pattern) > 1: diff --git a/taskiq/receiver/receiver.py b/taskiq/receiver/receiver.py index afc458d1..98994bf7 100644 --- a/taskiq/receiver/receiver.py +++ b/taskiq/receiver/receiver.py @@ -10,7 +10,7 @@ from enum import Enum, auto from logging import getLogger from time import time -from typing import Any, get_type_hints +from typing import Any, Literal, get_type_hints import anyio from taskiq_dependencies import DependencyGraph @@ -645,7 +645,13 @@ async def _close_prefetch_state( state.owns_delivery_slot = False return late_delivery - async def _notify_prefetch_hook(self, hook_name: str) -> None: + async def _notify_prefetch_hook( + self, + hook_name: Literal[ + "on_prefetch_queue_add", + "on_prefetch_queue_remove", + ], + ) -> None: """Run all prefetch hooks and preserve the first failure.""" first_error: BaseException | None = None for middleware in reversed(self.broker.middlewares): @@ -723,7 +729,18 @@ async def runner( await asyncio.wait(tasks, timeout=self.wait_tasks_timeout) logger.info("No more tasks to wait for. Shutting down.") break - started_callback = await self._start_callback(queued_message) + execution_semaphore = self.sem + owns_execution_slot = execution_semaphore is not None + if execution_semaphore is not None: + try: + await execution_semaphore.acquire() + except BaseException: + queue.put_nowait(queued_message) + raise + started_callback = await self._start_callback( + queued_message, + owns_execution_slot=owns_execution_slot, + ) tasks.add(started_callback.task) # We want the task to remove itself from the set when it's done. @@ -747,15 +764,13 @@ async def runner( async def _start_callback( self, message: _PrefetchedMessage, + *, + owns_execution_slot: bool, ) -> _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() @@ -767,6 +782,9 @@ async def _start_callback( owns_delivery_slot=owns_delivery_slot, ) except BaseException: + logger.warning( + "Discarding 1 prefetched delivery during Receiver cleanup.", + ) if owns_delivery_slot: self.sem_prefetch.release() if owns_execution_slot and self.sem is not None: diff --git a/tests/api/test_receiver_task.py b/tests/api/test_receiver_task.py index bae1df8d..e593f7cd 100644 --- a/tests/api/test_receiver_task.py +++ b/tests/api/test_receiver_task.py @@ -1,12 +1,30 @@ import asyncio import contextlib +from typing import Any import pytest from taskiq.api import run_receiver_task +from taskiq.receiver import Receiver from tests.utils import AsyncQueueBroker +class _UnexpectedReceiverRetry(BaseException): + """Signal that invalid configuration reached the Receiver retry loop.""" + + +class _ValidationProbeReceiver(Receiver): + """Fail deterministically if invalid configuration is retried.""" + + construction_attempts = 0 + + def __init__(self, *args: Any, **kwargs: Any) -> None: + type(self).construction_attempts += 1 + if type(self).construction_attempts > 1: + raise _UnexpectedReceiverRetry + super().__init__(*args, **kwargs) + + async def test_successful() -> None: broker = AsyncQueueBroker() kicked = 0 @@ -52,3 +70,18 @@ def test_func() -> None: with pytest.raises(asyncio.TimeoutError): await asyncio.wait_for(broker.wait_tasks(), 0.2) assert kicked == 1 + + +async def test_negative_prefetch_is_rejected_before_receiver_retry() -> None: + broker = AsyncQueueBroker() + _ValidationProbeReceiver.construction_attempts = 0 + + with pytest.raises(ValueError, match="max_prefetch cannot be negative"): + await run_receiver_task( + broker, + receiver_cls=_ValidationProbeReceiver, + max_prefetch=-1, + ) + + assert _ValidationProbeReceiver.construction_attempts == 0 + assert not broker.is_worker_process diff --git a/tests/cli/worker/test_args.py b/tests/cli/worker/test_args.py new file mode 100644 index 00000000..258d7e1c --- /dev/null +++ b/tests/cli/worker/test_args.py @@ -0,0 +1,37 @@ +import pytest + +from taskiq.cli.worker.args import WorkerArgs + + +@pytest.mark.parametrize("max_prefetch", [0, 3]) +def test_max_prefetch_accepts_non_negative_values(max_prefetch: int) -> None: + args = WorkerArgs.from_cli( + ["example:broker", "--max-prefetch", str(max_prefetch)], + ) + + assert args.max_prefetch == max_prefetch + + +def test_max_prefetch_rejects_negative_value( + capsys: pytest.CaptureFixture[str], +) -> None: + with pytest.raises(SystemExit) as exc_info: + WorkerArgs.from_cli( + ["example:broker", "--max-prefetch", "-1"], + ) + + assert exc_info.value.code == 2 + assert "max_prefetch cannot be negative" in capsys.readouterr().err + + +def test_max_prefetch_rejects_negative_default( + capsys: pytest.CaptureFixture[str], +) -> None: + with pytest.raises(SystemExit) as exc_info: + WorkerArgs.from_cli( + ["example:broker"], + defaults={"max_prefetch": -1}, + ) + + assert exc_info.value.code == 2 + assert "max_prefetch cannot be negative" in capsys.readouterr().err diff --git a/tests/receiver/test_receiver_listener.py b/tests/receiver/test_receiver_listener.py index dc498058..f8badd9f 100644 --- a/tests/receiver/test_receiver_listener.py +++ b/tests/receiver/test_receiver_listener.py @@ -215,6 +215,74 @@ async def task() -> None: await assert_semaphore_capacity(receiver.sem, 2) +async def test_runner_retains_prefetch_ownership_until_callback_handoff( + caplog: pytest.LogCaptureFixture, +) -> None: + broker = ControlledBroker() + prefetch_counter = PrefetchCounterMiddleware() + three_messages_added = asyncio.Event() + callback_started = asyncio.Event() + callback_finished = asyncio.Event() + release_callback = asyncio.Event() + added_messages = 0 + + class AddBarrierMiddleware(TaskiqMiddleware): + def on_prefetch_queue_add(self) -> None: + nonlocal added_messages + added_messages += 1 + if added_messages == 3: + three_messages_added.set() + + broker.with_middlewares(prefetch_counter, AddBarrierMiddleware()) + + @broker.task(task_name="receiver.prefetch.runner-owned") + 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=2, + run_startup=False, + ) + execution_capacity = ObservedSemaphore(1) + receiver.sem = execution_capacity + for _ in range(3): + await task.kiq() + + listen_task = asyncio.create_task(receiver.listen(asyncio.Event())) + try: + await asyncio.wait_for(callback_started.wait(), timeout=1) + assert await execution_capacity.acquire_attempts.get() == 1 + assert await execution_capacity.acquire_attempts.get() == 2 + await asyncio.wait_for(three_messages_added.wait(), timeout=1) + + assert prefetch_counter.queued_messages == 2 + + with caplog.at_level(logging.WARNING, logger="taskiq.receiver.receiver"): + listen_task.cancel() + with pytest.raises(asyncio.CancelledError): + await listen_task + + assert ( + "Discarding 2 prefetched deliveries during Receiver cleanup" in caplog.text + ) + assert prefetch_counter.queued_messages == 0 + finally: + if not listen_task.done(): + listen_task.cancel() + await asyncio.gather(listen_task, return_exceptions=True) + release_callback.set() + await asyncio.wait_for(callback_finished.wait(), timeout=1) + + await assert_semaphore_capacity(receiver.sem_prefetch, 3) + await assert_semaphore_capacity(execution_capacity, 1) + + async def test_unlimited_execution_keeps_zero_prefetch_handoff_progress() -> None: broker = ControlledBroker() two_callbacks_started = asyncio.Event() @@ -255,7 +323,7 @@ async def task() -> None: await asyncio.wait_for(callbacks_finished.wait(), timeout=1) -def test_delivery_capacity_uses_jittered_execution_limit() -> None: +async def test_delivery_capacity_uses_jittered_execution_limit() -> None: with unittest.mock.patch( "taskiq.receiver.receiver.random.randint", return_value=3, @@ -269,8 +337,8 @@ def test_delivery_capacity_uses_jittered_execution_limit() -> None: ) assert receiver.sem is not None - assert receiver.sem._value == 8 - assert receiver.sem_prefetch._value == 10 + await assert_semaphore_capacity(receiver.sem, 8) + await assert_semaphore_capacity(receiver.sem_prefetch, 10) def test_negative_prefetch_is_rejected_before_listener_startup() -> None: @@ -297,6 +365,23 @@ async def test_finish_wakes_prefetcher_blocked_on_capacity() -> None: assert observed_semaphore.locked() +async def test_finish_returns_concurrently_acquired_capacity_once() -> 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() + observed_semaphore.release() + finish_event.set() + await asyncio.wait_for(listen_task, timeout=1) + + assert broker.read_started.empty() + await assert_semaphore_capacity(observed_semaphore, 1) + + async def test_pending_read_is_cancelled_and_iterator_closed() -> None: broker = ControlledBroker() receiver = Receiver(broker, max_prefetch=1, run_startup=False) @@ -571,6 +656,25 @@ def opening_failure() -> AsyncGenerator[bytes | AckableMessage, None]: assert exc_info.value is listener_error +async def test_listener_exhaustion_stops_cleanly_and_releases_capacity() -> None: + listener_closed = asyncio.Event() + + async def exhausted_listener() -> AsyncGenerator[bytes | AckableMessage, None]: + try: + if False: # pragma: no branch + yield b"" + finally: + listener_closed.set() + + broker = ListenerBroker(exhausted_listener) + receiver = Receiver(broker, max_prefetch=0, run_startup=False) + + await asyncio.wait_for(receiver.listen(asyncio.Event()), timeout=1) + + assert listener_closed.is_set() + await assert_semaphore_capacity(receiver.sem_prefetch, 1) + + async def test_prefetch_add_hook_failure_preserves_delivery_and_capacity() -> None: hook_error = ReceiverLifecycleError("prefetch add hook failed") executed = asyncio.Event()