Use rust to improve reading .db file performance

Signed-off-by: Chris Marchbanks <[email protected]>

Author: csmarchbanks

prometheus_client/multiprocess.py

@@ -10,6 +10,13 @@
 from .samples import Sample
 from .utils import floatToGoString
 
+_speedups = False
+try:
+    import prometheus_client_python_speedups
+    _speedups = True
+except:
+    pass
+
 try:  # Python3
     FileNotFoundError
 except NameError:  # Python >= 2.5
@@ -155,6 +162,14 @@ def _accumulate_metrics(metrics, accumulate):
 
     def collect(self):
         files = glob.glob(os.path.join(self._path, '*.db'))
+        if _speedups:
+            metrics = prometheus_client_python_speedups.merge(files)
+            native_metrics = []
+            for metric in metrics:
+                native_metric = Metric(metric.name, metric.documentation, metric.typ)
+                native_metric.samples = [Sample(sample.name, sample.labels, sample.value) for sample in metric.samples]
+                native_metrics.append(native_metric)
+            return native_metrics
         return self.merge(files, accumulate=True)
 
 

tests/test_multiprocess.py

@@ -5,10 +5,13 @@
 import unittest
 import warnings
 
+import pytest
+
 from prometheus_client import mmap_dict, values
 from prometheus_client.core import (
     CollectorRegistry, Counter, Gauge, Histogram, Sample, Summary,
 )
+import prometheus_client.multiprocess as multiprocess
 from prometheus_client.multiprocess import (
     mark_process_dead, MultiProcessCollector,
 )
@@ -398,6 +401,56 @@ def test_remove_clear_warning(self):
             assert "Clearing labels has not been implemented" in str(w[-1].message)
 
 
+@pytest.fixture
+def tempdir():
+    tempdir = tempfile.mkdtemp()
+    os.environ['PROMETHEUS_MULTIPROC_DIR'] = tempdir
+    values.ValueClass = MultiProcessValue(lambda: 123)
+    yield tempdir
+    del os.environ['PROMETHEUS_MULTIPROC_DIR']
+    shutil.rmtree(tempdir)
+    values.ValueClass = MutexValue
+
+@pytest.fixture
+def registry() -> CollectorRegistry:
+    return CollectorRegistry()
+
+@pytest.fixture
+def collector(tempdir, registry) -> MultiProcessCollector:
+    return MultiProcessCollector(registry)
+
+@pytest.fixture
+def no_speedup():
+    tmp = multiprocess._speedups
+    multiprocess._speedups = False
+    yield tmp
+    multiprocess._speedups = tmp
+
+def setup_benchmark():
+    labels = {i: i for i in 'abcd'}
+    for pid in range(1000):
+        values.ValueClass = MultiProcessValue(lambda: pid)
+
+        c = Counter('c', 'help', labelnames=labels.keys(), registry=None)
+        g = Gauge('g', 'help', labelnames=labels.keys(), registry=None)
+        h = Histogram('h', 'help', labelnames=labels.keys(), registry=None)
+
+        c.labels(**labels).inc(1)
+        g.labels(**labels).set(1)
+        h.labels(**labels).observe(1)
+
+
+def test_native_collect_performance(benchmark, collector, no_speedup):
+    setup_benchmark()
+    benchmark(collector.collect)
+
+def test_speedup_collect_performance(benchmark, collector):
+    if not multiprocess._speedups:
+        pytest.skip("prometheus_client_python_speedups not installed")
+    setup_benchmark()
+    benchmark(collector.collect)
+
+
 class TestMmapedDict(unittest.TestCase):
     def setUp(self):
         fd, self.tempfile = tempfile.mkstemp()

tox.ini

@@ -5,11 +5,13 @@ envlist = coverage-clean,py{3.9,3.10,3.11,3.12,3.13,py3.9,3.9-nooptionals},cover
 deps =
     coverage
     pytest
+	pytest-benchmark
     attrs
     {py3.9,pypy3.9}: twisted
     # NOTE: Pinned due to https://github.com/prometheus/client_python/issues/1020
     py3.9: asgiref==3.7
     pypy3.9: asgiref==3.7
+	py3.13: prometheus_client_python_speedups==0.1.0
 commands = coverage run --parallel -m pytest {posargs}
 
 [testenv:py3.9-nooptionals]