Repository: prometheus/client_python
Pull request #302 (2018-09-05 11:48:04)

Refactor MultiProcessCollector.collect() to allow for arbitrary merging.
Factors out a merge() method from the previous collect() method, which
is parameterized, and thus can be used for arbitrary merging of samples.
For motivation, see discussion in issue #275 around merging dead worker's
data into a single mmaped file.

This basically allows us to parameterize the files to be merged, and
also whether to accumulate histograms or not. Accumulation is on by
default, as that is what the prometheus format expects. But it can now
be disabled, which allows merged values to be correctly written back to
an mmaped file. For the same reason, the order of labels is preserved
via OrderedDict.

brian-brazil: Can you resolve the conflict? I just merged a large reworking of the library to support openmetrics.
bloodearnest: Rebased on latest master

Patch:
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
index da3034d..cb1e7c5 100644
--- a/prometheus_client/core.py
+++ b/prometheus_client/core.py
@@ -185,7 +185,7 @@ class CollectorRegistry(object):
 REGISTRY = CollectorRegistry(auto_describe=True)
 '''The default registry.'''
 
-_METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram', 
+_METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram',
         'gaugehistogram', 'unknown', 'info', 'stateset')
 
 
@@ -378,8 +378,8 @@ class HistogramMetricFamily(Metric):
             exemplar = None
             if len(b) == 3:
                 exemplar = b[2]
