Switch multi-process support to custom data store.

brian-brazil · brian-brazil · commit f85ccd910fe3 · 2016-10-10T16:23:59.000+01:00
This uses an mmaped file to store the data.
Fsync is not called, and concurrent readers are permitted (which
some shelve backends did not permit).

Microbenchmarks indicate 1.2-1.5us per instrumentation event.
This compares to 3-13ms from the shelve solution, so 3 orders
of magnitude better.

The 1.2-1.5us is about double normal multi-process instrumentation, due
to the 2nd mutex required. Python doesn't offer an RWLock which would
let this be reduced.

Overall this performance is acceptable, as if someone was doing
something extremely performance critical with instrumentation they
wouldn't be using Python.
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -3,11 +3,12 @@
 from __future__ import unicode_literals
 
 import copy
-import math
 import json
+import math
+import mmap
 import os
 import re
-import shelve
+import struct
 import types
 from timeit import default_timer
 
@@ -26,6 +27,7 @@
 _RESERVED_METRIC_LABEL_NAME_RE = re.compile(r'^__.*$')
 _INF = float("inf")
 _MINUS_INF = float("-inf")
+_INITIAL_MMAP_SIZE = 1024*1024
 
 
 class CollectorRegistry(object):
@@ -241,50 +243,129 @@ def get(self):
       with self._lock:
           return self._value
 
+class _MmapedDict(object):
+    """A dict of doubles, backed by an mmapped file.
+
+    The file starts with a 4 byte int, indicating how much of it is used.
+    Then 4 bytes of padding.
+    There's then a number of entries, consisting of a 4 byte int which is the
+    side of the next field, a utf-8 encoded string key, padding to a 8 byte
+    alignment, and then a 8 byte float which is the value.
+    """
+    def __init__(self, filename):
+        self._lock = Lock()
+        self._f = open(filename, 'a+b')
+        if os.fstat(self._f.fileno()).st_size == 0:
+            self._f.truncate(_INITIAL_MMAP_SIZE)
+        self._capacity = os.fstat(self._f.fileno()).st_size
+        self._m = mmap.mmap(self._f.fileno(), self._capacity)
+
+        self._positions = {}
+        self._used = struct.unpack_from(b'i', self._m, 0)[0]
+        if self._used == 0:
+            self._used = 8
+            struct.pack_into(b'i', self._m, 0, self._used)
+        else:
+            for key, _, pos in self._read_all_values():
+                self._positions[key] = pos
+
+    def _init_value(self, key):
+        """Initilize a value. Lock must be held by caller."""
+        encoded = key.encode('utf-8')
+        # Pad to be 8-byte aligned.
+        padded = encoded + (b' ' * (8 - (len(encoded) + 4) % 8))
+        value = struct.pack('i{0}sd'.format(len(padded)).encode(), len(encoded), padded, 0.0)
+        while self._used + len(value) > self._capacity:
+            self._capacity *= 2
+            self._f.truncate(self._capacity * 2)
+            self._m = mmap.mmap(self._f.fileno(), self._capacity)
+        self._m[self._used:self._used + len(value)] = value
+
+        # Update how much space we've used.
+        self._used += len(value)
+        struct.pack_into(b'i', self._m, 0, self._used)
+        self._positions[key] = self._used - 8
+
+    def _read_all_values(self):
+        """Yield (key, value, pos). No locking is performed."""
+        pos = 8
+        while pos < self._used:
+            encoded_len = struct.unpack_from(b'i', self._m, pos)[0]
+            pos += 4
+            encoded = struct.unpack_from('{0}s'.format(encoded_len).encode(), self._m, pos)[0]
+            padded_len = encoded_len + (8 - (encoded_len + 4) % 8)
+            pos += padded_len
+            value = struct.unpack_from(b'd', self._m, pos)[0]
+            yield encoded.decode('utf-8'), value, pos
+            pos += 8
+
+    def read_all_values(self):
+        """Yield (key, value, pos). No locking is performed."""
+        for k, v, _ in self._read_all_values():
+            yield k, v
+
+    def read_value(self, key):
+        with self._lock:
+            if key not in self._positions:
+                self._init_value(key)
+        pos = self._positions[key]
+        # We assume that reading from an 8 byte aligned value is atomic
+        return struct.unpack_from(b'd', self._m, pos)[0]
+
+    def write_value(self, key, value):
+        with self._lock:
+            if key not in self._positions:
+                self._init_value(key)
+        pos = self._positions[key]
+        # We assume that writing to an 8 byte aligned value is atomic
+        struct.pack_into(b'd', self._m, pos, value)
+
+    def close(self):
+        if self._f:
+            self._f.close()
+            self._f = None
+
 
 def _MultiProcessValue(__pid=os.getpid()):
     pid = __pid
