diff --git a/prometheus_client/mmap_dict.py b/prometheus_client/mmap_dict.py index edd895cd..fb4a19eb 100644 --- a/prometheus_client/mmap_dict.py +++ b/prometheus_client/mmap_dict.py @@ -101,7 +101,8 @@ def _init_value(self, key): while self._used + len(value) > self._capacity: self._capacity *= 2 self._f.truncate(self._capacity) - self._m = mmap.mmap(self._f.fileno(), self._capacity) + self._m.close() + self._m = mmap.mmap(self._f.fileno(), self._capacity, access=mmap.ACCESS_WRITE) self._m[self._used:self._used + len(value)] = value # Update how much space we've used. @@ -131,12 +132,21 @@ def write_value(self, key, value, timestamp): _pack_two_doubles(self._m, pos, value, timestamp) def close(self): - if self._f: - self._m.close() + m = getattr(self, '_m', None) + if m is not None: + m.close() self._m = None - self._f.close() + f = getattr(self, '_f', None) + if f is not None: + f.close() self._f = None + def __del__(self): + try: + self.close() + except Exception: + return + def mmap_key(metric_name: str, name: str, labelnames: List[str], labelvalues: List[str], help_text: str) -> str: """Format a key for use in the mmap file.""" diff --git a/tests/test_mmap_dict_close.py b/tests/test_mmap_dict_close.py new file mode 100644 index 00000000..db5adcbe --- /dev/null +++ b/tests/test_mmap_dict_close.py @@ -0,0 +1,47 @@ +import gc +import os +import tempfile +import unittest +import warnings + +import prometheus_client.mmap_dict as mmap_dict +from prometheus_client.mmap_dict import MmapedDict + + +class TestMmapedDictClose(unittest.TestCase): + def test_discarded_dict_does_not_warn_about_an_open_file(self): + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "metrics.db") + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", ResourceWarning) + store = MmapedDict(path) + store.write_value("one", 1.0, 2.0) + del store + gc.collect() + self.assertEqual( + [w for w in caught if issubclass(w.category, ResourceWarning)], + [], + ) + + def test_growth_keeps_written_values(self): + mmap_dict._INITIAL_MMAP_SIZE = 128 + try: + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "metrics.db") + store = MmapedDict(path) + for i in range(40): + store.write_value("k" * 20 + str(i), float(i), 1.0) + store.close() + values = { + key: value + for key, value, _, _ in MmapedDict.read_all_values_from_file(path) + } + finally: + mmap_dict._INITIAL_MMAP_SIZE = 1 << 16 + self.assertEqual(values["k" * 20 + "0"], 0.0) + self.assertEqual(values["k" * 20 + "39"], 39.0) + self.assertEqual(len(values), 40) + + +if __name__ == "__main__": + unittest.main()