diff --git a/taskiq/middlewares/prometheus_middleware.py b/taskiq/middlewares/prometheus_middleware.py index c3968cf..bbc78f5 100644 --- a/taskiq/middlewares/prometheus_middleware.py +++ b/taskiq/middlewares/prometheus_middleware.py @@ -10,6 +10,18 @@ logger = getLogger("taskiq.prometheus") +try: + from prometheus_client import ( # noqa: WPS433 + REGISTRY, + CollectorRegistry, + Counter, + Histogram, + ) +except ImportError as imp_exc: + raise ImportError( + "Cannot initialize metrics. Please install 'taskiq[metrics]'.", + ) from imp_exc + class PrometheusMiddleware(TaskiqMiddleware): """ @@ -24,6 +36,7 @@ class PrometheusMiddleware(TaskiqMiddleware): def __init__( self, metrics_path: Optional[Path] = None, + metrics_registry: CollectorRegistry = REGISTRY, server_port: int = 9000, server_addr: str = "0.0.0.0", # noqa: S104 ) -> None: @@ -41,38 +54,37 @@ def __init__( logger.debug("Initializing metrics") - try: - from prometheus_client import Counter, Histogram # noqa: WPS433 - except ImportError as exc: - raise ImportError( - "Cannot initialize metrics. Please install 'taskiq[metrics]'.", - ) from exc - self.found_errors = Counter( "found_errors", "Number of found errors", ["task_name"], + registry=metrics_registry, ) self.received_tasks = Counter( "received_tasks", "Number of received tasks", ["task_name"], + registry=metrics_registry, ) self.success_tasks = Counter( "success_tasks", "Number of successfully executed tasks", ["task_name"], + registry=metrics_registry, ) self.saved_results = Counter( "saved_results", "Number of saved results in result backend", ["task_name"], + registry=metrics_registry, ) self.execution_time = Histogram( "execution_time", - "Tome of function execution", + "Time of function execution", ["task_name"], + registry=metrics_registry, ) + self.metrics_registry = metrics_registry self.server_port = server_port self.server_addr = server_addr @@ -87,7 +99,11 @@ def startup(self) -> None: if self.broker.is_worker_process: try: - start_http_server(port=self.server_port, addr=self.server_addr) + start_http_server( + port=self.server_port, + addr=self.server_addr, + registry=self.metrics_registry, + ) except OSError as exc: logger.debug("Cannot start prometheus server: %s", exc)