From 7d87f077318eb849ac6e6eeb71daf132ef1fb3a8 Mon Sep 17 00:00:00 2001 From: isra17 Date: Wed, 26 Jul 2023 17:38:13 -0400 Subject: [PATCH] Fix memory leak from Cancellable helper --- src/saturn_engine/utils/asyncutils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/saturn_engine/utils/asyncutils.py b/src/saturn_engine/utils/asyncutils.py index 3772027e..2601e199 100644 --- a/src/saturn_engine/utils/asyncutils.py +++ b/src/saturn_engine/utils/asyncutils.py @@ -441,4 +441,7 @@ async def __call__(self, *args: t.Any, **kwargs: t.Any) -> R: self.func(*args, **kwargs), name=f"cancellable({self.func})" ) self._tasks.add(task) - return await task + try: + return await task + finally: + self._tasks.remove(task)