[#674, #675] Mitigate performance regression by filtering restricted registry collections

prometheus_client/registry.py

@@ -135,10 +135,21 @@ def __init__(self, names, registry):
         self._registry = registry
 
     def collect(self):
-        for metric in self._registry.collect():
-            m = metric._restricted_metric(self._name_set)
-            if m:
-                yield m
+        collectors = set()
+        target_info_metric = None
+        with self._registry._lock:
+            if 'target_info' in self._name_set and self._registry._target_info:
+                target_info_metric = self._registry._target_info_metric()
+            for name in self._name_set:
+                if name != 'target_info' and name in self._registry._names_to_collectors:
+                    collectors.add(self._registry._names_to_collectors[name])
+        if target_info_metric:
+            yield target_info_metric
+        for collector in collectors:
+            for metric in collector.collect():
+                m = metric._restricted_metric(self._name_set)
+                if m:
+                    yield m
 
 
 REGISTRY = CollectorRegistry(auto_describe=True)

tests/test_core.py

@@ -1,6 +1,7 @@
 from __future__ import unicode_literals
 
 from concurrent.futures import ThreadPoolExecutor
+import sys
 import time
 
 import pytest
@@ -838,6 +839,33 @@ def test_target_info_restricted_registry(self):
         m.samples = [Sample('target_info', {'foo': 'bar'}, 1)]
         self.assertEqual([m], list(registry.restricted_registry(['target_info']).collect()))
 
+    @unittest.skipIf(sys.version_info < (3, 3), "Test requires Python 3.3+.")
+    def test_restricted_registry_does_not_call_extra(self):
+        from unittest.mock import MagicMock
+        registry = CollectorRegistry()
+        mock_collector = MagicMock()
+        mock_collector.describe.return_value = [Metric('foo', 'help', 'summary')]
+        registry.register(mock_collector)
+        Summary('s', 'help', registry=registry).observe(7)
+
+        m = Metric('s', 'help', 'summary')
+        m.samples = [Sample('s_sum', {}, 7)]
+        self.assertEqual([m], list(registry.restricted_registry(['s_sum']).collect()))
+        mock_collector.collect.assert_not_called()
+
+    def test_restricted_registry_does_not_yield_while_locked(self):
+        registry = CollectorRegistry(target_info={'foo': 'bar'})
+        Summary('s', 'help', registry=registry).observe(7)
+
+        m = Metric('s', 'help', 'summary')
+        m.samples = [Sample('s_sum', {}, 7)]
+        self.assertEqual([m], list(registry.restricted_registry(['s_sum']).collect()))
+
+        m = Metric('target', 'Target metadata', 'info')
+        m.samples = [Sample('target_info', {'foo': 'bar'}, 1)]
+        for _ in registry.restricted_registry(['target_info', 's_sum']).collect():
+            self.assertFalse(registry._lock.locked())
+
 
 if __name__ == '__main__':
     unittest.main()