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.

Author: riconnon

django_prometheus/exports.py

@@ -19,7 +19,16 @@
 logger = logging.getLogger(__name__)
 
 
-def SetupPrometheusEndpointOnPort(port, addr=""):
+def GetRegistry():
+    if "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
@@ -43,7 +52,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):
@@ -57,7 +66,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
@@ -83,8 +92,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
@@ -105,22 +116,19 @@ 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):
     """Exports /metrics as a Django view.
 
     You can use django_prometheus.urls to map /metrics to this view.
     """
-    if "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

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