Support MultiProcessCollector in RestrictedRegistry.

This change makes it so that the RestrictedRegistry will always attempt to collect metrics from a collector for which it couldn’t find any metrics name. Although this can be used generally, this is meant to be used with MultiProcessCollector.

This changes the current behavior of the code but should be somehow safe as it enables filtering in case where it was not working previously. If this is an issue, an alternative approach with an explicit flag could be used (set either in the MultiProcessCollector or in the registry).

The intent here is to allow collecting a subset of metrics from production fastapi servers (running in multiprocess mode). So not having to change the library usage in these servers is advantageous to have filtering work out-of-the-box with this change.

Author: mathias-kende

docs/content/multiprocess/_index.md

@@ -10,9 +10,10 @@ it's common to have processes rather than threads to handle large workloads.
 To handle this the client library can be put in multiprocess mode.
 This comes with a number of limitations:
 
-- Registries can not be used as normal, all instantiated metrics are exported
+- Registries can not be used as normal
   - Registering metrics to a registry later used by a `MultiProcessCollector`
     may cause duplicate metrics to be exported
+  - Filtering on metrics works but is inefficient
 - Custom collectors do not work (e.g. cpu and memory metrics)
 - Gauges cannot use `set_function`
 - Info and Enum metrics do not work

prometheus_client/registry.py

@@ -33,6 +33,7 @@ def __init__(self, auto_describe: bool = False, target_info: Optional[Dict[str,
         self._lock = Lock()
         self._target_info: Optional[Dict[str, str]] = {}
         self.set_target_info(target_info)
+        self._collectors_with_no_names: List[Collector] = []
 
     def register(self, collector: Collector) -> None:
         """Add a collector to the registry."""
@@ -46,6 +47,8 @@ def register(self, collector: Collector) -> None:
             for name in names:
                 self._names_to_collectors[name] = collector
             self._collector_to_names[collector] = names
+            if not names:
+                self._collectors_with_no_names.append(collector)
 
     def unregister(self, collector: Collector) -> None:
         """Remove a collector from the registry."""
@@ -148,7 +151,7 @@ def __init__(self, names: Iterable[str], registry: CollectorRegistry):
         self._registry = registry
 
     def collect(self) -> Iterable[Metric]:
-        collectors = set()
+        collectors = set(self._registry._collectors_with_no_names)
         target_info_metric = None
         with self._registry._lock:
             if 'target_info' in self._name_set and self._registry._target_info:

tests/test_asgi.py

@@ -223,3 +223,32 @@ def test_qs_parsing(self):
             asyncio.new_event_loop().run_until_complete(
                 self.communicator.wait()
             )
+
+    def test_qs_parsing_multi(self):
+        """Only metrics that match the 'name[]' query string param appear"""
+
+        app = make_asgi_app(self.registry)
+        metrics = [
+            ("asdf", "first test metric", 1),
+            ("bsdf", "second test metric", 2),
+            ("csdf", "third test metric", 3)
+        ]
+
+        for m in metrics:
+            self.increment_metrics(*m)
+
+        self.seed_app(app)
+        self.scope['query_string'] = "&".join([f"name[]={m[0]}_total" for m in metrics[0:2]]).encode("utf-8")
+        self.send_default_request()
+
+        outputs = self.get_all_output()
+        response_body = outputs[1]
+        output = response_body['body'].decode('utf8')
+
+        self.assert_metrics(output, *metrics[0])
+        self.assert_metrics(output, *metrics[1])
+        self.assert_not_metrics(output, *metrics[2])
+
+        asyncio.new_event_loop().run_until_complete(
+            self.communicator.wait()
+        )

tests/test_core.py

@@ -1024,6 +1024,15 @@ def test_restricted_registry_does_not_call_extra(self):
         self.assertEqual([m], list(registry.restricted_registry(['s_sum']).collect()))
         mock_collector.collect.assert_not_called()
 
+    def test_restricted_registry_collects_no_names_collectors(self):
+        from unittest.mock import MagicMock
+        registry = CollectorRegistry()
+        mock_collector = MagicMock()
+        mock_collector.describe.return_value = []
+        registry.register(mock_collector)
+        self.assertEqual(list(registry.restricted_registry(['metric']).collect()), [])
+        mock_collector.collect.assert_called()
+
     def test_restricted_registry_does_not_yield_while_locked(self):
         registry = CollectorRegistry(target_info={'foo': 'bar'})
         Summary('s', 'help', registry=registry).observe(7)

tests/test_multiprocess.py

@@ -301,6 +301,35 @@ def add_label(key, value):
 
         self.assertEqual(metrics['h'].samples, expected_histogram)
 
+    def test_restrict(self):
+        pid = 0
+        values.ValueClass = MultiProcessValue(lambda: pid)
+        labels = {i: i for i in 'abcd'}
+
+        def add_label(key, value):
+            l = labels.copy()
+            l[key] = value
+            return l
+
+        c = Counter('c', 'help', labelnames=labels.keys(), registry=None)
+        g = Gauge('g', 'help', labelnames=labels.keys(), registry=None)
+
+        c.labels(**labels).inc(1)
+        g.labels(**labels).set(1)
+
+        pid = 1
+
+        c.labels(**labels).inc(1)
+        g.labels(**labels).set(1)
+
+        metrics = {m.name: m for m in self.registry.restricted_registry(['c_total']).collect()}
+
+        self.assertEqual(metrics.keys(), {'c'})
+
+        self.assertEqual(
+            metrics['c'].samples, [Sample('c_total', labels, 2.0)]
+        )
+
     def test_collect_preserves_help(self):
         pid = 0
         values.ValueClass = MultiProcessValue(lambda: pid)