feat: add ReceiverObserver protocol and Prometheus receiver metrics

Add production observability for the Receiver via an observer protocol
that tracks prefetch queue depth, semaphore availability, active task
count, unknown task lookups, and deserialization errors.

- Add ReceiverObserver protocol (taskiq/receiver/observer.py)
- Instrument Receiver with guarded observer callbacks at 5 sites
- Add PrometheusReceiverObserver implementation with Gauges/Counters
- Wire observer from middleware to receiver via broker attribute
- Remove redundant in_flight_tasks gauge (replaced by active_tasks_count)

Author: Mohammed0tarek

taskiq/cli/worker/run.py

@@ -13,7 +13,7 @@
 from taskiq.cli.utils import import_object, import_tasks
 from taskiq.cli.worker.args import WorkerArgs
 from taskiq.cli.worker.process_manager import ProcessManager
-from taskiq.receiver import Receiver
+from taskiq.receiver import Receiver, ReceiverObserver
 
 try:
     import uvloop
@@ -163,6 +163,7 @@ def interrupt_handler(signum: int, _frame: Any) -> None:
             receiver = receiver_type(
                 broker=broker,
                 executor=pool,
+                observer=getattr(broker, "_receiver_observer", None),
                 validate_params=not args.no_parse,
                 max_async_tasks=args.max_async_tasks,
                 max_prefetch=args.max_prefetch,

taskiq/middlewares/prometheus_middleware.py

@@ -7,6 +7,7 @@
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
+from taskiq.receiver.observer import ReceiverObserver
 
 logger = getLogger("taskiq.prometheus")
 
@@ -75,12 +76,6 @@ def __init__(
             ["task_name"],
         )
 
-        self.in_flight_tasks = Gauge(
-            "in_flight_tasks",
-            "Number of tasks in flight",
-            ["task_name"],
-            multiprocess_mode="livesum",
-        )
         self.queue_wait_seconds = Histogram(
             "queue_wait_seconds",
             "time task spent in message queue",
@@ -166,7 +161,6 @@ def pre_execute(
                 time_delta,
             )
 
-        self.in_flight_tasks.labels(message.task_name).inc()
         self.received_tasks.labels(message.task_name).inc()
         return message
 
@@ -203,9 +197,12 @@ def post_execute(
             self.found_errors.labels(message.task_name).inc()
         else:
             self.success_tasks.labels(message.task_name).inc()
-        self.in_flight_tasks.labels(message.task_name).dec()
         self.execution_time.labels(message.task_name).observe(result.execution_time)
 
+    def set_broker(self, broker: "AsyncBroker") -> None:  # noqa: F821 pyright: ignore[reportUnknownVariableType]
+        super().set_broker(broker)
+        broker._receiver_observer = PrometheusReceiverObserver()
+
     def post_save(
         self,
         message: "TaskiqMessage",
@@ -218,3 +215,55 @@ def post_save(
         :param result: result of execution.
         """
         self.saved_results.labels(message.task_name).inc()
+
+
+class PrometheusReceiverObserver(ReceiverObserver):
+    """Receiver observer implementation for prometheus."""
+
+    def __init__(self) -> None:
+        try:
+            from prometheus_client import Counter, Gauge  # noqa: PLC0415
+        except ImportError as exc:
+            raise ImportError(
+                "Cannot initialize metrics. Please install 'taskiq[metrics]'.",
+            ) from exc
+
+        self.prefetch_queue_size = Gauge(
+            "prefetch_queue_size",
+            "The number of task in the prefetch queue.",
+            multiprocess_mode="livesum",
+        )
+        self.semaphore_available = Gauge(
+            "semaphore_available",
+            "Number of semaphore slots available in broker",
+            multiprocess_mode="livesum",
+        )
+        self.active_tasks_count = Gauge(
+            "worker_active_tasks_count",
+            "Number of active tasks in worker",
+            multiprocess_mode="livesum",
+        )
+        self.task_not_found_total = Counter(
+            "task_not_found_total",
+            "Number of times the worker got a task not registered",
+            ["task_name"],
+        )
+        self.deserialize_error = Counter(
+            "deserialize_error_count",
+            "Number of times broker faced a desrialization error",
+        )
+
+    def on_prefetch_queue_size(self, size: int) -> None:
+        self.prefetch_queue_size.set(size)
+
+    def on_semaphore_status(self, available: int) -> None:
+        self.semaphore_available.set(available)
+
+    def on_active_tasks_count(self, count: int) -> None:
+        self.active_tasks_count.set(count)
+
+    def on_task_not_found(self, task_name: str) -> None:
+        self.task_not_found_total.labels(task_name).inc()
+
+    def on_deserialize_error(self, raw: bytes, error: Exception) -> None:
+        self.deserialize_error.inc()

taskiq/receiver/__init__.py

@@ -1,5 +1,6 @@
 """Package for message receiver."""
 
 from taskiq.receiver.receiver import Receiver
+from taskiq.receiver.observer import ReceiverObserver
 
-__all__ = ["Receiver"]
+__all__ = ["Receiver", "ReceiverObserver"]

taskiq/receiver/observer.py

@@ -0,0 +1,21 @@
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class ReceiverObserver(Protocol):
+    """
+    Observer for reciever stats.
+
+    This classs is used to observe/collect metrics for the receiver.
+    This includes semaphore usage, tasks in queue, etc.
+
+    metrics tracked:
+    - Number of tasks in queue
+    - Number of taks in execution (Semaphore uusage)
+    """
+
+    def on_prefetch_queue_size(self, size: int) -> None: ...
+    def on_semaphore_status(self, available: int) -> None: ...
+    def on_active_tasks_count(self, count: int) -> None: ...
+    def on_task_not_found(self, task_name: str) -> None: ...
+    def on_deserialize_error(self, raw: bytes, error: Exception) -> None: ...

taskiq/receiver/receiver.py

@@ -18,6 +18,7 @@
 from taskiq.context import Context
 from taskiq.exceptions import NoResultError
 from taskiq.message import TaskiqMessage
+from taskiq.receiver.observer import ReceiverObserver
 from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
 from taskiq.state import TaskiqState
@@ -35,6 +36,7 @@ def __init__(
         self,
         broker: AsyncBroker,
         executor: Executor | None = None,
+        observer: ReceiverObserver | None = None,
         validate_params: bool = True,
         max_async_tasks: "int | None" = None,
         max_prefetch: int = 0,
@@ -54,6 +56,7 @@ def __init__(
         self.dependency_graphs: dict[str, DependencyGraph] = {}
         self.propagate_exceptions = propagate_exceptions
         self.on_exit = on_exit
+        self.observer = observer
         self.ack_time = ack_type or AcknowledgeType.WHEN_SAVED
         self.known_tasks: set[str] = set()
         self.max_tasks_to_execute = max_tasks_to_execute
@@ -92,6 +95,11 @@ async def callback(  # noqa: C901, PLR0912
             taskiq_msg = self.broker.formatter.loads(message=message_data)
             taskiq_msg.parse_labels()
         except Exception as exc:
+            if self.observer is not None:
+                self.observer.on_deserialize_error(
+                    raw=message_data,
+                    error=exc,
+                )
             logger.warning(
                 "Cannot parse message: %s. Skipping execution.\n %s",
                 message_data,
@@ -102,6 +110,11 @@ async def callback(  # noqa: C901, PLR0912
         logger.debug(f"Received message: {taskiq_msg}")
         task = self.broker.find_task(taskiq_msg.task_name)
         if task is None:
+            if self.observer is not None:
+                self.observer.on_task_not_found(
+                    taskiq_msg.task_name,
+                )
+
             logger.warning(
                 'task "%s" is not found. Maybe you forgot to import it?',
                 taskiq_msg.task_name,
@@ -363,6 +376,7 @@ async def prefetcher(
                 break
             try:
                 await self.sem_prefetch.acquire()
+
                 if (
                     self.max_tasks_to_execute
                     and fetched_tasks >= self.max_tasks_to_execute
@@ -376,13 +390,20 @@ async def prefetcher(
                 # and continue the loop. So it will check if finished event was set.
                 if not done:
                     self.sem_prefetch.release()
+
                     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)
+
+                if self.observer is not None:
+                    self.observer.on_prefetch_queue_size(
+                        queue.qsize(),
+                    )
+
             except (asyncio.CancelledError, StopAsyncIteration):
                 break
         # We don't want to fetch new messages if we are shutting down.
@@ -413,17 +434,35 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
             :param task: finished task
             """
             tasks.discard(task)
+            if self.observer is not None:
+                self.observer.on_active_tasks_count(
+                    len(tasks),
+                )
+
             if self.sem is not None:
                 self.sem.release()
 
+                if self.observer is not None:
+                    self.observer.on_semaphore_status(
+                        self.sem._value  # noqa
+                    )
+
         while True:
             try:
                 # Waits for semaphore to be released.
                 if self.sem is not None:
                     await self.sem.acquire()
+                    if self.observer is not None:
+                        self.observer.on_semaphore_status(
+                            self.sem._value  # noqa
+                        )
 
                 self.sem_prefetch.release()
                 message = await queue.get()
+                if self.observer is not None:
+                    self.observer.on_prefetch_queue_size(
+                        queue.qsize()  # noqa
+                    )
                 if message is QUEUE_DONE:
                     # asyncio.wait will throw an error if there is nothing to wait for
                     if tasks:
@@ -438,7 +477,10 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                     self.callback(message=message, raise_err=False),
                 )
                 tasks.add(task)
-
+                if self.observer is not None:
+                    self.observer.on_active_tasks_count(
+                        len(tasks),
+                    )
                 # We want the task to remove itself from the set when it's done.
                 #
                 # Because if we won't save it anywhere,