8000 Fix bug of multiprocess mode by `path` argument · prometheus/client_python@ec84762 · GitHub
[go: up one dir, main page]

Skip to content

Commit ec84762

Browse files
committed
Fix bug of multiprocess mode by path argument
Signed-off-by: Vladimír Kašpar <vladimir.kaspar@heureka.cz>
1 parent 2206f15 commit ec84762

File tree

3 files changed

+37
-13
lines changed

3 files changed

+37
-13
lines changed

prometheus_client/multiprocess.py

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,23 @@
1313
from .utils import floatToGoString
1414

1515

16+
PATH = None
17+
18+
1619
class MultiProcessCollector(object):
1720
"""Collector for files for multi-process mode."""
1821

1922
def __init__(self, registry, path=None):
20-
if path is None:
21-
path = os.environ.get('prometheus_multiproc_dir')
22-
if not path or not os.path.isdir(path):
23-
raise ValueError('env prometheus_multiproc_dir is not set or not a directory')
24-
self._path = path
23+
global PATH
24+
25+
if path is not None:
26+
PATH = path
27+
elif 'prometheus_multiproc_dir' in os.environ:
28+
PATH = os.environ['prometheus_multiproc_dir']
29+
30+
if not PATH or not os.path.isdir(PATH):
31+
raise ValueError('Any prometheus multiprocess dictionary is not set or not a directory.')
32+
2533
if registry:
2634
registry.register(self)
2735

@@ -114,7 +122,7 @@ def merge(files, accumulate=True):
114122
return metrics.values()
115123

116124
def collect(self):
117-
files = glob.glob(os.path.join(self._path, '*.db'))
125+
files = glob.glob(os.path.join(PATH, '*.db'))
118126
return self.merge(files, accumulate=True)
119127

120128

prometheus_client/values.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from threading import Lock
55

66
from .mmap_dict import mmap_key, MmapedDict
7+
from . import multiprocess
78

89

910
class MutexValue(object):
@@ -57,7 +58,7 @@ def __reset(self):
5758
file_prefix = typ
5859
if file_prefix not in files:
5960
filename = os.path.join(
60-
os.environ['prometheus_multiproc_dir'],
61+
multiprocess.PATH,
6162
'{0}_{1}.db'.format(file_prefix, pid['value']))
6263

6364
files[file_prefix] = MmapedDict(filename)
@@ -101,7 +102,7 @@ def get_value_class():
101102
# This needs to be chosen before the first metric is constructed,
102103
# and as that may be in some arbitrary library the user/admin has
103104
# no control over we use an environment variable.
104-
if 'prometheus_multiproc_dir' in os.environ:
105+
if multiprocess.PATH:
105106
return MultiProcessValue()
106107
else:
107108
return MutexValue

tests/test_multiprocess.py

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ def test_gauge_all(self):
8080
self.assertEqual(0, self.registry.get_sample_value('g', {'pid': '456'}))
8181
g1.set(1)
8282
g2.set(2)
83-
mark_process_dead(123, os.environ['prometheus_multiproc_dir'])
83+
mark_process_dead(123, self.tempdir)
8484
self.assertEqual(1, self.registry.get_sample_value('g', {'pid': '123'}))
8585
self.assertEqual(2, self.registry.get_sample_value('g', {'pid': '456'}))
8686

@@ -94,7 +94,7 @@ def test_gauge_liveall(self):
9494
g2.set(2)
9595
self.assertEqual(1, self.registry.get_sample_value('g', {'pid': '123'}))
9696
self.assertEqual(2, self.registry.get_sample_value('g', {'pid': '456'}))
97-
mark_process_dead(123, os.environ['prometheus_multiproc_dir'])
97+
mark_process_dead(123, self.tempdir)
9898
self.assertEqual(None, self.registry.get_sample_value('g', {'pid': '123'}))
9999
self.assertEqual(2, self.registry.get_sample_value('g', {'pid': '456'}))
100100

@@ -124,7 +124,7 @@ def test_gauge_livesum(self):
124124
g1.set(1)
125125
g2.set(2)
126126
self.assertEqual(3, self.registry.get_sample_value('g'))
127-
mark_process_dead(123, os.environ['prometheus_multiproc_dir'])
127+
mark_process_dead(123, self.tempdir)
128128
self.assertEqual(2, self.registry.get_sample_value('g'))
129129

130130
def test_namespace_subsystem(self):
@@ -151,7 +151,7 @@ def test_initialization_detects_pid_change(self):
151151
# can not inspect the files cache directly, as it's a closure, so we
152152
# check for the actual files themselves
153153
def files():
154-
fs = os.listdir(os.environ['prometheus_multiproc_dir'])
154+
fs = os.listdir(self.tempdir)
155155
fs.sort()
156156
return fs
157157

@@ -240,7 +240,7 @@ def add_label(key, value):
240240
pid = 1
241241
h.labels(**labels).observe(5)
242242

243-
path = os.path.join(os.environ['prometheus_multiproc_dir'], '*.db')
243+
path = os.path.join(self.tempdir, '*.db')
244244
files = glob.glob(path)
245245
metrics = dict(
246246
(m.name, m) for m in self.collector.merge(files, accumulate=False)
@@ -271,6 +271,21 @@ def add_label(key, value):
271271
self.assertEqual(metrics['h'].samples, expected_histogram)
272272

273273

274+
class TestMultiProcessByPathArgument(TestMultiProcess):
275+
def setUp(self):
276+
if 'prometheus_multiproc_dir' in os.environ:
277+
del os.environ['prometheus_multiproc_dir']
278+
279+
self.tempdir = tempfile.mkdtemp()
280+
values.ValueClass = MultiProcessValue(lambda: 123)
281+
self.registry = CollectorRegistry()
282+
self.collector = MultiProcessCollector(self.registry, path=self.tempdir)
283+
284+
def tearDown(self):
285+
shutil.rmtree(self.tempdir)
286+
values.ValueClass = MutexValue
287+
288+
274289
class TestMmapedDict(unittest.TestCase):
275290
def setUp(self):
276291
fd, self.tempfile = tempfile.mkstemp()

0 commit comments

Comments
 (0)
0