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

Author: GefMar

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:

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:

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:

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

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

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()