Repository: prometheus/client_python
Pull request #680 (2021-07-01 18:35:27)

[#674, #675] Mitigate performance regression by filtering restricted registry collections
#675 introduces a significant performance regression in restricted registry collecting. This attempts to reduce the regression by re-introducing filter code removed in the previous version.

pavel-lexyr: py2.7 passes locally. What do the logs say on the CI?
csmarchbanks: ```
==================================== ERRORS ====================================
_____________________ ERROR collecting tests/test_core.py ______________________
ImportError while importing test module '/home/circleci/project/tests/test_core.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
tests/test_core.py:5: in <module>
    from unittest.mock import MagicMock
E   ImportError: No module named mock
```

You can use a skipIf and import only for the test such as done in https://github.com/prometheus/client_python/blob/master/tests/openmetrics/test_parser.py#L492.
pavel-lexyr: Aha. That worked - thank you!
brian-brazil: :+1: 

Patch:
diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py
index 0eb003b..67cde45 100644
--- a/prometheus_client/registry.py
+++ b/prometheus_client/registry.py
@@ -135,10 +135,21 @@ class RestrictedRegistry(object):
         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)


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index d36aef0..4ccf93f 100644
--- a/tests/test_core.py
+++ b/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 @@ class TestCollectorRegistry(unittest.TestCase):
         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()