diff --git a/tests/test_tasks.py b/tests/test_tasks.py index aefe679..9de7157 100644 --- a/tests/test_tasks.py +++ b/tests/test_tasks.py @@ -13,10 +13,9 @@ async def test_wait(): t = Tasks('tasks') for _ in range(20): t.start(asyncio.sleep(.001), 'task') + len(t._tasks) == 20 await t.wait() - assert len(t._tasks) == 20 - for task in t._tasks: - assert task.done() + assert len(t._tasks) == 0 async def test_cancel(): diff --git a/walnats/__init__.py b/walnats/__init__.py index 1fe1676..49f4f44 100644 --- a/walnats/__init__.py +++ b/walnats/__init__.py @@ -7,7 +7,7 @@ from ._services import Services -__version__ = '1.2.5' +__version__ = '1.2.6' __all__ = [ # classes 'Actor', diff --git a/walnats/_tasks.py b/walnats/_tasks.py index 1dfab92..1e54955 100644 --- a/walnats/_tasks.py +++ b/walnats/_tasks.py @@ -8,29 +8,23 @@ class Tasks: """Supervise multiple async tasks. """ - __slots__ = ('_tasks', '_cleanup_every', '_since_cleanup', '_name', '_done') + __slots__ = ('_tasks', '_name', '_done') _name: str - _tasks: list[asyncio.Task] - _cleanup_every: int - _since_cleanup: int + _tasks: set[asyncio.Task] _done: bool def __init__(self, name: str) -> None: self._name = name - self._tasks = [] - self._cleanup_every = 100 - self._since_cleanup = 0 + self._tasks = set() self._done = False def start(self, coro: Coroutine[None, None, None], name: str) -> None: """Create a new task and track it in the supervisor. """ assert not self._done - self._since_cleanup = (self._since_cleanup + 1) % self._cleanup_every - if self._since_cleanup == 0: - self._tasks = [t for t in self._tasks if not t.done()] task = asyncio.create_task(coro, name=name) - self._tasks.append(task) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) def cancel(self) -> None: """Cancel all supervised tasks.