-    samples = {}
-    samples_lock = Lock()
+    files = {}
+    files_lock = Lock()
 
-    class _ShelveValue(object):
-        '''A float protected by a mutex backed by a per-process shelve.'''
+    class _MmapedValue(object):
+        '''A float protected by a mutex backed by a per-process mmaped file.'''
 
         _multiprocess = True
 
         def __init__(self, typ, metric_name, name, labelnames, labelvalues, multiprocess_mode='', **kwargs):
-            with samples_lock:
-                if typ == 'gauge':
-                    file_prefix = typ + '_' +  multiprocess_mode
-                else:
-                    file_prefix = typ
-                if file_prefix not in samples:
-                    filename = os.path.join(os.environ['prometheus_multiproc_dir'], '{0}_{1}.db'.format(file_prefix, pid))
-                    samples[file_prefix] = shelve.open(filename)
-            self._samples = samples[file_prefix]
+            if typ == 'gauge':
+                file_prefix = typ + '_' +  multiprocess_mode
+            else:
+                file_prefix = typ
+            with files_lock:
+                if file_prefix not in files:
+                    filename = os.path.join(
+                            os.environ['prometheus_multiproc_dir'], '{0}_{1}.db'.format(file_prefix, pid))
+                    files[file_prefix] = _MmapedDict(filename)
+            self._file = files[file_prefix]
             self._key = json.dumps((metric_name, name, labelnames, labelvalues))
-            self._value = self._samples.get(self._key, 0.0)
-            self._samples[self._key] = self._value
-            self._samples.sync()
+            self._value = self._file.read_value(self._key)
             self._lock = Lock()
 
         def inc(self, amount):
             with self._lock:
                 self._value += amount
-                self._samples[self._key] = self._value
-                self._samples.sync()
+                self._file.write_value(self._key, self._value)
 
         def set(self, value):
             with self._lock:
                 self._value = value
-                self._samples[self._key] = self._value
-                self._samples.sync()
+                self._file.write_value(self._key, self._value)
 
         def get(self):
             with self._lock:
                 return self._value
 
-    return _ShelveValue
+    return _MmapedValue
 
 
 # Should we enable multi-process mode?
diff --git a/prometheus_client/multiprocess.py b/prometheus_client/multiprocess.py
@@ -21,7 +21,8 @@ def collect(self):
         for f in glob.glob(os.path.join(self._path, '*.db')):
             parts = os.path.basename(f).split('_')
             typ = parts[0]
-            for key, value in shelve.open(f).items():
+            d = core._MmapedDict(f)
+            for key, value in d.read_all_values():
                 metric_name, name, labelnames, labelvalues = json.loads(key)
                 metrics.setdefault(metric_name, core.Metric(metric_name, 'Multiprocess metric', typ))
                 metric = metrics[metric_name]
@@ -30,8 +31,9 @@ def collect(self):
                     metric._multiprocess_mode = parts[1]
                     metric.add_sample(name, tuple(zip(labelnames, labelvalues)) + (('pid', pid), ), value)
                 else:
-                    # The deplucates and labels are fixed in the next for.
+                    # The duplicates and labels are fixed in the next for.
                     metric.add_sample(name, tuple(zip(labelnames, labelvalues)), value)
+            d.close()
 
         for metric in metrics.values():
             samples = {}
diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py
@@ -109,3 +109,32 @@ def test_gauge_livesum(self):
         self.assertEqual(3, self.registry.get_sample_value('g'))
         mark_process_dead(123, os.environ['prometheus_multiproc_dir'])
         self.assertEqual(2, self.registry.get_sample_value('g'))
+
+class TestMmapedDict(unittest.TestCase):
+    def setUp(self):
+        fd, self.tempfile = tempfile.mkstemp()
+        os.close(fd)
+        self.d = core._MmapedDict(self.tempfile)
+
+    def test_process_restart(self):
+        self.d.write_value('abc', 123.0)
+        self.d.close()
+        self.d = core._MmapedDict(self.tempfile)
+        self.assertEqual(123, self.d.read_value('abc'))
+        self.assertEqual([('abc', 123.0)], list(self.d.read_all_values()))
+
+    def test_expansion(self):
+        key = 'a' * core._INITIAL_MMAP_SIZE
+        self.d.write_value(key, 123.0)
+        self.assertEqual([(key, 123.0)], list(self.d.read_all_values()))
+
+    def test_multi_expansion(self):
+        key = 'a' * core._INITIAL_MMAP_SIZE * 4
+        self.d.write_value('abc', 42.0)
+        self.d.write_value(key, 123.0)
+        self.d.write_value('def', 17.0)
+        self.assertEqual([('abc', 42.0), (key, 123.0), ('def', 17.0)],
+                list(self.d.read_all_values()))
+
+    def tearDown(self):
+        os.unlink(self.tempfile)
