fix: enable prometheus multiprocessing feature (#357)

Author: Panaetius

bases/renku_data_services/data_api/main.py

@@ -7,7 +7,6 @@
 
 import sentry_sdk
 import uvloop
-from prometheus_sanic import monitor
 from sanic import Sanic
 from sanic.log import logger
 from sanic.worker.loader import AppLoader
@@ -17,6 +16,7 @@
 from renku_data_services.app_config import Config
 from renku_data_services.authz.admin_sync import sync_admins_from_keycloak
 from renku_data_services.data_api.app import register_all_handlers
+from renku_data_services.data_api.prometheus import collect_system_metrics, setup_app_metrics, setup_prometheus
 from renku_data_services.errors.errors import (
     ForbiddenError,
     MissingResourceError,
@@ -31,11 +31,14 @@
     import sentry_sdk._types
 
 
-async def _send_messages() -> None:
+async def _send_messages(app: Sanic) -> None:
     config = Config.from_env()
     while True:
         try:
             await config.event_repo.send_pending_events()
+            # we need to collect metrics for this background process separately from the task we add to the
+            # server processes
+            await collect_system_metrics(app, "send_events_worker")
             await asyncio.sleep(1)
         except (asyncio.CancelledError, KeyboardInterrupt) as e:
             logger.warning(f"Exiting: {e}")
@@ -45,14 +48,15 @@ async def _send_messages() -> None:
             raise
 
 
-def send_pending_events() -> None:
+def send_pending_events(app_name: str) -> None:
     """Send pending messages in case sending in a handler failed."""
-    _ = Sanic("send_events")  # we need a dummy app for logging to work.
+    app = Sanic(app_name)
+    setup_app_metrics(app)
 
     logger.info("running events sending loop.")
 
     asyncio.set_event_loop(uvloop.new_event_loop())
-    asyncio.run(_send_messages())
+    asyncio.run(_send_messages(app))
 
 
 def create_app() -> Sanic:
@@ -104,9 +108,7 @@ async def setup_sentry(_: Sanic) -> None:
     logger.info(f"REAL_IP_HEADER = {app.config.REAL_IP_HEADER}")
 
     app = register_all_handlers(app, config)
-
-    # Setup prometheus
-    monitor(app, endpoint_type="url", multiprocess_mode="all", is_middleware=True).expose_endpoint()
+    setup_prometheus(app)
 
     if environ.get("CORS_ALLOW_ALL_ORIGINS", "false").lower() == "true":
         from sanic_ext import Extend
@@ -134,7 +136,7 @@ async def setup_rclone_validator(app: Sanic) -> None:
     async def ready(app: Sanic) -> None:
         """Application ready event handler."""
         logger.info("starting events background job.")
-        app.manager.manage("SendEvents", send_pending_events, {}, transient=True)
+        app.manager.manage("SendEvents", send_pending_events, {"app_name": config.app_name}, transient=True)
 
     return app
 

bases/renku_data_services/data_api/prometheus.py

@@ -0,0 +1,73 @@
+"""Prometheus Metrics."""
+
+import asyncio
+import resource
+
+import aiofiles
+from prometheus_client import Gauge
+from prometheus_sanic import monitor
+from prometheus_sanic.constants import BaseMetrics
+from prometheus_sanic.metrics import init
+from sanic import Sanic
+
+_PAGESIZE = resource.getpagesize()
+PROMETHEUS_VIRTUAL_MEMORY = "sanic_process_virtual_memory_bytes"
+PROMETHEUS_RESIDENT_MEMORY = "sanic_process_resident_memory_bytes"
+PROMETHEUS_METRICS_LIST = [
+    (
+        PROMETHEUS_VIRTUAL_MEMORY,
+        Gauge(PROMETHEUS_VIRTUAL_MEMORY, "Virtual memory size in bytes.", ["worker"]),
+    ),
+    (
+        PROMETHEUS_RESIDENT_MEMORY,
+        Gauge(PROMETHEUS_RESIDENT_MEMORY, "Resident memory size in bytes.", ["worker"]),
+    ),
+]
+
+
+async def collect_system_metrics(app: Sanic, name: str) -> None:
+    """Collect prometheus system metrics in a background task.
+
+    This is similar to the official prometheus_client implementation, which doesn't support CPU/Mem metrics
+    in multiprocess mode
+    """
+    try:
+        async with aiofiles.open("/proc/self/stat", "rb") as stat:
+            content = await stat.read()
+            parts = content.split(b")")[-1].split()
+        app.ctx.metrics[PROMETHEUS_VIRTUAL_MEMORY].labels({name}).set(float(parts[20]))
+        app.ctx.metrics[PROMETHEUS_RESIDENT_MEMORY].labels({name}).set(float(parts[21]) * _PAGESIZE)
+    except OSError:
+        pass
+
+
+async def collect_system_metrics_task(app: Sanic) -> None:
+    """Background task to collect metrics."""
+    while True:
+        await collect_system_metrics(app, app.m.name)
+        await asyncio.sleep(5)
+
+
+def setup_prometheus(app: Sanic) -> None:
+    """Setup prometheus monitoring.
+
+    We add custom metrics collection wo sanic workers and to the send_messages background job, since
+    prometheus does not collect cpy/memory metrics when in multiprocess mode.
+    """
+    app.add_task(collect_system_metrics_task)  # type:ignore[arg-type]
+    monitor(
+        app,
+        endpoint_type="url",
+        multiprocess_mode="all",
+        is_middleware=True,
+        metrics_list=PROMETHEUS_METRICS_LIST,
+    ).expose_endpoint()
+
+
+def setup_app_metrics(app: Sanic) -> None:
+    """Setup metrics for a Sanic app.
+
+    NOTE: this should only be called for manually created workers (with app.manager.manage(...))
+    """
+    app.ctx.metrics = {}
+    init(app, metrics_list=PROMETHEUS_METRICS_LIST, metrics=BaseMetrics)