| 1 | import asyncio |
| 2 | from dataclasses import dataclass |
| 3 | import threading |
| 4 | from concurrent.futures import Future, InvalidStateError |
| 5 | from typing import Any, Callable, Optional, Coroutine, TypeVar, Awaitable |
| 6 | |
| 7 | T = TypeVar("T") |
| 8 | |
| 9 | THREAD_BACKGROUND = "Background" |
| 10 | |
| 11 | |
| 12 | class EventLoopThread: |
| 13 | _instances: dict[str, "EventLoopThread"] = {} |
| 14 | _lock = threading.Lock() |
| 15 | |
| 16 | def __init__(self, thread_name: str = THREAD_BACKGROUND) -> None: |
| 17 | """Initialize the event loop thread.""" |
| 18 | self.thread_name = thread_name |
| 19 | self._start() |
| 20 | |
| 21 | def __new__(cls, thread_name: str = THREAD_BACKGROUND): |
| 22 | with cls._lock: |
| 23 | if thread_name not in cls._instances: |
| 24 | instance = super(EventLoopThread, cls).__new__(cls) |
| 25 | cls._instances[thread_name] = instance |
| 26 | return cls._instances[thread_name] |
| 27 | |
| 28 | def _start(self): |
| 29 | if not hasattr(self, "loop") or not self.loop: |
| 30 | self.loop = asyncio.new_event_loop() |
| 31 | if not hasattr(self, "thread") or not self.thread: |
| 32 | self.thread = threading.Thread( |
| 33 | target=self._run_event_loop, daemon=True, name=self.thread_name |
| 34 | ) |
| 35 | self.thread.start() |
| 36 | |
| 37 | def _run_event_loop(self): |
| 38 | if not self.loop: |
| 39 | raise RuntimeError("Event loop is not initialized") |
| 40 | asyncio.set_event_loop(self.loop) |
| 41 | self.loop.run_forever() |
| 42 | |
| 43 | def terminate(self): |
| 44 | loop = getattr(self, "loop", None) |
| 45 | thread = getattr(self, "thread", None) |
| 46 | |
| 47 | if not loop: |
| 48 | return |
| 49 | |
| 50 | if loop.is_running(): |
| 51 | if thread and thread is threading.current_thread(): |
| 52 | loop.stop() |
| 53 | else: |
| 54 | loop.call_soon_threadsafe(loop.stop) |
| 55 | if thread: |
| 56 | thread.join() |
| 57 | elif thread and thread.is_alive() and thread is not threading.current_thread(): |
| 58 | thread.join() |
| 59 | |
| 60 | if not loop.is_closed(): |
| 61 | loop.close() |
| 62 | |
| 63 | with self.__class__._lock: |
| 64 | if self.thread_name in self.__class__._instances: |
| 65 | del self.__class__._instances[self.thread_name] |
| 66 | |
| 67 | self.loop = None |
| 68 | self.thread = None |
| 69 | |
| 70 | def run_coroutine(self, coro): |
| 71 | self._start() |
| 72 | if not self.loop: |
| 73 | raise RuntimeError("Event loop is not initialized") |
| 74 | return asyncio.run_coroutine_threadsafe(coro, self.loop) |
| 75 | |
| 76 | |
| 77 | @dataclass |
| 78 | class ChildTask: |
| 79 | task: "DeferredTask" |
| 80 | terminate_thread: bool |
| 81 | |
| 82 | |
| 83 | class DeferredTask: |
| 84 | def __init__( |
| 85 | self, |
| 86 | thread_name: str = THREAD_BACKGROUND, |
| 87 | ): |
| 88 | self.event_loop_thread = EventLoopThread(thread_name) |
| 89 | self._future: Optional[Future] = None |
| 90 | self.children: list[ChildTask] = [] |
| 91 | self.func: Optional[Callable[..., Coroutine[Any, Any, Any]]] = None |
| 92 | self.args: tuple[Any, ...] = () |
| 93 | self.kwargs: dict[str, Any] = {} |
| 94 | |
| 95 | def start_task( |
| 96 | self, func: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any |
| 97 | ): |
| 98 | self.func = func |
| 99 | self.args = args |
| 100 | self.kwargs = kwargs |
| 101 | self._start_task() |
| 102 | return self |
| 103 | |
| 104 | def add_done_callback(self, callback: Callable[[Future], Any]) -> None: |
| 105 | if not self._future: |
| 106 | raise RuntimeError("Task hasn't been started") |
| 107 | self._future.add_done_callback(callback) |
| 108 | |
| 109 | def __del__(self): |
| 110 | self.kill() |
| 111 | |
| 112 | def _start_task(self): |
| 113 | if self.func is None: |
| 114 | raise RuntimeError("Task callable is no longer available") |
| 115 | |
| 116 | self._future = self.event_loop_thread.run_coroutine( |
| 117 | self._run(self.func, self.args, self.kwargs) |
| 118 | ) |
| 119 | if self._future: |
| 120 | self._future.add_done_callback(self._on_task_done) |
| 121 | |
| 122 | def _on_task_done(self, future: Future): |
| 123 | # Ensure child background tasks are always cleaned up once the parent finishes |
| 124 | if future is self._future: |
| 125 | self.kill_children() |
| 126 | self._clear_call() |
| 127 | |
| 128 | def _clear_call(self) -> None: |
| 129 | self.func = None |
| 130 | self.args = () |
| 131 | self.kwargs = {} |
| 132 | |
| 133 | @staticmethod |
| 134 | async def _run(func, args, kwargs): |
| 135 | return await func(*args, **kwargs) |
| 136 | |
| 137 | def is_ready(self) -> bool: |
| 138 | return self._future.done() if self._future else False |
| 139 | |
| 140 | def result_sync(self, timeout: Optional[float] = None) -> Any: |
| 141 | if not self._future: |
| 142 | raise RuntimeError("Task hasn't been started") |
| 143 | try: |
| 144 | return self._future.result(timeout) |
| 145 | except TimeoutError: |
| 146 | raise TimeoutError( |
| 147 | "The task did not complete within the specified timeout." |
| 148 | ) |
| 149 | |
| 150 | async def result(self, timeout: Optional[float] = None) -> Any: |
| 151 | if not self._future: |
| 152 | raise RuntimeError("Task hasn't been started") |
| 153 | |
| 154 | loop = asyncio.get_running_loop() |
| 155 | |
| 156 | def _get_result(): |
| 157 | try: |
| 158 | result = self._future.result(timeout) # type: ignore |
| 159 | # self.kill() |
| 160 | return result |
| 161 | except TimeoutError: |
| 162 | raise TimeoutError( |
| 163 | "The task did not complete within the specified timeout." |
| 164 | ) |
| 165 | |
| 166 | return await loop.run_in_executor(None, _get_result) |
| 167 | |
| 168 | def kill(self, terminate_thread: bool = False) -> None: |
| 169 | """Kill the task and optionally terminate its thread.""" |
| 170 | self.kill_children() |
| 171 | if self._future and not self._future.done(): |
| 172 | self._future.cancel() |
| 173 | self._clear_call() |
| 174 | |
| 175 | if terminate_thread and self.event_loop_thread.loop: |
| 176 | if self.event_loop_thread.loop.is_running(): |
| 177 | try: |
| 178 | cleanup_future = asyncio.run_coroutine_threadsafe( |
| 179 | self._drain_event_loop_tasks(), self.event_loop_thread.loop |
| 180 | ) |
| 181 | cleanup_future.result() |
| 182 | except Exception: |
| 183 | pass |
| 184 | |
| 185 | self.event_loop_thread.terminate() |
| 186 | |
| 187 | def kill_children(self) -> None: |
| 188 | for child in self.children: |
| 189 | child.task.kill(terminate_thread=child.terminate_thread) |
| 190 | self.children = [] |
| 191 | |
| 192 | def is_alive(self) -> bool: |
| 193 | return self._future and not self._future.done() # type: ignore |
| 194 | |
| 195 | def restart(self, terminate_thread: bool = False) -> None: |
| 196 | if self.func is None: |
| 197 | raise RuntimeError("Completed task cannot be restarted") |
| 198 | func, args, kwargs = self.func, self.args, self.kwargs |
| 199 | self.kill(terminate_thread=terminate_thread) |
| 200 | self.start_task(func, *args, **kwargs) |
| 201 | |
| 202 | def add_child_task( |
| 203 | self, task: "DeferredTask", terminate_thread: bool = False |
| 204 | ) -> None: |
| 205 | self.children.append(ChildTask(task, terminate_thread)) |
| 206 | |
| 207 | async def _execute_in_task_context( |
| 208 | self, func: Callable[..., T], *args, **kwargs |
| 209 | ) -> T: |
| 210 | """Execute a function in the task's context and return its result.""" |
| 211 | result = func(*args, **kwargs) |
| 212 | if asyncio.iscoroutine(result): |
| 213 | return await result |
| 214 | return result |
| 215 | |
| 216 | def execute_inside(self, func: Callable[..., T], *args, **kwargs) -> Awaitable[T]: |
| 217 | if not self.event_loop_thread.loop: |
| 218 | raise RuntimeError("Event loop is not initialized") |
| 219 | |
| 220 | future: Future = Future() |
| 221 | |
| 222 | def set_result(result: Any) -> None: |
| 223 | try: |
| 224 | future.set_result(result) |
| 225 | except InvalidStateError: |
| 226 | pass |
| 227 | |
| 228 | def set_exception(exception: BaseException) -> None: |
| 229 | try: |
| 230 | future.set_exception(exception) |
| 231 | except InvalidStateError: |
| 232 | pass |
| 233 | |
| 234 | async def wrapped(): |
| 235 | if not self.event_loop_thread.loop: |
| 236 | raise RuntimeError("Event loop is not initialized") |
| 237 | try: |
| 238 | result = await self._execute_in_task_context(func, *args, **kwargs) |
| 239 | # Keep awaiting until we get a concrete value |
| 240 | while isinstance(result, Awaitable): |
| 241 | result = await result |
| 242 | self.event_loop_thread.loop.call_soon_threadsafe( |
| 243 | set_result, result |
| 244 | ) |
| 245 | except Exception as e: |
| 246 | self.event_loop_thread.loop.call_soon_threadsafe( |
| 247 | set_exception, e |
| 248 | ) |
| 249 | |
| 250 | asyncio.run_coroutine_threadsafe(wrapped(), self.event_loop_thread.loop) |
| 251 | return asyncio.wrap_future(future) |
| 252 | |
| 253 | @staticmethod |
| 254 | async def _drain_event_loop_tasks(): |
| 255 | """Cancel and await all pending tasks on the current event loop.""" |
| 256 | loop = asyncio.get_running_loop() |
| 257 | current_task = asyncio.current_task(loop=loop) |
| 258 | pending = [ |
| 259 | task |
| 260 | for task in asyncio.all_tasks(loop=loop) |
| 261 | if task is not current_task |
| 262 | ] |
| 263 | if not pending: |
| 264 | return |
| 265 | for task in pending: |
| 266 | task.cancel() |
| 267 | await asyncio.gather(*pending, return_exceptions=True) |