Skip to content

Commit 36b643e

Browse files
author
Brian Maissy
authored
Remove memory leak from storing completed tasks
To ensure that we don't end up with lots of "this task was deleted but never ran" warnings we hold onto the tasks we end up creating and ensure they are cancelled and ran in the loop before they are garbage collected. Prior to this we were doing that activity at the end of the run, which leads to large memory usage in large test suites This change makes it so that we perform this garbage collection as we proceed with the tests so that memory can be released during the run of the test suite
1 parent 602766d commit 36b643e

File tree

1 file changed

+28
-1
lines changed

1 file changed

+28
-1
lines changed

alt_pytest_asyncio/converter.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,33 @@ def __init__(self) -> None:
2121
defaultdict(list)
2222
)
2323

24+
def _cleanup_completed_tasks(self) -> None:
25+
"""
26+
Remove references to completed tasks to they can be garbage collected, including the
27+
return (or yielded) values of the fixture functions, which would otherwise leak memory.
28+
"""
29+
for loop, tasks in list(self._test_tasks.items()):
30+
if loop.is_closed():
31+
continue
32+
33+
remaining: list[asyncio.Task[object]] = []
34+
finished: list[asyncio.Task[object]] = []
35+
for t in tasks:
36+
if not t.done():
37+
remaining.append(t)
38+
else:
39+
t.cancel()
40+
finished.append(t)
41+
42+
self._test_tasks[loop] = remaining
43+
44+
if finished:
45+
loop.run_until_complete(asyncio.tasks.gather(*finished, return_exceptions=True))
46+
47+
def _add_new_task(self, loop: asyncio.AbstractEventLoop, task: asyncio.Task[object]) -> None:
48+
self._cleanup_completed_tasks()
49+
self._test_tasks[loop].append(task)
50+
2451
def sessionfinish(self) -> None:
2552
for loop, tasks in self._test_tasks.items():
2653
ts = []
@@ -227,7 +254,7 @@ def silent_done_task(res: asyncio.Future[protocols.T_Ret | None] | None) -> None
227254
self._async_runner(async_timeout, func, args, kwargs), context=self._ctx
228255
)
229256
task.add_done_callback(silent_done_task)
230-
self._test_tasks[loop].append(task)
257+
self._add_new_task(loop, task)
231258

232259
return loop.run_until_complete(task)
233260

0 commit comments

Comments
 (0)