diff --git a/taskiq/middlewares/prometheus_middleware.py b/taskiq/middlewares/prometheus_middleware.py index a7713d3e..ae86a7c2 100644 --- a/taskiq/middlewares/prometheus_middleware.py +++ b/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"], diff --git a/tests/middlewares/test_prometheus_middleware.py b/tests/middlewares/test_prometheus_middleware.py new file mode 100644 index 00000000..0af4a269 --- /dev/null +++ b/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)