From 6b37e07ad4359151cb72bd23c444473ad6b63a25 Mon Sep 17 00:00:00 2001
From: Chris Marchbanks <csmarchbanks@gmail.com>
Date: Fri, 24 Jan 2025 13:29:08 -0700
Subject: [PATCH] Use rust to improve reading .db file performance

Signed-off-by: Chris Marchbanks <csmarchbanks@gmail.com>
---
 mypy.ini                          |  3 ++
 prometheus_client/multiprocess.py | 15 ++++++++
 tests/test_multiprocess.py        | 58 +++++++++++++++++++++++++++++++
 tox.ini                           |  2 ++
 4 files changed, 78 insertions(+)

diff --git a/mypy.ini b/mypy.ini
index fe372d07..983b1760 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -5,3 +5,6 @@ disallow_incomplete_defs = True
 
 [mypy-prometheus_client.decorator]
 follow_imports = skip
+
+[mypy-prometheus_client_python_speedups]
+ignore_missing_imports = True
diff --git a/prometheus_client/multiprocess.py b/prometheus_client/multiprocess.py
index 2682190a..e9685c74 100644
--- a/prometheus_client/multiprocess.py
+++ b/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)
 
 
diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py
index 77fd3d81..e626b0a7 100644
--- a/tests/test_multiprocess.py
+++ b/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,61 @@ 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):
+    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()
diff --git a/tox.ini b/tox.ini
index 157a8bb2..924b8e68 100644
--- a/tox.ini
+++ b/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]