fix(receiver): simplify prefetch listener lifecycle

- remove undocumented prefetch middleware hooks and metrics
- preserve listener errors during task group failures
- restore capacity when callback handoff fails
- strengthen shutdown and cleanup coverage

Refs: #528

Author: GefMar

taskiq/middlewares/opentelemetry_middleware.py

@@ -230,12 +230,6 @@ def __init__(
             unit="1",
             description="Number of tasks currently executing in the worker.",
         )
-        # 9- Number of tasks executing
-        self.number_of_broker_prefetched_tasks = self._meter.create_up_down_counter(
-            "worker_prefetched_tasks",
-            unit="1",
-            description="Number of tasks currently prefetched in the worker.",
-        )
 
     def _observe_memory(self, options: Any) -> Generator[Observation, None, None]:
         if self.broker and self.broker.is_worker_process:
@@ -447,11 +441,3 @@ def post_execute(
             -1,
             attributes={"task_name": message.task_name},
         )
-
-    def on_prefetch_queue_add(self) -> None:
-        """This hook is called after task is added to the worker prefetch queue."""
-        self.number_of_broker_prefetched_tasks.add(1)
-
-    def on_prefetch_queue_remove(self) -> None:
-        """This hook is called after task is removed from the worker prefetch queue."""
-        self.number_of_broker_prefetched_tasks.add(-1)

taskiq/receiver/receiver.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextlib
 import contextvars
 import functools
 import inspect
@@ -10,7 +11,7 @@
 from enum import Enum, auto
 from logging import getLogger
 from time import time
-from typing import Any, Literal, get_type_hints
+from typing import Any, get_type_hints
 
 import anyio
 from taskiq_dependencies import DependencyGraph
@@ -442,7 +443,7 @@ async def listen(self, finish_event: asyncio.Event) -> None:  # pragma: no cover
                 logger.error(
                     "A Receiver listener lifecycle error was recorded before "
                     "the task group failed.",
-                    exc_info=(type(error), error, error.__traceback__),
+                    exc_info=error,
                 )
             raise
 
@@ -463,12 +464,7 @@ async def prefetcher(
         :param queue: queue for prefetched data.
         :param finish_event: event to indicate that we need to stop prefetching.
         """
-        try:
-            state = _PrefetchState(iterator=self.broker.listen())
-        except BaseException as exc:
-            self._record_listen_error(exc)
-            queue.put_nowait(_QueueSignal.DONE)
-            return
+        state = _PrefetchState(iterator=self.broker.listen())
 
         fetched_tasks = 0
         finish_waiter = asyncio.create_task(finish_event.wait())
@@ -494,13 +490,6 @@ async def prefetcher(
                     ),
                 )
                 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):
@@ -509,7 +498,8 @@ async def prefetcher(
                 finally:
                     queue.put_nowait(_QueueSignal.DONE)
                     finish_waiter.cancel()
-                    await asyncio.gather(finish_waiter, return_exceptions=True)
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await finish_waiter
 
     def _should_stop_prefetch(
         self,
@@ -563,15 +553,15 @@ async def _acquire_delivery_slot(
         """Acquire one delivery-admission slot unless shutdown wins the race."""
         acquire_task = asyncio.create_task(self.sem_prefetch.acquire())
         try:
-            done, _ = await asyncio.wait(
+            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():
+        if finish_waiter.done() or finish_event.is_set():
             await self._settle_delivery_acquire(acquire_task)
             return False
 
@@ -605,10 +595,6 @@ async def _enqueue_late_prefetched_message(
             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,
@@ -645,38 +631,11 @@ async def _close_prefetch_state(
         state.owns_delivery_slot = False
         return late_delivery
 
-    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):
-            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."""
+        """Release capacity for abandoned queue entries."""
         discarded_messages = 0
         while True:
             try:
@@ -694,18 +653,6 @@ async def _discard_queued_messages(
                 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[_PrefetchedMessage | _QueueSignal]",
@@ -737,7 +684,7 @@ async def runner(
                     except BaseException:
                         queue.put_nowait(queued_message)
                         raise
-                started_callback = await self._start_callback(
+                started_callback = self._start_callback(
                     queued_message,
                     owns_execution_slot=owns_execution_slot,
                 )
@@ -761,7 +708,7 @@ async def runner(
                 break
         logger.info("The runner is stopped.")
 
-    async def _start_callback(
+    def _start_callback(
         self,
         message: _PrefetchedMessage,
         *,
@@ -770,8 +717,6 @@ async def _start_callback(
         """Transfer execution and delivery capacity to a callback task."""
         owns_delivery_slot = message.owns_delivery_slot
         try:
-            await self._notify_prefetch_hook("on_prefetch_queue_remove")
-
             if self.sem is None and owns_delivery_slot:
                 self.sem_prefetch.release()
                 owns_delivery_slot = False

tests/opentelemetry/test_metrics.py

@@ -189,19 +189,6 @@ async def test() -> None:
             "tests.opentelemetry.taskiq_test_tasks:task_add",
         )
 
-    def test_prefetch_queue_counter(self) -> None:
-        middleware = next(
-            m for m in broker.middlewares if isinstance(m, OpenTelemetryMiddleware)
-        )
-        middleware.on_prefetch_queue_add()
-        middleware.on_prefetch_queue_add()
-        middleware.on_prefetch_queue_add()
-        middleware.on_prefetch_queue_remove()
-
-        points = self._get_data_points("worker_prefetched_tasks")
-        self.assertEqual(len(points), 1)
-        self.assertEqual(points[0].value, 2)
-
     def test_worker_resource_metrics_when_worker_process(self) -> None:
         middleware = next(
             m for m in broker.middlewares if isinstance(m, OpenTelemetryMiddleware)

tests/receiver/receiver_listener_support.py

@@ -3,28 +3,11 @@
 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
+    """Marker error for listener lifecycle failures."""
 
 
 class ObservedSemaphore(asyncio.Semaphore):