feat: add queue_wait_seconds, task_errors_by_type, and pre_send timestamp to PrometheusMiddleware

Author: Mohammed0tarek

taskiq/abc/broker.py

@@ -32,10 +32,13 @@
 from taskiq.utils import maybe_awaitable
 from taskiq.warnings import TaskiqDeprecationWarning
 
-if sys.version_info >= (3, 11):
+if sys.version_info >= (
+    3,
+    11,
+):  # Check which python version are we running to import correctly
     from typing import Self
 else:
-    from typing_extensions import Self
+    from typing_extensions import Self  # pragma: no cover
 
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -46,6 +49,7 @@
 _FuncParams = ParamSpec("_FuncParams")
 _ReturnType = TypeVar("_ReturnType")
 
+# an event handler can be either a sync or an async function that has one parameter of type TaskiqState
 EventHandler: TypeAlias = Callable[[TaskiqState], Awaitable[None] | None]
 
 logger = getLogger("taskiq")

taskiq/cli/worker/process_manager.py

@@ -169,7 +169,7 @@ def __init__(
             for path_to_watch in watch_paths:
                 logger.debug(f"Watching directory: {path_to_watch}")
                 observer.schedule(
-                    FileWatcher(
+                    FileWatcher(  # type: ignore
                         callback=schedule_workers_reload,
                         path=Path(path_to_watch),
                         use_gitignore=not args.no_gitignore,

taskiq/middlewares/prometheus_middleware.py

@@ -1,9 +1,9 @@
+import datetime
 import os
 from logging import getLogger
 from pathlib import Path
 from tempfile import gettempdir
 from typing import Any
-
 from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.message import TaskiqMessage
 from taskiq.result import TaskiqResult
@@ -43,7 +43,7 @@ def __init__(
         logger.debug("Initializing metrics")
 
         try:
-            from prometheus_client import Counter, Histogram  # noqa: PLC0415
+            from prometheus_client import Counter, Histogram, Gauge  # noqa: PLC0415
         except ImportError as exc:
             raise ImportError(
                 "Cannot initialize metrics. Please install 'taskiq[metrics]'.",
@@ -74,6 +74,24 @@ def __init__(
             "Time of function execution",
             ["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",
+            ["task_name"],
+        )
+        self.task_errors_by_type = Counter(
+            "task_errors_by_type",
+            "Number of errors raised in tasks by their type",
+            ["task_name", "error_type"],
+        )
+
         self.server_port = server_port
         self.server_addr = server_addr
 
@@ -104,6 +122,24 @@ def startup(self) -> None:
             except OSError as exc:
                 logger.debug("Cannot start prometheus server: %s", exc)
 
+    def pre_send(
+        self,
+        message: "TaskiqMessage",
+    ) -> "TaskiqMessage":
+        """
+        Function to track the time a task spend in queue.
+
+        This function tracks the time a task spends in a queue until it is executed.
+
+        :param message: current message.
+        :return: message
+        """
+        if not message.labels.get("_taskiq_enqueue_timestamp"):
+            message.labels["_taskiq_enqueue_timestamp"] = datetime.datetime.now(
+                datetime.UTC,
+            ).isoformat()  # Might conside using timezones too
+        return message
+
     def pre_execute(
         self,
         message: "TaskiqMessage",
@@ -117,9 +153,41 @@ def pre_execute(
         :param message: current message.
         :return: message
         """
+        if message.labels.get(
+            "_taskiq_enqueue_timestamp",
+        ):  # Handle case where the sender doesn't use the prometheus middleware
+            time_delta = datetime.datetime.now(
+                datetime.UTC,
+            ) - datetime.datetime.fromisoformat(
+                message.labels["_taskiq_enqueue_timestamp"],
+            )
+            time_delta = max(0, time_delta.total_seconds())
+            self.queue_wait_seconds.labels(message.task_name).observe(
+                time_delta,
+            )
+
+        self.in_flight_tasks.labels(message.task_name).inc()
         self.received_tasks.labels(message.task_name).inc()
         return message
 
+    def on_error(
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult[Any],  # pylint: disable=unused-argument
+        exception: BaseException,
+    ) -> None:
+        """
+        This function tracks the number of errors raised by tasks.
+
+        :param message: the received task message
+        :param result: the result of task
+        :param exception: exception raised
+        """
+        self.task_errors_by_type.labels(
+            message.task_name,
+            type(exception).__name__,
+        ).inc()
+
     def post_execute(
         self,
         message: "TaskiqMessage",
@@ -135,6 +203,7 @@ 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 post_save(