SS-1322 monitor connection pool (#546)

Author: hamzaimran08

.env.template

@@ -57,6 +57,8 @@ APP_PROBE_STATUSES=Running,Deleted
 # Rate limiting configuration for the auth endpoint
 AUTH_RATE_LIMIT_VALUE=
 AUTH_RATE_LIMIT_WHITELIST=
+DB_POOL_STATS_LOGGING_ENABLED=false
+PROMETHEUS_METRICS_ENABLED=false
 
 # Caching/logging configuration.
 # SESSION_CACHE_ENABLED uses Django's cached_db session backend and requires Redis.

.pre-commit-config.yaml

@@ -52,4 +52,4 @@ repos:
   hooks:
   - id: mypy
     args: [--config-file, pyproject.toml]
-    additional_dependencies: [types-requests==2.25.9, pydantic==2.12.5]
+    additional_dependencies: [types-requests==2.25.9, pydantic==2.12.5, prometheus-client==0.24.1]

poetry.lock

@@ -37,6 +37,7 @@ wiki = "==0.11.2"
 
 # Other Python and project-related libraries
 flower = "^2.0.1"
+prometheus-client = "==0.24.1"
 requests = "==2.31.0"
 amqp = "==5.1.1"
 psycopg = {extras = ["binary", "pool", "c"], version = "^3.2.1"}

pyproject.toml

@@ -0,0 +1,156 @@
+import os
+from typing import Any
+
+from django.conf import settings
+from django.http import HttpRequest, HttpResponse
+from prometheus_client import (
+    CONTENT_TYPE_LATEST,
+    Counter,
+    Gauge,
+    Histogram,
+    generate_latest,
+)
+
+from studio.middleware import get_db_pool_stats
+
+DB_POOL_ENABLED = Gauge(
+    "django_db_pool_enabled",
+    "Whether Django database connection pooling is enabled.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_OPENED = Gauge(
+    "django_db_pool_opened",
+    "Whether the Django database connection pool has been opened in this process.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_SIZE = Gauge(
+    "django_db_pool_size",
+    "Current number of connections in the Django database connection pool.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_AVAILABLE = Gauge(
+    "django_db_pool_available",
+    "Current number of available connections in the Django database connection pool.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_MIN = Gauge(
+    "django_db_pool_min",
+    "Configured minimum number of connections in the Django database connection pool.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_MAX = Gauge(
+    "django_db_pool_max",
+    "Configured maximum number of connections in the Django database connection pool.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_REQUESTS_WAITING = Gauge(
+    "django_db_pool_requests_waiting",
+    "Current number of requests waiting for a Django database connection pool slot.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_REQUESTS_NUM = Gauge(
+    "django_db_pool_requests_total",
+    "Total number of requests served by the Django database connection pool in this process.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+DB_POOL_REQUESTS_ERRORS = Gauge(
+    "django_db_pool_requests_errors_total",
+    "Total number of Django database connection pool request errors in this process.",
+    ["alias", "pod", "pid", "pool_id"],
+)
+HTTP_REQUESTS_TOTAL = Counter(
+    "django_http_requests_total",
+    "Total number of Django HTTP requests handled by this process.",
+    ["method", "route", "status_code", "pod"],
+)
+HTTP_REQUEST_DURATION_SECONDS = Histogram(
+    "django_http_request_duration_seconds",
+    "Django HTTP request duration in seconds for this process.",
+    ["method", "route", "status_code", "pod"],
+)
+HTTP_REQUEST_EXCEPTIONS_TOTAL = Counter(
+    "django_http_request_exceptions_total",
+    "Total number of Django HTTP requests that raised an exception in this process.",
+    ["method", "route", "exception", "pod"],
+)
+
+
+def _metric_value(stats: dict[str, Any], key: str) -> float:
+    value = stats.get(key)
+    if isinstance(value, bool):
+        return 1.0 if value else 0.0
+    if isinstance(value, int | float):
+        return float(value)
+    return 0.0
+
+
+def update_db_pool_metrics() -> None:
+    stats = get_db_pool_stats()
+    labels = {
+        "alias": str(stats.get("alias", "")),
+        "pod": str(stats.get("pod", "")),
+        "pid": str(stats.get("pid", "")),
+        "pool_id": str(stats.get("pool_id", "")),
+    }
+
+    DB_POOL_ENABLED.labels(**labels).set(_metric_value(stats, "enabled"))
+    DB_POOL_OPENED.labels(**labels).set(_metric_value(stats, "opened"))
+    DB_POOL_SIZE.labels(**labels).set(_metric_value(stats, "pool_size"))
+    DB_POOL_AVAILABLE.labels(**labels).set(_metric_value(stats, "pool_available"))
+    DB_POOL_MIN.labels(**labels).set(_metric_value(stats, "pool_min"))
+    DB_POOL_MAX.labels(**labels).set(_metric_value(stats, "pool_max"))
+    DB_POOL_REQUESTS_WAITING.labels(**labels).set(_metric_value(stats, "requests_waiting"))
+    DB_POOL_REQUESTS_NUM.labels(**labels).set(_metric_value(stats, "requests_num"))
+    DB_POOL_REQUESTS_ERRORS.labels(**labels).set(_metric_value(stats, "requests_errors"))
+
+
+def get_http_route_label(request: HttpRequest) -> str:
+    resolver_match = getattr(request, "resolver_match", None)
+    if resolver_match is None:
+        return "unknown"
+    view_name = getattr(resolver_match, "view_name", "")
+    if isinstance(view_name, str) and view_name:
+        return view_name
+    route = getattr(resolver_match, "route", "")
+    if isinstance(route, str) and route:
+        return route
+    return "unknown"
+
+
+def _pod_label() -> str:
+    return os.environ.get("HOSTNAME", "")
+
+
+def record_http_request_metrics(request: HttpRequest, response: HttpResponse, duration_seconds: float) -> None:
+    labels = {
+        "method": request.method,
+        "route": get_http_route_label(request),
+        "status_code": str(response.status_code),
+        "pod": _pod_label(),
+    }
+    HTTP_REQUESTS_TOTAL.labels(**labels).inc()
+    HTTP_REQUEST_DURATION_SECONDS.labels(**labels).observe(duration_seconds)
+
+
+def record_http_exception_metrics(request: HttpRequest, exception: Exception, duration_seconds: float) -> None:
+    HTTP_REQUEST_EXCEPTIONS_TOTAL.labels(
+        method=request.method,
+        route=get_http_route_label(request),
+        exception=exception.__class__.__name__,
+        pod=_pod_label(),
+    ).inc()
+    labels = {
+        "method": request.method,
+        "route": get_http_route_label(request),
+        "status_code": "exception",
+        "pod": _pod_label(),
+    }
+    HTTP_REQUEST_DURATION_SECONDS.labels(**labels).observe(duration_seconds)
+
+
+def metrics_view(request: HttpRequest) -> HttpResponse:
+    if not settings.PROMETHEUS_METRICS_ENABLED:
+        return HttpResponse(status=404)
+
+    update_db_pool_metrics()
+    return HttpResponse(generate_latest(), content_type=CONTENT_TYPE_LATEST)

studio/metrics.py

@@ -1,13 +1,47 @@
-import logging
+import os
 import sys
+import time
 import traceback
 from typing import Any, Callable
 
+from django.conf import settings
+from django.db import connections
 from django.http import HttpRequest, HttpResponse
 
 from studio.utils import get_logger
 
 logger = get_logger(__name__)
+_last_db_pool_stats: dict[str, Any] | None = None
+DB_POOL_STATS_CHANGE_KEYS = (
+    "enabled",
+    "opened",
+    "pool_size",
+    "pool_available",
+    "requests_waiting",
+    "requests_errors",
+)
+METRICS_PATHS = ("/metrics", "/metrics/")
+
+
+def get_db_pool_stats(alias: str = "default") -> dict[str, Any]:
+    connection = connections[alias]
+    pool_enabled = bool(connection.settings_dict["OPTIONS"].get("pool"))
+    pools = getattr(connection, "_connection_pools", {})
+    pool = pools.get(alias)
+    pod_name = os.environ.get("HOSTNAME", "")
+    pid = os.getpid()
+
+    stats = {
+        "alias": alias,
+        "enabled": pool_enabled,
+        "opened": pool is not None,
+        "pid": pid,
+        "pod": pod_name,
+        "pool_id": f"{pod_name}:{pid}:{alias}",
+    }
+    if pool is not None:
+        stats.update(pool.get_stats())
+    return stats
 
 
 class ExceptionLoggingMiddleware:
@@ -33,3 +67,73 @@ def process_exception(self, request: HttpRequest, exception: Exception) -> None:
         msg += "".join(traceback.format_tb(stacktrace)).replace("\n", "\\n")
         logger.error(msg)
         return None
+
+
+class PrometheusHttpMetricsMiddleware:
+    """
+    Records Django HTTP request metrics for the Prometheus endpoint.
+    """
+
+    def __init__(self, get_response: Callable[[HttpRequest], Any]):
+        self.get_response = get_response
+
+    def __call__(self, request: HttpRequest) -> HttpResponse:
+        if not settings.PROMETHEUS_METRICS_ENABLED or request.path in METRICS_PATHS:
+            return self.get_response(request)
+
+        start = time.perf_counter()
+        try:
+            response = self.get_response(request)
+        except Exception as exception:
+            from studio.metrics import record_http_exception_metrics
+
+            record_http_exception_metrics(request, exception, time.perf_counter() - start)
+            raise
+
+        from studio.metrics import record_http_request_metrics
+
+        record_http_request_metrics(request, response, time.perf_counter() - start)
+        return response
+
+
+class DatabasePoolStatsLoggingMiddleware:
+    """
+    Logs psycopg pool stats from the Django worker process that handled the request.
+    """
+
+    def __init__(self, get_response: Callable[[HttpRequest], Any]):
+        self.get_response = get_response
+
+    def __call__(self, request: HttpRequest) -> HttpResponse:
+        response = self.get_response(request)
+
+        if settings.DB_POOL_STATS_LOGGING_ENABLED:
+            pool_stats = get_db_pool_stats()
+            if _db_pool_stats_changed(pool_stats):
+                logger.info(
+                    "Django DB pool stats path=%s status=%s pod=%s pid=%s pool_id=%s opened=%s "
+                    "pool_size=%s pool_available=%s requests_waiting=%s requests_num=%s requests_errors=%s stats=%s",
+                    request.path,
+                    response.status_code,
+                    pool_stats.get("pod"),
+                    pool_stats.get("pid"),
+                    pool_stats.get("pool_id"),
+                    pool_stats.get("opened"),
+                    pool_stats.get("pool_size"),
+                    pool_stats.get("pool_available"),
+                    pool_stats.get("requests_waiting"),
+                    pool_stats.get("requests_num"),
+                    pool_stats.get("requests_errors"),
+                    pool_stats,
+                )
+
+        return response
+
+
+def _db_pool_stats_changed(stats: dict[str, Any]) -> bool:
+    global _last_db_pool_stats
+
+    current_stats = {key: stats.get(key) for key in DB_POOL_STATS_CHANGE_KEYS}
+    previous_stats = _last_db_pool_stats
+    _last_db_pool_stats = current_stats
+    return previous_stats != current_stats

studio/middleware.py

@@ -122,6 +122,8 @@
         "corsheaders.middleware.CorsMiddleware",
         "axes.middleware.AxesMiddleware",
         "studio.middleware.ExceptionLoggingMiddleware",
+        "studio.middleware.PrometheusHttpMetricsMiddleware",
+        "studio.middleware.DatabasePoolStatsLoggingMiddleware",
         "django_htmx.middleware.HtmxMiddleware",
         "waffle.middleware.WaffleMiddleware",
     ]
