From 9aa8decf403c618092b00db0f8b76c9917987b6b Mon Sep 17 00:00:00 2001
From: Kevin Morris <kevr@0cost.org>
Date: Mon, 1 Nov 2021 14:18:19 -0700
Subject: [PATCH] fix(fastapi): use metrics in cases where
 PROMETHEUS_MULTIPROC_DIR is defined

Previously, we restricted this to gunicorn to get it working on aur-dev.
This change makes it usable through any backend, and also no-op if
PROMETHEUS_MULTIPROC_DIR is not defined.

Signed-off-by: Kevin Morris <kevr@0cost.org>
---
 aurweb/routers/html.py        | 4 +---
 docker-compose.yml            | 3 +++
 docker/scripts/run-pytests.sh | 3 +++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/aurweb/routers/html.py b/aurweb/routers/html.py
index 4cee5f99..525fb626 100644
--- a/aurweb/routers/html.py
+++ b/aurweb/routers/html.py
@@ -209,9 +209,7 @@ async def index(request: Request):
 @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.
+    if os.environ.get("PROMETHEUS_MULTIPROC_DIR", None):  # pragma: no cover
         multiprocess.MultiProcessCollector(registry)
     data = generate_latest(registry)
     headers = {
diff --git a/docker-compose.yml b/docker-compose.yml
index 5dffe5d3..225e5b9b 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -241,6 +241,7 @@ services:
     environment:
       - AUR_CONFIG=conf/config
       - TEST_RECURSION_LIMIT=${TEST_RECURSION_LIMIT}
+      - PROMETHEUS_MULTIPROC_DIR=/tmp_prometheus
     entrypoint: /docker/test-mysql-entrypoint.sh
     command: /docker/scripts/run-pytests.sh clean
     stdin_open: true
@@ -267,6 +268,7 @@ services:
     environment:
       - AUR_CONFIG=conf/config.sqlite
       - TEST_RECURSION_LIMIT=${TEST_RECURSION_LIMIT}
+      - PROMETHEUS_MULTIPROC_DIR=/tmp_prometheus
     entrypoint: /docker/test-sqlite-entrypoint.sh
     command: setup-sqlite.sh run-pytests.sh clean
     stdin_open: true
@@ -289,6 +291,7 @@ services:
     environment:
       - AUR_CONFIG=conf/config
       - TEST_RECURSION_LIMIT=${TEST_RECURSION_LIMIT}
+      - PROMETHEUS_MULTIPROC_DIR=/tmp_prometheus
     entrypoint: /docker/tests-entrypoint.sh
     command: setup-sqlite.sh run-tests.sh
     stdin_open: true
diff --git a/docker/scripts/run-pytests.sh b/docker/scripts/run-pytests.sh
index d992bf06..ee546fb7 100755
--- a/docker/scripts/run-pytests.sh
+++ b/docker/scripts/run-pytests.sh
@@ -22,6 +22,9 @@ while [ $# -ne 0 ]; do
     esac
 done
 
+rm -rf $PROMETHEUS_MULTIPROC_DIR
+mkdir -p $PROMETHEUS_MULTIPROC_DIR
+
 # Initialize the new database; ignore errors.
 python -m aurweb.initdb 2>/dev/null || \
     (echo "Error: aurweb.initdb failed; already initialized?" && /bin/true)