fix: make PrometheusMiddleware import-safe (fix #397)

taskiq/middlewares/prometheus_middleware.py

@@ -18,9 +18,15 @@ class PrometheusMiddleware(TaskiqMiddleware):
     This middleware starts wsgi server with prometheus metrics.
     Also it updates metrics on events.
 
+    The middleware is import-safe: creating multiple instances
+    in the same process (e.g. when the broker module is imported
+    more than once during task discovery) reuses already registered
+    collectors instead of raising
+    ``ValueError: Duplicated timeseries in CollectorRegistry``.
+
     :param server_port: The port to listen on.
     :param server_addr: The address to listen on.
-    :paam metrics_path: The path to store metrics for multiproc env.
+    :param metrics_path: The path to store metrics for multiproc env.
     """
 
     def __init__(
@@ -43,33 +49,65 @@ def __init__(
         logger.debug("Initializing metrics")
 
         try:
-            from prometheus_client import Counter, Histogram  # noqa: PLC0415
+            from prometheus_client import (  # noqa: PLC0415
+                REGISTRY,
+                Counter,
+                Histogram,
+            )
         except ImportError as exc:
             raise ImportError(
                 "Cannot initialize metrics. Please install 'taskiq[metrics]'.",
             ) from exc
 
-        self.found_errors = Counter(
+        def _get_or_create_counter(
+            name: str,
+            documentation: str,
+            labelnames: list[str],
+        ) -> Counter:
+            """Return existing counter or create a new one."""
+            try:
+                return Counter(name, documentation, labelnames)
+            except ValueError:
+                existing = REGISTRY._names_to_collectors.get(name)  # noqa: SLF001
+                if existing is None or not isinstance(existing, Counter):
+                    raise
+                return existing
+
+        def _get_or_create_histogram(
+            name: str,
+            documentation: str,
+            labelnames: list[str],
+        ) -> Histogram:
+            """Return existing histogram or create a new one."""
+            try:
+                return Histogram(name, documentation, labelnames)
+            except ValueError:
+                existing = REGISTRY._names_to_collectors.get(name)  # noqa: SLF001
+                if existing is None or not isinstance(existing, Histogram):
+                    raise
+                return existing
+
+        self.found_errors = _get_or_create_counter(
             "found_errors",
             "Number of found errors",
             ["task_name"],
         )
-        self.received_tasks = Counter(
+        self.received_tasks = _get_or_create_counter(
             "received_tasks",
             "Number of received tasks",
             ["task_name"],
         )
-        self.success_tasks = Counter(
+        self.success_tasks = _get_or_create_counter(
             "success_tasks",
             "Number of successfully executed tasks",
             ["task_name"],
         )
-        self.saved_results = Counter(
+        self.saved_results = _get_or_create_counter(
             "saved_results",
             "Number of saved results in result backend",
             ["task_name"],
         )
-        self.execution_time = Histogram(
+        self.execution_time = _get_or_create_histogram(
             "execution_time",
             "Time of function execution",
             ["task_name"],

tests/middlewares/test_prometheus_middleware.py

@@ -0,0 +1,90 @@
+import pytest
+
+from taskiq import PrometheusMiddleware
+from taskiq.message import TaskiqMessage
+from taskiq.result import TaskiqResult
+
+pytest.importorskip("prometheus_client")
+
+
+def _make_message(task_name: str = "test_task") -> TaskiqMessage:
+    return TaskiqMessage(
+        task_id="test_id",
+        task_name=task_name,
+        labels={},
+        args=[],
+        kwargs={},
+    )
+
+
+def test_multiple_instances_do_not_raise_duplicate_timeseries() -> None:
+    """Regression test for https://github.com/taskiq-python/taskiq/issues/397."""
+    first = PrometheusMiddleware(server_port=19001)
+    second = PrometheusMiddleware(server_port=19001)
+
+    assert first.found_errors is second.found_errors
+    assert first.received_tasks is second.received_tasks
+    assert first.success_tasks is second.success_tasks
+    assert first.saved_results is second.saved_results
+    assert first.execution_time is second.execution_time
+
+
+def test_metrics_still_work_after_reuse() -> None:
+    first = PrometheusMiddleware(server_port=19002)
+    second = PrometheusMiddleware(server_port=19002)
+
+    message = _make_message()
+
+    second.pre_execute(message)
+    second.post_execute(
+        message,
+        TaskiqResult(is_err=False, return_value=None, execution_time=0.01),
+    )
+    second.post_execute(
+        message,
+        TaskiqResult(is_err=True, return_value=None, execution_time=0.02),
+    )
+    second.post_save(
+        message,
+        TaskiqResult(is_err=False, return_value=None, execution_time=0.01),
+    )
+
+    # Both instances share collectors, so increments are visible either way.
+    assert first.received_tasks.labels(message.task_name)._value.get() >= 1
+    assert first.success_tasks.labels(message.task_name)._value.get() >= 1
+    assert first.found_errors.labels(message.task_name)._value.get() >= 1
+    assert first.saved_results.labels(message.task_name)._value.get() >= 1
+
+
+def test_counter_type_mismatch_reraises() -> None:
+    """Wrong-type collector for a Counter name must re-raise ValueError."""
+    from prometheus_client import REGISTRY, Histogram  # noqa: PLC0415
+
+    existing = REGISTRY._names_to_collectors.get("found_errors")
+    if existing is not None:
+        REGISTRY.unregister(existing)
+    wrong = Histogram("found_errors", "Number of found errors", ["task_name"])
+    try:
+        with pytest.raises(ValueError):
+            PrometheusMiddleware(server_port=19003)
+    finally:
+        REGISTRY.unregister(wrong)
+        # Restore healthy state for other tests.
+        PrometheusMiddleware(server_port=19003)
+
+
+def test_histogram_type_mismatch_reraises() -> None:
+    """Wrong-type collector for a Histogram name must re-raise ValueError."""
+    from prometheus_client import REGISTRY, Counter  # noqa: PLC0415
+
+    existing = REGISTRY._names_to_collectors.get("execution_time")
+    if existing is not None:
+        REGISTRY.unregister(existing)
+    wrong = Counter("execution_time", "Time of function execution", ["task_name"])
+    try:
+        with pytest.raises(ValueError):
+            PrometheusMiddleware(server_port=19004)
+    finally:
+        REGISTRY.unregister(wrong)
+        # Restore healthy state for other tests.
+        PrometheusMiddleware(server_port=19004)