From f4e5e2ab02c0233e758c1e97519e421b32642509 Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 2 Jul 2026 13:13:06 +0300 Subject: [PATCH] Harden Prometheus metrics --- Dockerfile | 5 +- pyproject.toml | 1 + src/infrastructure/metrics.py | 237 ++++++++++++---------------------- uv.lock | 12 +- 4 files changed, 101 insertions(+), 154 deletions(-) diff --git a/Dockerfile b/Dockerfile index 06c8580..64a8d4e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,8 +21,9 @@ COPY --from=builder /app/src /app/src ENV PATH="/app/.venv/bin:$PATH" \ PYTHONUNBUFFERED=1 \ PYTHONDONTWRITEBYTECODE=1 \ - PYTHONPATH=/app + PYTHONPATH=/app \ + PROMETHEUS_MULTIPROC_DIR=/tmp/prometheus_multiproc EXPOSE 8000 -CMD ["sh", "-c", "granian --interface asgi ${APP_MODULE:-src.main:app} --host ${APP_HOST:-0.0.0.0} --port ${APP_PORT:-8000} --workers ${APP_WORKERS:-1} --loop uvloop"] +CMD ["sh", "-c", "mkdir -p "$PROMETHEUS_MULTIPROC_DIR" && rm -f "$PROMETHEUS_MULTIPROC_DIR"/* && granian --interface asgi ${APP_MODULE:-src.main:app} --host ${APP_HOST:-0.0.0.0} --port ${APP_PORT:-8000} --workers ${APP_WORKERS:-1} --loop uvloop"] diff --git a/pyproject.toml b/pyproject.toml index 5eff682..54b3186 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "hvac==2.4.0", "itsdangerous==2.2.0", "orjson==3.11.7", + "prometheus-client==0.23.1", "pydantic-settings==2.12.0", "python-jose==3.5.0", "python-ulid==3.1.0", diff --git a/src/infrastructure/metrics.py b/src/infrastructure/metrics.py index c1e09a1..5dd36f7 100644 --- a/src/infrastructure/metrics.py +++ b/src/infrastructure/metrics.py @@ -1,11 +1,14 @@ from __future__ import annotations +import asyncio +import os import time -from collections import defaultdict +from collections.abc import Awaitable, Callable from datetime import timedelta -from typing import Iterable +from typing import Any from fastapi import Request, Response +from prometheus_client import CollectorRegistry, Counter, Histogram, CONTENT_TYPE_LATEST, generate_latest, multiprocess from sqlalchemy import distinct, func, select from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.responses import Response as StarletteResponse @@ -14,183 +17,115 @@ from src.infrastructure.database.context import async_session_maker from src.infrastructure.database.models.sessions import Session from src.infrastructure.database.models.user import UserModel - -SERVICE_NAME = 'auth' -WINDOWS: tuple[tuple[str, timedelta | None], ...] = ( - ('all', None), - ('30d', timedelta(days=30)), - ('7d', timedelta(days=7)), - ('24h', timedelta(days=1)), - ('1h', timedelta(hours=1)), -) +SERVICE_NAME = "auth" +BUSINESS_CACHE_TTL_SECONDS = 30 LATENCY_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0) +WINDOWS = (("all", None), ("30d", timedelta(days=30)), ("7d", timedelta(days=7)), ("24h", timedelta(days=1)), ("1h", timedelta(hours=1))) -_http_requests: dict[tuple[str, str, str], int] = defaultdict(int) -_http_latency_sum: dict[tuple[str, str], float] = defaultdict(float) -_http_latency_count: dict[tuple[str, str], int] = defaultdict(int) -_http_latency_buckets: dict[tuple[str, str, float], int] = defaultdict(int) - - -def _format_labels(labels: dict[str, str]) -> str: - return ','.join(f'{key}="{value}"' for key, value in labels.items()) - - -def _sample(name: str, value: int | float, labels: dict[str, str] | None = None) -> str: - if not labels: - return f'{name} {value}' - return f'{name}{{{_format_labels(labels)}}} {value}' +HTTP_REQUESTS = Counter("auth_http_requests_total", "Total HTTP requests handled by auth service.", ("service", "method", "route", "status_code")) +HTTP_LATENCY = Histogram("auth_http_request_duration_seconds", "HTTP request duration in seconds for auth service.", ("service", "method", "route"), buckets=LATENCY_BUCKETS) def _route_template(request: Request) -> str: - route = request.scope.get('route') - path = getattr(route, 'path', None) - if path: - return str(path) - return request.url.path + route = request.scope.get("route") + path = getattr(route, "path", None) + return str(path) if path else "__unmatched__" class PrometheusMetricsMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> StarletteResponse: - if request.url.path == '/metrics': + if request.url.path == "/metrics": return await call_next(request) - started = time.perf_counter() - status_code = '500' - route = request.url.path - + status_code = "500" try: response = await call_next(request) status_code = str(response.status_code) return response finally: elapsed = time.perf_counter() - started - route = _route_template(request) or route - method = request.method - - _http_requests[(method, route, status_code)] += 1 - _http_latency_sum[(method, route)] += elapsed - _http_latency_count[(method, route)] += 1 - for bucket in LATENCY_BUCKETS: - if elapsed <= bucket: - _http_latency_buckets[(method, route, bucket)] += 1 + HTTP_REQUESTS.labels(SERVICE_NAME, request.method, _route_template(request), status_code).inc() + HTTP_LATENCY.labels(SERVICE_NAME, request.method, _route_template(request)).observe(elapsed) -async def _count_scalar(statement) -> int: - async with async_session_maker() as session: - result = await session.execute(statement) - return int(result.scalar_one() or 0) +class CachedBusinessMetrics: + def __init__(self, ttl_seconds: int, collector: Callable[[], Awaitable[list[str]]]) -> None: + self._ttl_seconds = ttl_seconds + self._collector = collector + self._expires_at = 0.0 + self._lines: list[str] = [] + self._lock = asyncio.Lock() + + async def get(self) -> list[str]: + now = time.monotonic() + if now < self._expires_at: + return self._lines + async with self._lock: + now = time.monotonic() + if now < self._expires_at: + return self._lines + self._lines = await self._collector() + self._expires_at = now + self._ttl_seconds + return self._lines + + +def _labels(labels: dict[str, str]) -> str: + escaped = {key: value.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") for key, value in labels.items()} + return ",".join(f'{key}="{value}"' for key, value in escaped.items()) + + +def _sample(name: str, value: int | float, labels: dict[str, str] | None = None) -> str: + return f"{name}{{{_labels(labels)}}} {value}" if labels else f"{name} {value}" + + +async def _scalar(session: Any, statement: Any) -> int: + result = await session.execute(statement) + return int(result.scalar_one() or 0) async def _collect_user_metrics() -> list[str]: - lines: list[str] = [] - - users_total = await _count_scalar(select(func.count()).select_from(UserModel)) - lines.append(_sample('auth_users_total', users_total, {'service': SERVICE_NAME})) - - active_sessions = await _count_scalar( - select(func.count()) - .select_from(Session) - .where(Session.revoked_at.is_(None)) - ) - lines.append(_sample('auth_sessions_active_total', active_sessions, {'service': SERVICE_NAME})) - - for window_name, window_delta in WINDOWS: - created_conditions = [] - active_conditions = [Session.revoked_at.is_(None)] - if window_delta is not None: - created_conditions.append(UserModel.created_at >= func.now() - window_delta) - active_conditions.append(Session.last_seen_at >= func.now() - window_delta) - - new_users = await _count_scalar( - select(func.count()) - .select_from(UserModel) - .where(*created_conditions) - ) - active_users = await _count_scalar( - select(func.count(distinct(Session.user_id))) - .select_from(Session) - .where(*active_conditions) - ) - - labels = {'service': SERVICE_NAME, 'window': window_name} - lines.append(_sample('auth_user_registrations_total', new_users, labels)) - lines.append(_sample('auth_active_unique_users_total', active_users, labels)) - + lines: list[str] = [ + "# TYPE auth_users_total gauge", + "# TYPE auth_sessions_active_total gauge", + "# TYPE auth_user_registrations_total gauge", + "# TYPE auth_active_unique_users_total gauge", + ] + async with async_session_maker() as session: + users_total = await _scalar(session, select(func.count()).select_from(UserModel)) + active_sessions = await _scalar(session, select(func.count()).select_from(Session).where(Session.revoked_at.is_(None))) + lines.append(_sample("auth_users_total", users_total, {"service": SERVICE_NAME})) + lines.append(_sample("auth_sessions_active_total", active_sessions, {"service": SERVICE_NAME})) + for window_name, window_delta in WINDOWS: + created_conditions = [] + active_conditions = [Session.revoked_at.is_(None)] + if window_delta is not None: + created_conditions.append(UserModel.created_at >= func.now() - window_delta) + active_conditions.append(Session.last_seen_at >= func.now() - window_delta) + new_users = await _scalar(session, select(func.count()).select_from(UserModel).where(*created_conditions)) + active_users = await _scalar(session, select(func.count(distinct(Session.user_id))).select_from(Session).where(*active_conditions)) + labels = {"service": SERVICE_NAME, "window": window_name} + lines.append(_sample("auth_user_registrations_total", new_users, labels)) + lines.append(_sample("auth_active_unique_users_total", active_users, labels)) + lines.append(_sample("auth_metrics_scrape_success", 1, {"service": SERVICE_NAME})) return lines -def _collect_http_metrics() -> Iterable[str]: - yield '# TYPE auth_http_requests_total counter' - for (method, route, status_code), value in sorted(_http_requests.items()): - yield _sample( - 'auth_http_requests_total', - value, - { - 'service': SERVICE_NAME, - 'method': method, - 'route': route, - 'status_code': status_code, - }, - ) +_business_metrics_cache = CachedBusinessMetrics(BUSINESS_CACHE_TTL_SECONDS, _collect_user_metrics) - yield '# TYPE auth_http_request_duration_seconds histogram' - route_keys = sorted(_http_latency_count) - for method, route in route_keys: - cumulative = 0 - for bucket in LATENCY_BUCKETS: - cumulative = _http_latency_buckets.get((method, route, bucket), cumulative) - yield _sample( - 'auth_http_request_duration_seconds_bucket', - cumulative, - { - 'service': SERVICE_NAME, - 'method': method, - 'route': route, - 'le': str(bucket), - }, - ) - yield _sample( - 'auth_http_request_duration_seconds_bucket', - _http_latency_count[(method, route)], - { - 'service': SERVICE_NAME, - 'method': method, - 'route': route, - 'le': '+Inf', - }, - ) - yield _sample( - 'auth_http_request_duration_seconds_sum', - _http_latency_sum[(method, route)], - {'service': SERVICE_NAME, 'method': method, 'route': route}, - ) - yield _sample( - 'auth_http_request_duration_seconds_count', - _http_latency_count[(method, route)], - {'service': SERVICE_NAME, 'method': method, 'route': route}, - ) + +def _client_metrics() -> bytes: + if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + return generate_latest(registry) + return generate_latest() async def metrics_response() -> Response: - lines = [ - '# HELP auth_users_total Total registered users in auth database.', - '# TYPE auth_users_total gauge', - '# HELP auth_user_registrations_total Registered users counted from the database for the selected rolling window.', - '# TYPE auth_user_registrations_total gauge', - '# HELP auth_active_unique_users_total Unique users with non-revoked sessions seen in the selected rolling window.', - '# TYPE auth_active_unique_users_total gauge', - '# HELP auth_sessions_active_total Non-revoked auth sessions.', - '# TYPE auth_sessions_active_total gauge', - ] - + lines: list[str] = [] try: - lines.extend(await _collect_user_metrics()) - lines.append(_sample('auth_metrics_scrape_success', 1, {'service': SERVICE_NAME})) + lines.extend(await _business_metrics_cache.get()) except Exception: - lines.append(_sample('auth_metrics_scrape_success', 0, {'service': SERVICE_NAME})) - - lines.extend(_collect_http_metrics()) - return Response( - content='\n'.join(lines) + '\n', - media_type='text/plain; version=0.0.4; charset=utf-8', - ) + lines.append(_sample("auth_metrics_scrape_success", 0, {"service": SERVICE_NAME})) + body = _client_metrics().decode("utf-8") + "\n".join(lines) + "\n" + return Response(body, media_type=CONTENT_TYPE_LATEST) diff --git a/uv.lock b/uv.lock index a323c4c..7ef814a 100644 --- a/uv.lock +++ b/uv.lock @@ -289,6 +289,7 @@ dependencies = [ { name = "hvac" }, { name = "itsdangerous" }, { name = "orjson" }, + { name = "prometheus-client" }, { name = "pydantic-settings" }, { name = "python-jose" }, { name = "python-ulid" }, @@ -312,6 +313,7 @@ requires-dist = [ { name = "hvac", specifier = "==2.4.0" }, { name = "itsdangerous", specifier = "==2.2.0" }, { name = "orjson", specifier = "==3.11.7" }, + { name = "prometheus-client", specifier = "==0.23.1" }, { name = "pydantic-settings", specifier = "==2.12.0" }, { name = "python-jose", specifier = "==3.5.0" }, { name = "python-ulid", specifier = "==3.1.0" }, @@ -593,7 +595,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f9/c8/9d76a66421d1ae24340dfae7e79c313957f6e3195c144d2c73333b5bfe34/greenlet-3.3.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:7e806ca53acf6d15a888405880766ec84721aa4181261cd11a457dfe9a7a4975", size = 276443, upload-time = "2026-01-23T15:30:10.066Z" }, { url = "https://files.pythonhosted.org/packages/81/99/401ff34bb3c032d1f10477d199724f5e5f6fbfb59816ad1455c79c1eb8e7/greenlet-3.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d842c94b9155f1c9b3058036c24ffb8ff78b428414a19792b2380be9cecf4f36", size = 597359, upload-time = "2026-01-23T16:00:57.394Z" }, { url = "https://files.pythonhosted.org/packages/2b/bc/4dcc0871ed557792d304f50be0f7487a14e017952ec689effe2180a6ff35/greenlet-3.3.1-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:20fedaadd422fa02695f82093f9a98bad3dab5fcda793c658b945fcde2ab27ba", size = 607805, upload-time = "2026-01-23T16:05:28.068Z" }, - { url = "https://files.pythonhosted.org/packages/3b/cd/7a7ca57588dac3389e97f7c9521cb6641fd8b6602faf1eaa4188384757df/greenlet-3.3.1-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c620051669fd04ac6b60ebc70478210119c56e2d5d5df848baec4312e260e4ca", size = 622363, upload-time = "2026-01-23T16:15:54.754Z" }, { url = "https://files.pythonhosted.org/packages/cf/05/821587cf19e2ce1f2b24945d890b164401e5085f9d09cbd969b0c193cd20/greenlet-3.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14194f5f4305800ff329cbf02c5fcc88f01886cadd29941b807668a45f0d2336", size = 609947, upload-time = "2026-01-23T15:32:51.004Z" }, { url = "https://files.pythonhosted.org/packages/a4/52/ee8c46ed9f8babaa93a19e577f26e3d28a519feac6350ed6f25f1afee7e9/greenlet-3.3.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7b2fe4150a0cf59f847a67db8c155ac36aed89080a6a639e9f16df5d6c6096f1", size = 1567487, upload-time = "2026-01-23T16:04:22.125Z" }, { url = "https://files.pythonhosted.org/packages/8f/7c/456a74f07029597626f3a6db71b273a3632aecb9afafeeca452cfa633197/greenlet-3.3.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:49f4ad195d45f4a66a0eb9c1ba4832bb380570d361912fa3554746830d332149", size = 1636087, upload-time = "2026-01-23T15:33:47.486Z" }, @@ -839,6 +840,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/94/448f037fb0ffd0e8a63b625cf9f5b13494b88d15573a987be8aaa735579d/progressbar2-4.5.0-py3-none-any.whl", hash = "sha256:625c94a54e63915b3959355e6d4aacd63a00219e5f3e2b12181b76867bf6f628", size = 57132, upload-time = "2024-08-28T22:50:10.264Z" }, ] +[[package]] +name = "prometheus-client" +version = "0.23.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/23/53/3edb5d68ecf6b38fcbcc1ad28391117d2a322d9a1a3eff04bfdb184d8c3b/prometheus_client-0.23.1.tar.gz", hash = "sha256:6ae8f9081eaaaf153a2e959d2e6c4f4fb57b12ef76c8c7980202f1e57b48b2ce", size = 80481, upload-time = "2025-09-18T20:47:25.043Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/db/14bafcb4af2139e046d03fd00dea7873e48eafe18b7d2797e73d6681f210/prometheus_client-0.23.1-py3-none-any.whl", hash = "sha256:dd1913e6e76b59cfe44e7a4b83e01afc9873c1bdfd2ed8739f1e76aeca115f99", size = 61145, upload-time = "2025-09-18T20:47:23.875Z" }, +] + [[package]] name = "propcache" version = "0.4.1"