Feature: user defined pidprovider

prometheus_client/mmap_dict.py

@@ -37,7 +37,7 @@ def _read_all_values(data, used=0):
         if encoded_len + pos > used:
             raise RuntimeError('Read beyond file size detected, file is corrupted.')
         pos += 4
-        encoded_key = data[pos : pos + encoded_len]
+        encoded_key = data[pos: pos + encoded_len]
         padded_len = encoded_len + (8 - (encoded_len + 4) % 8)
         pos += padded_len
         value = _unpack_double(data, pos)[0]

prometheus_client/pidprovider.py

@@ -0,0 +1,9 @@
+import os
+
+
+class Pidprovider(object):
+    source = os.getpid
+
+    @classmethod
+    def getpid(cls):
+        return cls.source()

prometheus_client/values.py

@@ -4,6 +4,7 @@
 from threading import Lock
 
 from .mmap_dict import mmap_key, MmapedDict
+from .pidprovider import Pidprovider
 
 
 class MutexValue(object):
@@ -28,7 +29,7 @@ def get(self):
             return self._value
 
 
-def MultiProcessValue(_pidFunc=os.getpid):
+def MultiProcessValue(_pidFunc=Pidprovider.getpid):
     files = {}
     values = []
     pid = {'value': _pidFunc()}

tests/test_pidprovider.py

@@ -0,0 +1,45 @@
+import os
+import shutil
+import sys
+import tempfile
+
+from prometheus_client import values
+from prometheus_client.core import Counter
+from prometheus_client.values import MultiProcessValue, Pidprovider
+
+if sys.version_info < (2, 7):
+    # We need the skip decorators from unittest2 on Python 2.6.
+    import unittest2 as unittest
+else:
+    import unittest
+
+
+class TestPidprovider(unittest.TestCase):
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        os.environ['prometheus_multiproc_dir'] = self.tempdir
+        self.originalValueClass = values.ValueClass
+        self.originalPidproviderSource = Pidprovider.source
+        values.ValueClass = MultiProcessValue()
+
+    def tearDown(self):
+        del os.environ['prometheus_multiproc_dir']
+        shutil.rmtree(self.tempdir)
+        values.ValueClass = self.originalValueClass
+        Pidprovider.source = self.originalPidproviderSource
+
+    # can not inspect the files cache directly, as it's a closure, so we
+    # check for the actual files themselves
+    def _files(self):
+        fs = os.listdir(self.tempdir)
+        fs.sort()
+        return fs
+
+    def test_with_default_pidprovider(self):
+        Counter('c1', 'c1', registry=None)
+        self.assertEqual(self._files(), ['counter_{0}.db'.format(os.getpid())])
+
+    def test_with_user_defined_pidprovider(self):
+        Pidprovider.source = staticmethod(lambda: 1234)
+        Counter('c1', 'c1', registry=None)
+        self.assertEqual(self._files(), ['counter_1234.db'])