@@ -378,6 +380,16 @@ def postgres_session_options() -> str:
 AUTH_RATE_LIMIT_VALUE = os.environ.get("AUTH_RATE_LIMIT_VALUE", None)
 AUTH_RATE_LIMIT_WHITELIST_RAW = os.environ.get("AUTH_RATE_LIMIT_WHITELIST", "")
 AUTH_RATE_LIMIT_WHITELIST = [ip.strip() for ip in AUTH_RATE_LIMIT_WHITELIST_RAW.split(",")]
+DB_POOL_STATS_LOGGING_ENABLED = os.environ.get("DB_POOL_STATS_LOGGING_ENABLED", "false").lower() in (
+    "1",
+    "true",
+    "yes",
+)
+PROMETHEUS_METRICS_ENABLED = os.environ.get("PROMETHEUS_METRICS_ENABLED", "false").lower() in (
+    "1",
+    "true",
+    "yes",
+)
 
 AUTH_PERMISSION_CACHE_ENABLED = os.environ.get("AUTH_PERMISSION_CACHE_ENABLED", "false").lower() in ("1", "true", "yes")
 AUTH_PERMISSION_CACHE_TIMEOUT = int(os.environ.get("AUTH_PERMISSION_CACHE_TIMEOUT", "30"))

studio/settings.py

@@ -20,10 +20,11 @@
 from django.conf.urls.static import static
 from django.contrib import admin
 from django.contrib.staticfiles.urls import staticfiles_urlpatterns
-from django.urls import include, path
+from django.urls import include, path, re_path
 from django_altcha import AltchaChallengeView
 
 from apps.views import record_lookup
+from studio.metrics import metrics_view
 
 from . import views
 
@@ -41,6 +42,7 @@
             name="delete_account_post_handler",
         ),
         path("status", views.status_view, name="status"),
+        re_path(r"^metrics/?$", metrics_view, name="metrics"),
         path("user/account-deleted/<int:user_id>", views.account_deleted, name="account_deleted"),
         path("auth/", views.AuthView.as_view()),
         # API paths using NamespaceVersioning