Fix memory leak from Cancellable helper

Flared · Jul 26, 2023 · 7d87f07 · 7d87f07
1 parent 7ec1ede
commit 7d87f07
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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)