feat: add worker queue and active task OTel metrics

- Add worker_active_tasks UpDownCounter driven by pre/post_execute hooks
- Add worker_prefetched_tasks UpDownCounter via on_prefetch_queue_add/remove hooks in receiver
- Add worker_cpu_utilization and worker_memory_utilization observable gauges (worker process only)
- Add tests for all new metrics

Author: Mohammed0tarek

taskiq/middlewares/opentelemetry_middleware.py

@@ -1,10 +1,11 @@
 import logging
+from collections.abc import Generator
 from contextlib import AbstractContextManager
 from datetime import datetime, timezone
 from importlib.metadata import version
-from typing import Any, Generator, TypeVar
-import psutil
+from typing import Any, TypeVar
 
+import psutil
 from packaging.version import Version, parse
 
 try:
@@ -223,11 +224,24 @@ def __init__(
             description="Worker memory utilization in bytes. Only for worker processes",
         )
 
-    def _observe_memory(self, _options: Any) -> Generator[Observation, None, None]:
+        # 8- Number of tasks executing
+        self.number_of_broker_active_tasks = self._meter.create_up_down_counter(
+            "worker_active_tasks",
+            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:
             yield Observation(self._process.memory_info().rss)
 
-    def _observe_cpu(self, _options: Any) -> Generator[Observation, None, None]:
+    def _observe_cpu(self, options: Any) -> Generator[Observation, None, None]:
         if self.broker and self.broker.is_worker_process:
             yield Observation(self._process.cpu_percent())
 
@@ -298,6 +312,10 @@ def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:
         activation.__enter__()  # pylint: disable=E1101
         attach_context(message, span, activation, token)
         message.labels[_TASK_RECEIVED_TIME_KEY] = datetime.now(timezone.utc).timestamp()
+        self.number_of_broker_active_tasks.add(
+            1,
+            attributes={"task_name": message.task_name},
+        )
         return message
 
     def post_save(  # pylint: disable=R6301
@@ -424,3 +442,16 @@ def post_execute(
                 amount=task_receive_time - task_send_time,
                 attributes={"task_name": message.task_name},
             )
+
+        self.number_of_broker_active_tasks.add(
+            -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

@@ -383,6 +383,12 @@ async def prefetcher(
                 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
+                        )
             except (asyncio.CancelledError, StopAsyncIteration):
                 break
         # We don't want to fetch new messages if we are shutting down.
@@ -434,6 +440,13 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                         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),
                 )

tests/opentelemetry/test_metrics.py

@@ -6,6 +6,7 @@
 from opentelemetry.test.test_base import TestBase
 
 from taskiq.instrumentation import TaskiqInstrumentor
+from taskiq.middlewares.opentelemetry_middleware import OpenTelemetryMiddleware
 
 from .taskiq_test_tasks import (
     broker,
@@ -58,6 +59,7 @@ async def test() -> None:
             "task_success",
             "task_execution_time",
             "task_wait_time",
+            "worker_active_tasks",
         }
         found = {
             metric.name
@@ -167,3 +169,56 @@ async def test() -> None:
         points = self._get_data_points("tasks_sent")
         self.assertEqual(len(points), 1)
         self.assertEqual(points[0].value, 3)
+
+    def test_active_tasks_counter(self) -> None:
+        async def test() -> None:
+            for _ in range(3):
+                await task_add.kiq(1, 2)
+            await broker.wait_all()
+
+        asyncio.run(test())
+
+        points = self._get_data_points("worker_active_tasks")
+        # all 3 tasks share the same task_name so they aggregate into one data point
+        self.assertEqual(len(points), 1)
+        # net zero: pre_execute incremented, post_execute decremented for each task
+        self.assertEqual(points[0].value, 0)
+        self.assertIn("task_name", points[0].attributes)
+        self.assertEqual(
+            points[0].attributes.get("task_name"),
+            "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)
+        )
+        middleware.set_broker(broker)
+        broker.is_worker_process = True
+        try:
+            metrics_data = self.reader.get_metrics_data()
+            self.assertIsNotNone(metrics_data)
+            found = {
+                metric.name
+                for rm in metrics_data.resource_metrics  # type: ignore[union-attr]
+                for sm in rm.scope_metrics
+                for metric in sm.metrics
+            }
+            self.assertIn("worker_cpu_utilization", found)
+            self.assertIn("worker_memory_utilization", found)
+        finally:
+            broker.is_worker_process = False
+            middleware.set_broker(None)  # type: ignore[arg-type]