|
| 1 | +import asyncio |
| 2 | +import logging |
| 3 | +from functools import partial |
| 4 | +from typing import Any, Coroutine |
| 5 | + |
| 6 | +import utils.app as app |
| 7 | + |
| 8 | +TASKS_FINISH_CHECK_TIME = 1 |
| 9 | + |
| 10 | +_logger = logging.getLogger("task_manager") |
| 11 | + |
| 12 | +_tasks: dict[asyncio.Task[Any] | None, list[asyncio.Task[Any]]] = {} |
| 13 | + |
| 14 | + |
| 15 | +def _on_parent_done(parent_task: asyncio.Task[Any], task: asyncio.Task[Any]) -> None: |
| 16 | + """Callback when parent task is done to cancel child task if it's still running""" |
| 17 | + if not task.done(): |
| 18 | + _logger.error(f"Cancelling task '{task.get_name()}' as parent task is done") |
| 19 | + task.cancel() |
| 20 | + |
| 21 | + |
| 22 | +def create_task( |
| 23 | + coro: Coroutine[Any, Any, Any], parent_task: asyncio.Task[Any] | None = None |
| 24 | +) -> asyncio.Task[Any]: |
| 25 | + """Create a task that will be executed in the background with an optional 'parent' attribute. If |
| 26 | + the parent task is done while the child task is running, the child task will be canceled""" |
| 27 | + task = asyncio.create_task(coro, name=coro.__name__) |
| 28 | + _tasks.setdefault(parent_task, []).append(task) |
| 29 | + |
| 30 | + if parent_task is not None: |
| 31 | + parent_task.add_done_callback(partial(_on_parent_done, task=task)) |
| 32 | + |
| 33 | + return task |
| 34 | + |
| 35 | + |
| 36 | +def _clear_completed() -> None: |
| 37 | + """Remove completed tasks from the global task list""" |
| 38 | + global _tasks |
| 39 | + _tasks = { |
| 40 | + parent: [task for task in tasks if not task.done()] for parent, tasks in _tasks.items() |
| 41 | + } |
| 42 | + _tasks = {parent: tasks for parent, tasks in _tasks.items() if len(tasks) > 0} |
| 43 | + |
| 44 | + |
| 45 | +async def wait_for_tasks( |
| 46 | + parent_task: asyncio.Task[Any] | None, timeout: float | None = None, cancel: bool = False |
| 47 | +) -> bool: |
| 48 | + """Wait for all running tasks started by the parent task to finish. If all tasks finish before |
| 49 | + the timeout, the function will return True. If the timeout is 'None', the function will wait |
| 50 | + until all tasks finish. If the timeout is reached, the function will return False. |
| 51 | + If cancel is True, all pending tasks will be cancelled on timeout.""" |
| 52 | + tasks = _tasks.get(parent_task, []) |
| 53 | + if len(tasks) == 0: |
| 54 | + return True |
| 55 | + |
| 56 | + done, pending = await asyncio.wait(tasks, timeout=timeout) |
| 57 | + |
| 58 | + if len(pending) > 0: |
| 59 | + if cancel: |
| 60 | + for task in pending: |
| 61 | + _logger.info(f"Task '{task.get_name()}' timed out") |
| 62 | + task.cancel() |
| 63 | + return False |
| 64 | + |
| 65 | + return True |
| 66 | + |
| 67 | + |
| 68 | +async def wait_for_all_tasks(timeout: float | None = None, cancel: bool = False) -> None: |
| 69 | + """Wait for all running tasks to finish""" |
| 70 | + for parent_task in _tasks.keys(): |
| 71 | + await wait_for_tasks(parent_task=parent_task, timeout=timeout, cancel=cancel) |
| 72 | + |
| 73 | + |
| 74 | +def _count_running(tasks: dict[Any, list[asyncio.Task[Any]]]) -> int: |
| 75 | + """Count the number of running tasks""" |
| 76 | + running_tasks = 0 |
| 77 | + for task_list in tasks.values(): |
| 78 | + running_tasks += len([task for task in task_list if not task.done()]) |
| 79 | + return running_tasks |
| 80 | + |
| 81 | + |
| 82 | +async def _wait_to_finish(tasks: dict[Any, list[asyncio.Task[Any]]]) -> None: |
| 83 | + """Wait for all running tasks to finish""" |
| 84 | + while True: |
| 85 | + running_tasks = _count_running(tasks) |
| 86 | + if running_tasks == 0: |
| 87 | + break |
| 88 | + _logger.info(f"Waiting for {running_tasks} tasks to finish") |
| 89 | + await asyncio.sleep(TASKS_FINISH_CHECK_TIME) |
| 90 | + |
| 91 | + |
| 92 | +async def run() -> None: |
| 93 | + _logger.info("Task manager running") |
| 94 | + |
| 95 | + while app.running(): |
| 96 | + _clear_completed() |
| 97 | + await app.sleep(60) |
| 98 | + |
| 99 | + _logger.info("Finishing") |
| 100 | + await _wait_to_finish(_tasks) |
0 commit comments