File size: 1,033 Bytes
deea75c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | import logging
import asyncio
from typing import Dict
logger = logging.getLogger("s2s_server")
class TaskSupervisor:
def __init__(self, session_id: str):
self.session_id = session_id
self._tasks: Dict[str, asyncio.Task] = {}
def register_task(self, name: str, task: asyncio.Task):
# Cancel if already registered
self.cancel_task(name)
self._tasks[name] = task
def cancel_task(self, name: str):
task = self._tasks.get(name)
if task and not task.done():
logger.info(f"[{self.session_id}] TaskSupervisor: Cancelling task '{name}'")
task.cancel()
if name in self._tasks:
del self._tasks[name]
async def cancel_tasks(self, *names: str):
for name in names:
self.cancel_task(name)
def cancel_all(self):
logger.info(f"[{self.session_id}] TaskSupervisor: Cancelling all tasks ({list(self._tasks.keys())})")
for name in list(self._tasks.keys()):
self.cancel_task(name)
|