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.

Author: GefMar

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` - maximum number of deliveries buffered in addition to the listener's single in-flight delivery. The default `0` disables additional buffering while still allowing one delivery to move from the broker listener toward execution. The value must be non-negative. This option is 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.

taskiq/api/receiver.py

@@ -45,8 +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 deliveries buffered in addition to
-        the listener's single in-flight delivery. Must be non-negative.
+    :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.

taskiq/cli/worker/args.py

@@ -225,7 +225,7 @@ def from_cli(
             dest="max_prefetch",
             default=0,
             help=(
-                "Maximum deliveries buffered beyond the in-flight delivery. "
+                "Maximum deliveries waiting beyond async execution capacity. "
                 "Must be non-negative. "
             ),
         )

taskiq/receiver/receiver.py

@@ -38,10 +38,10 @@ class _QueueSignal(Enum):
 
 @dataclass(frozen=True, slots=True)
 class _PrefetchedMessage:
-    """A delivery and its prefetch-capacity ownership."""
+    """A delivery and its admission-capacity ownership."""
 
     data: bytes | AckableMessage
-    owns_prefetch_slot: bool
+    owns_delivery_slot: bool
 
 
 @dataclass(slots=True)
@@ -50,7 +50,15 @@ class _PrefetchState:
 
     iterator: AsyncGenerator[bytes | AckableMessage, None]
     current_message: asyncio.Task[bytes | AckableMessage] | None = None
-    owns_prefetch_slot: bool = False
+    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(
@@ -103,9 +111,13 @@ def __init__(
         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
@@ -116,16 +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",
             )
-        if max_prefetch < 0:
-            raise ValueError("max_prefetch cannot be negative.")
-        # One slot tracks the listener's in-flight delivery. The configured
-        # value bounds the additional deliveries buffered ahead of execution.
-        self.sem_prefetch = asyncio.Semaphore(max_prefetch + 1)
+        self.sem_prefetch = asyncio.Semaphore(delivery_capacity)
         self.is_process_pool = isinstance(executor, ProcessPoolExecutor)
 
     async def callback(  # noqa: C901, PLR0912
@@ -481,10 +490,10 @@ async def prefetcher(
                 queue.put_nowait(
                     _PrefetchedMessage(
                         data=message,
-                        owns_prefetch_slot=True,
+                        owns_delivery_slot=True,
                     ),
                 )
-                state.owns_prefetch_slot = False
+                state.owns_delivery_slot = False
                 try:
                     await self._notify_prefetch_hook("on_prefetch_queue_add")
                 except asyncio.CancelledError:
@@ -522,7 +531,7 @@ async def _get_prefetched_message(
         finish_waiter: asyncio.Task[bool],
     ) -> bytes | AckableMessage | None:
         """Acquire capacity and wait for one delivery or the finish signal."""
-        if not await self._acquire_prefetch_slot(
+        if not await self._acquire_delivery_slot(
             state,
             finish_event,
             finish_waiter,
@@ -540,50 +549,49 @@ async def _get_prefetched_message(
                 state.current_message = None
                 return current_message.result()
 
-            self._release_prefetch_slot(state)
             return None
         except BaseException:
-            self._release_prefetch_slot(state)
+            self._release_delivery_slot(state)
             raise
 
-    async def _acquire_prefetch_slot(
+    async def _acquire_delivery_slot(
         self,
         state: _PrefetchState,
         finish_event: asyncio.Event,
         finish_waiter: asyncio.Task[bool],
     ) -> bool:
-        """Acquire one prefetch slot unless shutdown wins the race."""
+        """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_prefetch_acquire(acquire_task)
+            await self._settle_delivery_acquire(acquire_task)
             raise
 
         if finish_waiter in done or finish_event.is_set():
-            await self._settle_prefetch_acquire(acquire_task)
+            await self._settle_delivery_acquire(acquire_task)
             return False
 
         acquire_task.result()
-        state.owns_prefetch_slot = True
+        state.owns_delivery_slot = True
         return True
 
-    async def _settle_prefetch_acquire(self, acquire_task: asyncio.Task[bool]) -> None:
+    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_prefetch_slot(self, state: _PrefetchState) -> None:
+    def _release_delivery_slot(self, state: _PrefetchState) -> None:
         """Release the slot currently owned by the prefetch state."""
-        if not state.owns_prefetch_slot:
+        if not state.owns_delivery_slot:
             return
-        state.owns_prefetch_slot = False
+        state.owns_delivery_slot = False
         self.sem_prefetch.release()
 
     async def _enqueue_late_prefetched_message(
@@ -592,16 +600,11 @@ async def _enqueue_late_prefetched_message(
         state: _PrefetchState,
     ) -> None:
         """Close listener state and retain a delivery completed during stop."""
-        late_message = await self._close_prefetch_state(state)
-        if late_message is None:
+        late_delivery = await self._close_prefetch_state(state)
+        if late_delivery is None:
             return
 
-        queue.put_nowait(
-            _PrefetchedMessage(
-                data=late_message,
-                owns_prefetch_slot=False,
-            ),
-        )
+        queue.put_nowait(late_delivery)
         try:
             await self._notify_prefetch_hook("on_prefetch_queue_add")
         except BaseException as exc:
@@ -610,9 +613,8 @@ async def _enqueue_late_prefetched_message(
     async def _close_prefetch_state(
         self,
         state: _PrefetchState,
-    ) -> bytes | AckableMessage | None:
+    ) -> _PrefetchedMessage | None:
         """Close the pending read and iterator, retaining their first error."""
-        self._release_prefetch_slot(state)
         late_message: bytes | AckableMessage | None = None
         current_message = state.current_message
         state.current_message = None
@@ -631,7 +633,17 @@ async def _close_prefetch_state(
             pass
         except BaseException as exc:
             self._record_listen_error(exc)
-        return late_message
+
+        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."""
@@ -667,7 +679,7 @@ async def _discard_queued_messages(
                 break
             if isinstance(queued_message, _PrefetchedMessage):
                 discarded_messages += 1
-                if queued_message.owns_prefetch_slot:
+                if queued_message.owns_delivery_slot:
                     self.sem_prefetch.release()
 
         if discarded_messages:
@@ -699,20 +711,6 @@ 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:
                 queued_message = await queue.get()
@@ -725,16 +723,22 @@ 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
-                task = await self._start_callback(queued_message)
-                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.
                 #
                 # Because if we won't save it anywhere,
                 # python's GC can silently cancel task
                 # and this behaviour considered to be a Hisenbug.
                 # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
-                task.add_done_callback(task_cb)
+                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
@@ -743,29 +747,46 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
     async def _start_callback(
         self,
         message: _PrefetchedMessage,
-    ) -> asyncio.Task[Any]:
-        """Transfer prefetch and execution capacity to a callback task."""
-        owns_prefetch_slot = message.owns_prefetch_slot
+    ) -> _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 owns_prefetch_slot:
+            if self.sem is None and owns_delivery_slot:
                 self.sem_prefetch.release()
-                owns_prefetch_slot = False
-            return asyncio.create_task(
-                self.callback(message=message.data, raise_err=False),
+                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_prefetch_slot:
+            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:

tests/receiver/receiver_listener_support.py

@@ -33,9 +33,13 @@ class ObservedSemaphore(asyncio.Semaphore):
     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()