fix(fastapi): utilize PROMETHEUS_MULTIPROC_DIR in our own /metrics

Signed-off-by: Kevin Morris <kevr@0cost.org>

aurweb/asgi.py

@@ -9,6 +9,7 @@ from urllib.parse import quote_plus
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
+from prometheus_client import multiprocess
 from sqlalchemy import and_, or_
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.sessions import SessionMiddleware
@@ -29,7 +30,7 @@ app = FastAPI(exception_handlers=errors.exceptions)
 # library with custom collectors and expose /metrics.
 instrumentator().add(http_api_requests_total())
 instrumentator().add(http_requests_total())
-instrumentator().instrument(app).expose(app)
+instrumentator().instrument(app)
 
 
 @app.on_event("startup")
@@ -79,6 +80,12 @@ async def app_startup():
     get_engine()
 
 
+def child_exit(server, worker):  # pragma: no cover
+    """ This function is required for gunicorn customization
+    of prometheus multiprocessing. """
+    multiprocess.mark_process_dead(worker.pid)
+
+
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):
     """

aurweb/routers/html.py

@@ -1,11 +1,14 @@
 """ AURWeb's primary routing module. Define all routes via @app.app.{get,post}
 decorators in some way; more complex routes should be defined in their
 own modules and imported here. """
+import os
+
 from datetime import datetime
 from http import HTTPStatus
 
-from fastapi import APIRouter, Form, HTTPException, Request
+from fastapi import APIRouter, Form, HTTPException, Request, Response
 from fastapi.responses import HTMLResponse, RedirectResponse
+from prometheus_client import CONTENT_TYPE_LATEST, CollectorRegistry, generate_latest, multiprocess
 from sqlalchemy import and_, case, or_
 
 import aurweb.config
@@ -203,7 +206,21 @@ async def index(request: Request):
     return render_template(request, "index.html", context)
 
 
-# A route that returns a error 503. For testing purposes.
+@router.get("/metrics")
+async def metrics(request: Request):
+    registry = CollectorRegistry()
+    if os.environ.get("FASTAPI_BACKEND", "") == "gunicorn":  # pragma: no cover
+        # This case only ever happens in production, when we are running
+        # gunicorn. We don't test with gunicorn, so we don't cover this path.
+        multiprocess.MultiProcessCollector(registry)
+    data = generate_latest(registry)
+    headers = {
+        "Content-Type": CONTENT_TYPE_LATEST,
+        "Content-Length": str(len(data))
+    }
+    return Response(data, headers=headers)
+
+
 @router.get("/raisefivethree", response_class=HTMLResponse)
 async def raise_service_unavailable(request: Request):
     raise HTTPException(status_code=503)

test/test_html.py

@@ -117,3 +117,10 @@ def test_get_successes():
 """
     successes = get_successes(html)
     assert successes[0].text.strip() == "Test"
+
+
+def test_metrics(client: TestClient):
+    with client as request:
+        resp = request.get("/metrics")
+    assert resp.status_code == int(HTTPStatus.OK)
+    assert resp.headers.get("Content-Type").startswith("text/plain")