From cfefbd069230617bad60034a83ed67dd23465d1f Mon Sep 17 00:00:00 2001
From: Richard Connon <richard@connon.me.uk>
Date: Wed, 18 Aug 2021 23:39:49 +0100
Subject: [PATCH] Fix use of dedicated listeners with multiprocess

Extract the code to select the registry used for the exporter from the
Django view and provide the registry as a parameter to clients for
deciated listeners as well.
---
 django_prometheus/exports.py            | 33 ++++++++++++++++---------
 django_prometheus/tests/test_exports.py |  6 +++--
 2 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/django_prometheus/exports.py b/django_prometheus/exports.py
index 94c04156..c1ddcdbd 100644
--- a/django_prometheus/exports.py
+++ b/django_prometheus/exports.py
@@ -18,7 +18,19 @@
 logger = logging.getLogger(__name__)
 
 
-def SetupPrometheusEndpointOnPort(port, addr=""):
+def GetRegistry():
+    if (
+        "PROMETHEUS_MULTIPROC_DIR" in os.environ
+        or "prometheus_multiproc_dir" in os.environ
+    ):
+        registry = prometheus_client.CollectorRegistry()
+        multiprocess.MultiProcessCollector(registry)
+    else:
+        registry = prometheus_client.REGISTRY
+    return registry
+
+
+def SetupPrometheusEndpointOnPort(registry, port, addr=""):
     """Exports Prometheus metrics on an HTTPServer running in its own thread.
 
     The server runs on the given port and is by default listenning on
@@ -42,7 +54,7 @@ def SetupPrometheusEndpointOnPort(port, addr=""):
         "autoreloader is active. Use the URL exporter, or start django "
         "with --noreload. See documentation/exports.md."
     )
-    prometheus_client.start_http_server(port, addr=addr)
+    prometheus_client.start_http_server(port, addr=addr, registry=registry)
 
 
 class PrometheusEndpointServer(threading.Thread):
@@ -56,7 +68,7 @@ def run(self):
         self.httpd.serve_forever()
 
 
-def SetupPrometheusEndpointOnPortRange(port_range, addr=""):
+def SetupPrometheusEndpointOnPortRange(registry, port_range, addr=""):
     """Like SetupPrometheusEndpointOnPort, but tries several ports.
 
     This is useful when you're running Django as a WSGI application
@@ -82,8 +94,10 @@ def SetupPrometheusEndpointOnPortRange(port_range, addr=""):
         "with --noreload. See documentation/exports.md."
     )
     for port in port_range:
+        handler = prometheus_client.MetricsHandler
+        handler.registry = registry
         try:
-            httpd = HTTPServer((addr, port), prometheus_client.MetricsHandler)
+            httpd = HTTPServer((addr, port), handler)
         except OSError:
             # Python 2 raises socket.error, in Python 3 socket.error is an
             # alias for OSError
@@ -102,10 +116,11 @@ def SetupPrometheusExportsFromConfig():
     port = getattr(settings, "PROMETHEUS_METRICS_EXPORT_PORT", None)
     port_range = getattr(settings, "PROMETHEUS_METRICS_EXPORT_PORT_RANGE", None)
     addr = getattr(settings, "PROMETHEUS_METRICS_EXPORT_ADDRESS", "")
+    registry = GetRegistry()
     if port_range:
-        SetupPrometheusEndpointOnPortRange(port_range, addr)
+        SetupPrometheusEndpointOnPortRange(registry, port_range, addr)
     elif port:
-        SetupPrometheusEndpointOnPort(port, addr)
+        SetupPrometheusEndpointOnPort(registry, port, addr)
 
 
 def ExportToDjangoView(request):
@@ -113,10 +128,6 @@ def ExportToDjangoView(request):
 
     You can use django_prometheus.urls to map /metrics to this view.
     """
-    if "PROMETHEUS_MULTIPROC_DIR" in os.environ or "prometheus_multiproc_dir" in os.environ:
-        registry = prometheus_client.CollectorRegistry()
-        multiprocess.MultiProcessCollector(registry)
-    else:
-        registry = prometheus_client.REGISTRY
+    registry = GetRegistry()
     metrics_page = prometheus_client.generate_latest(registry)
     return HttpResponse(metrics_page, content_type=prometheus_client.CONTENT_TYPE_LATEST)
diff --git a/django_prometheus/tests/test_exports.py b/django_prometheus/tests/test_exports.py
index 47d8abb0..76e20244 100644
--- a/django_prometheus/tests/test_exports.py
+++ b/django_prometheus/tests/test_exports.py
@@ -2,6 +2,8 @@
 import socket
 from unittest.mock import ANY, MagicMock, call, patch
 
+from prometheus_client import REGISTRY
+
 from django_prometheus.exports import SetupPrometheusEndpointOnPortRange
 
 
@@ -10,7 +12,7 @@ def test_port_range_available(httpserver_mock):
     """Test port range setup with an available port."""
     httpserver_mock.side_effect = [socket.error, MagicMock()]
     port_range = [8000, 8001]
-    port_chosen = SetupPrometheusEndpointOnPortRange(port_range)
+    port_chosen = SetupPrometheusEndpointOnPortRange(REGISTRY, port_range)
     assert port_chosen in port_range
 
     expected_calls = [call(("", 8000), ANY), call(("", 8001), ANY)]
@@ -22,7 +24,7 @@ def test_port_range_unavailable(httpserver_mock):
     """Test port range setup with no available ports."""
     httpserver_mock.side_effect = [socket.error, socket.error]
     port_range = [8000, 8001]
-    port_chosen = SetupPrometheusEndpointOnPortRange(port_range)
+    port_chosen = SetupPrometheusEndpointOnPortRange(REGISTRY, port_range)
 
     expected_calls = [call(("", 8000), ANY), call(("", 8001), ANY)]
     assert httpserver_mock.mock_calls == expected_calls