1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
# --------------------------------------------------------------------------------------
# Copyright (c) 2023-2024, Nucleic Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# --------------------------------------------------------------------------------------
import gc
import os
import pickle
import sys
import time
import tracemalloc
from multiprocessing import Process
import pytest
from atom.api import Atom, DefaultDict, Dict, Int, List, Set, atomref
try:
import psutil
PSUTIL_UNAVAILABLE = False
except ImportError:
PSUTIL_UNAVAILABLE = True
TIMEOUT = 6
class DictObj(Atom):
data = Dict(default={"a": 0})
class DefaultDictObj(Atom):
data = DefaultDict(value=Int(), default={1: 1})
class ListObj(Atom):
data = List(default=[1, 2, 3])
class SetObj(Atom):
data = Set(default={1, 2, 3})
class RefObj(Atom):
data = Int()
MEM_TESTS = {
"dict": DictObj,
"defaultdict": DefaultDictObj,
"list": ListObj,
"set": SetObj,
"atomref": RefObj,
}
PICKLE_MEM_TESTS = {
"dict": DictObj,
"defaultdict": DefaultDictObj,
"list": ListObj,
"set": SetObj,
}
def memtest(cls):
# Create object in a loop
# Memory usage should settle out and not change
while True:
obj = cls()
obj.data # Force creation
del obj
gc.collect()
def atomreftest(cls):
obj = cls()
obj.data
while True:
ref = atomref(obj)
del ref
gc.collect()
@pytest.mark.skipif(
"CI" in os.environ and sys.platform.startswith("darwin"),
reason="Flaky on MacOS CI runners",
)
@pytest.mark.skipif(sys.platform == "gnu0", reason="Flaky on Debian GNU/Hurd")
@pytest.mark.skipif(PSUTIL_UNAVAILABLE, reason="psutil is not installed")
@pytest.mark.parametrize("label", MEM_TESTS.keys())
def test_mem_usage(label):
TestClass = MEM_TESTS[label]
if "atomref" in label:
target = atomreftest
else:
target = memtest
p = Process(target=target, args=(TestClass,))
p.start()
try:
stats = psutil.Process(p.pid)
time.sleep(TIMEOUT * 1 / 4)
first_info = stats.memory_info()
time.sleep(TIMEOUT * 3 / 4)
last_info = stats.memory_info()
# Allow slight memory decrease over time to make tests more resilient
if first_info != last_info:
assert first_info.rss >= last_info.rss >= 0, (
"Memory leaked:\n {}\n {}".format(first_info, last_info)
)
assert first_info.vms >= last_info.vms >= 0, (
"Memory leaked:\n {}\n {}".format(first_info, last_info)
)
finally:
p.kill()
p.join()
@pytest.mark.parametrize("label", PICKLE_MEM_TESTS.keys())
def test_pickle_mem_usage(label):
TestClass = PICKLE_MEM_TESTS[label]
obj = TestClass()
for _ in range(100):
pickle.loads(pickle.dumps(obj))
gc.collect()
tracemalloc.start()
for i in range(10000):
pck = pickle.dumps(obj)
pickle.loads(pck)
del pck
gc.collect()
for stat in (
tracemalloc.take_snapshot()
.filter_traces(
[
tracemalloc.Filter(True, "*/atom/*"),
tracemalloc.Filter(False, "*/tests/*"),
]
)
.statistics("lineno")
):
# not sure why I sometimes see a 2 here but the original buggy version
# reported values > 50
assert stat.count < 5
|
