diff --git a/prometheus_client/mmap_dict.py b/prometheus_client/mmap_dict.py index 274b7389..95ee3c62 100644 --- a/prometheus_client/mmap_dict.py +++ b/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] diff --git a/prometheus_client/pidprovider.py b/prometheus_client/pidprovider.py new file mode 100644 index 00000000..b41bc3f1 --- /dev/null +++ b/prometheus_client/pidprovider.py @@ -0,0 +1,9 @@ +import os + + +class Pidprovider(object): + source = os.getpid + + @classmethod + def getpid(cls): + return cls.source() diff --git a/prometheus_client/values.py b/prometheus_client/values.py index 2831665a..2d7c26cb 100644 --- a/prometheus_client/values.py +++ b/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()} diff --git a/tests/test_pidprovider.py b/tests/test_pidprovider.py new file mode 100644 index 00000000..4296691a --- /dev/null +++ b/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'])