-            self.samples.append(Sample(self.name + '_bucket', 
-                dict(list(zip(self._labelnames, labels)) + [('le', bucket)]), 
+            self.samples.append(Sample(self.name + '_bucket',
+                dict(list(zip(self._labelnames, labels)) + [('le', bucket)]),
                 value, timestamp, exemplar))
         # +Inf is last and provides the count value.
         self.samples.append(Sample(self.name + '_count', dict(zip(self._labelnames, labels)), buckets[-1][1], timestamp))
@@ -411,7 +411,7 @@ class GaugeHistogramMetricFamily(Metric):
         '''
         for bucket, value in buckets:
             self.samples.append(Sample(
-                self.name + '_bucket', 
+                self.name + '_bucket',
                 dict(list(zip(self._labelnames, labels)) + [('le', bucket)]),
                 value, timestamp))
 
@@ -438,7 +438,7 @@ class InfoMetricFamily(Metric):
           labels: A list of label values
           value: A dict of labels
         '''
-        self.samples.append(Sample(self.name + '_info', 
+        self.samples.append(Sample(self.name + '_info',
             dict(dict(zip(self._labelnames, labels)), **value), 1, timestamp))
 
 
@@ -586,6 +586,13 @@ class _MmapedDict(object):
             self._f = None
 
 
+def _mmap_key(metric_name, name, labelnames, labelvalues):
+    """Format a key for use in the mmap file."""
+    # ensure labels are in consistent order for identity
+    labels = dict(zip(labelnames, labelvalues))
+    return json.dumps([metric_name, name, labels], sort_keys=True)
+
+
 def _MultiProcessValue(_pidFunc=os.getpid):
     files = {}
     values = []
@@ -618,7 +625,7 @@ def _MultiProcessValue(_pidFunc=os.getpid):
                     '{0}_{1}.db'.format(file_prefix, pid['value']))
                 files[file_prefix] = _MmapedDict(filename)
             self._file = files[file_prefix]
-            self._key = json.dumps((metric_name, name, labelnames, labelvalues))
+            self._key = _mmap_key(metric_name, name, labelnames, labelvalues)
             self._value = self._file.read_value(self._key)
 
         def __check_for_pid_change(self):
@@ -1143,7 +1150,7 @@ class Enum(object):
      Example usage:
         from prometheus_client import Enum
 
-        e = Enum('task_state', 'Description of enum', 
+        e = Enum('task_state', 'Description of enum',
           states=['starting', 'running', 'stopped'])
         e.state('running')
 
diff --git a/prometheus_client/multiprocess.py b/prometheus_client/multiprocess.py
index 7dd74b2..5521315 100644
--- a/prometheus_client/multiprocess.py
+++ b/prometheus_client/multiprocess.py
@@ -23,13 +23,24 @@ class MultiProcessCollector(object):
             registry.register(self)
 
     def collect(self):
+        files = glob.glob(os.path.join(self._path, '*.db'))
+        return self.merge(files, accumulate=True)
+
+    def merge(self, files, accumulate=True):
+        """Merge metrics from given mmap files.
+
+        By default, histograms are accumulated, as per prometheus wire format.
+        But if writing the merged data back to mmap files, use
+        accumulate=False to avoid compound accumulation.
+        """
         metrics = {}
-        for f in glob.glob(os.path.join(self._path, '*.db')):
+        for f in files:
             parts = os.path.basename(f).split('_')
             typ = parts[0]
             d = core._MmapedDict(f, read_mode=True)
             for key, value in d.read_all_values():
-                metric_name, name, labelnames, labelvalues = json.loads(key)
+                metric_name, name, labels = json.loads(key)
+                labels_key = tuple(sorted(labels.items()))
 
                 metric = metrics.get(metric_name)
                 if metric is None:
@@ -39,10 +50,10 @@ class MultiProcessCollector(object):
                 if typ == 'gauge':
                     pid = parts[2][:-3]
                     metric._multiprocess_mode = parts[1]
-                    metric.add_sample(name, tuple(zip(labelnames, labelvalues)) + (('pid', pid), ), value)
+                    metric.add_sample(name, labels_key + (('pid', pid), ), value)
                 else:
                     # The duplicates and labels are fixed in the next for.
-                    metric.add_sample(name, tuple(zip(labelnames, labelvalues)), value)
+                    metric.add_sample(name, labels_key, value)
             d.close()
 
         for metric in metrics.values():
@@ -86,9 +97,17 @@ class MultiProcessCollector(object):
                 for labels, values in buckets.items():
                     acc = 0.0
                     for bucket, value in sorted(values.items()):
-                        acc += value
-                        samples[(metric.name + '_bucket', labels + (('le', core._floatToGoString(bucket)), ))] = acc
-                    samples[(metric.name + '_count', labels)] = acc
+                        sample_key = (
+                            metric.name + '_bucket',
+                            labels + (('le', core._floatToGoString(bucket)), ),
+                        )
+                        if accumulate:
+                            acc += value
+                            samples[sample_key] = acc
+                        else:
+                            samples[sample_key] = value
+                    if accumulate:
+                        samples[(metric.name + '_count', labels)] = acc
 
             # Convert to correct sample format.
             metric.samples = [core.Sample(name, dict(labels), value) for (name, labels), value in samples.items()]


Test patch:
diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py
index ca84913..501cb62 100644
--- a/tests/test_multiprocess.py
+++ b/tests/test_multiprocess.py
@@ -1,9 +1,17 @@
 from __future__ import unicode_literals
 
+import glob
 import os
 import shutil
+import sys
 import tempfile
-import unittest
+
+if sys.version_info < (2, 7):
+    # We need the skip decorators from unittest2 on Python 2.6.
+    import unittest2 as unittest
+else:
+    import unittest
+
 
 from prometheus_client import core
 from prometheus_client.core import (
@@ -11,6 +19,7 @@ from prometheus_client.core import (
     Counter,
     Gauge,
     Histogram,
+    Sample,
     Summary,
 )
 from prometheus_client.multiprocess import (
@@ -25,7 +34,7 @@ class TestMultiProcess(unittest.TestCase):
         os.environ['prometheus_multiproc_dir'] = self.tempdir
         core._ValueClass = core._MultiProcessValue(lambda: 123)
         self.registry = CollectorRegistry()
-        MultiProcessCollector(self.registry, self.tempdir)
+        self.collector = MultiProcessCollector(self.registry, self.tempdir)
 
     def tearDown(self):
         del os.environ['prometheus_multiproc_dir']
@@ -137,6 +146,113 @@ class TestMultiProcess(unittest.TestCase):
         self.assertEqual(3, self.registry.get_sample_value('c_total'))
         self.assertEqual(1, c1._value.get())
 
+    @unittest.skipIf(sys.version_info < (2, 7), "Test requires Python 2.7+.")
+    def test_collect(self):
+        pid = 0
+        core._ValueClass = core._MultiProcessValue(lambda: pid)
+        labels = dict((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)
+        h = Histogram('h', 'help', labelnames=labels.keys(), registry=None)
+
+        c.labels(**labels).inc(1)
+        g.labels(**labels).set(1)
+        h.labels(**labels).observe(1)
+
+        pid = 1
+
+        c.labels(**labels).inc(1)
+        g.labels(**labels).set(1)
+        h.labels(**labels).observe(5)
+
+        metrics = dict((m.name, m) for m in self.collector.collect())
+
+        self.assertEqual(
+            metrics['c'].samples, [Sample('c_total', labels, 2.0)]
+        )
+        metrics['g'].samples.sort(key=lambda x: x[1]['pid'])
+        self.assertEqual(metrics['g'].samples, [
+            Sample('g', add_label('pid', '0'), 1.0),
+            Sample('g', add_label('pid', '1'), 1.0),
+        ])
+
+        metrics['h'].samples.sort(
+            key=lambda x: (x[0], float(x[1].get('le', 0)))
+        )
+        expected_histogram = [
+            Sample('h_bucket', add_label('le', '0.005'), 0.0),
+            Sample('h_bucket', add_label('le', '0.01'), 0.0),
+            Sample('h_bucket', add_label('le', '0.025'), 0.0),
+            Sample('h_bucket', add_label('le', '0.05'), 0.0),
+            Sample('h_bucket', add_label('le', '0.075'), 0.0),
+            Sample('h_bucket', add_label('le', '0.1'), 0.0),
+            Sample('h_bucket', add_label('le', '0.25'), 0.0),
+            Sample('h_bucket', add_label('le', '0.5'), 0.0),
+            Sample('h_bucket', add_label('le', '0.75'), 0.0),
+            Sample('h_bucket', add_label('le', '1.0'), 1.0),
+            Sample('h_bucket', add_label('le', '2.5'), 1.0),
+            Sample('h_bucket', add_label('le', '5.0'), 2.0),
+            Sample('h_bucket', add_label('le', '7.5'), 2.0),
+            Sample('h_bucket', add_label('le', '10.0'), 2.0),
+            Sample('h_bucket', add_label('le', '+Inf'), 2.0),
+            Sample('h_count', labels, 2.0),
+            Sample('h_sum', labels, 6.0),
+        ]
+
+        self.assertEqual(metrics['h'].samples, expected_histogram)
+
+    @unittest.skipIf(sys.version_info < (2, 7), "Test requires Python 2.7+.")
+    def test_merge_no_accumulate(self):
+        pid = 0
+        core._ValueClass = core._MultiProcessValue(lambda: pid)
+        labels = dict((i, i) for i in 'abcd')
+
+        def add_label(key, value):
+            l = labels.copy()
+            l[key] = value
+            return l
+
+        h = Histogram('h', 'help', labelnames=labels.keys(), registry=None)
+        h.labels(**labels).observe(1)
+        pid = 1
+        h.labels(**labels).observe(5)
+
+        path = os.path.join(os.environ['prometheus_multiproc_dir'], '*.db')
+        files = glob.glob(path)
+        metrics = dict(
+            (m.name, m) for m in self.collector.merge(files, accumulate=False)
+        )
+
+        metrics['h'].samples.sort(
+            key=lambda x: (x[0], float(x[1].get('le', 0)))
+        )
+        expected_histogram = [
+            Sample('h_bucket', add_label('le', '0.005'), 0.0),
+            Sample('h_bucket', add_label('le', '0.01'), 0.0),
+            Sample('h_bucket', add_label('le', '0.025'), 0.0),
+            Sample('h_bucket', add_label('le', '0.05'), 0.0),
+            Sample('h_bucket', add_label('le', '0.075'), 0.0),
+            Sample('h_bucket', add_label('le', '0.1'), 0.0),
+            Sample('h_bucket', add_label('le', '0.25'), 0.0),
+            Sample('h_bucket', add_label('le', '0.5'), 0.0),
+            Sample('h_bucket', add_label('le', '0.75'), 0.0),
+            Sample('h_bucket', add_label('le', '1.0'), 1.0),
+            Sample('h_bucket', add_label('le', '2.5'), 0.0),
+            Sample('h_bucket', add_label('le', '5.0'), 1.0),
+            Sample('h_bucket', add_label('le', '7.5'), 0.0),
+            Sample('h_bucket', add_label('le', '10.0'), 0.0),
+            Sample('h_bucket', add_label('le', '+Inf'), 0.0),
+            Sample('h_sum', labels, 6.0),
+        ]
+
+        self.assertEqual(metrics['h'].samples, expected_histogram)
+
 
 class TestMmapedDict(unittest.TestCase):
     def setUp(self):