TaskManager: use a dictionary instead of a set to store tasks

This commit is contained in:
Aleix Conchillo Flaqué
2025-05-05 21:27:49 -07:00
parent 3824da7261
commit 9cc498b1fa
2 changed files with 10 additions and 7 deletions

View File

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Fixed ### Fixed
- Fixed a `TaskManager` that was causing dangling tasks to be reported.
- Fixed an issue that could cause data to be sent to the transports when they - Fixed an issue that could cause data to be sent to the transports when they
were still not ready. were still not ready.

View File

@@ -6,7 +6,7 @@
import asyncio import asyncio
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Coroutine, Optional, Set from typing import Coroutine, Dict, Optional, Sequence, Set
from loguru import logger from loguru import logger
@@ -69,14 +69,14 @@ class BaseTaskManager(ABC):
pass pass
@abstractmethod @abstractmethod
def current_tasks(self) -> Set[asyncio.Task]: def current_tasks(self) -> Sequence[asyncio.Task]:
"""Returns the list of currently created/registered tasks.""" """Returns the list of currently created/registered tasks."""
pass pass
class TaskManager(BaseTaskManager): class TaskManager(BaseTaskManager):
def __init__(self) -> None: def __init__(self) -> None:
self._tasks: Set[asyncio.Task] = set() self._tasks: Dict[str, asyncio.Task] = {}
self._loop: Optional[asyncio.AbstractEventLoop] = None self._loop: Optional[asyncio.AbstractEventLoop] = None
def set_event_loop(self, loop: asyncio.AbstractEventLoop): def set_event_loop(self, loop: asyncio.AbstractEventLoop):
@@ -179,16 +179,17 @@ class TaskManager(BaseTaskManager):
finally: finally:
self._remove_task(task) self._remove_task(task)
def current_tasks(self) -> Set[asyncio.Task]: def current_tasks(self) -> Sequence[asyncio.Task]:
"""Returns the list of currently created/registered tasks.""" """Returns the list of currently created/registered tasks."""
return self._tasks return list(self._tasks.values())
def _add_task(self, task: asyncio.Task): def _add_task(self, task: asyncio.Task):
self._tasks.add(task) name = task.get_name()
self._tasks[name] = task
def _remove_task(self, task: asyncio.Task): def _remove_task(self, task: asyncio.Task):
name = task.get_name() name = task.get_name()
try: try:
self._tasks.remove(task) del self._tasks[name]
except KeyError as e: except KeyError as e:
logger.trace(f"{name}: unable to remove task (already removed?): {e}") logger.trace(f"{name}: unable to remove task (already removed?): {e}")