Defer fix
Fix for event loops in defer.py
frdel committed
Sep 13, 2024 at 14:30 UTC
62e244f96fba209eb31914b6cef483f0d280e317
1 file changed
+38
-47
python/helpers/defer.py
+38
-47
@@ -1,70 +1,61 @@
1
import asyncio
2
import threading
3
from concurrent.futures import Future
4
+from typing import Any, Callable, Optional, Coroutine
5
5
-class DeferredTask:
6
- def __init__(self, func, *args, **kwargs):
7
- self._loop: asyncio.AbstractEventLoop = None # type: ignore
8
- self._task = None
9
- self._future = Future()
10
- self._task_initialized = threading.Event()
11
- self._start_task(func, *args, **kwargs)
6
+class EventLoopThread:
7
+ _instance = None
8
+ _lock = threading.Lock()
9
13
- def _start_task(self, func, *args, **kwargs):
14
- def run_in_thread():
15
- self._loop = asyncio.new_event_loop()
16
- asyncio.set_event_loop(self._loop)
17
- self._task = self._loop.create_task(self._run(func, *args, **kwargs))
18
- self._task_initialized.set()
19
- self._loop.run_forever()
10
+ def __init__(self) -> None:
11
+ self.loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
12
+ self.thread: threading.Thread = threading.Thread(target=self._run_event_loop, daemon=True)
13
+ self.thread.start()
14
21
- self._thread = threading.Thread(target=run_in_thread)
22
- self._thread.start()
15
+ def __new__(cls) -> 'EventLoopThread':
16
+ with cls._lock:
17
+ if cls._instance is None:
18
+ cls._instance = super().__new__(cls)
19
+ cls._instance.__init__()
20
+ return cls._instance
21
24
- async def _run(self, func, *args, **kwargs):
25
- try:
26
- result = await func(*args, **kwargs)
27
- self._future.set_result(result)
28
- except Exception as e:
29
- self._future.set_exception(e)
30
- finally:
31
- self._loop.call_soon_threadsafe(self._cleanup)
22
+ def _run_event_loop(self):
23
+ asyncio.set_event_loop(self.loop)
24
+ self.loop.run_forever()
25
33
- def _cleanup(self):
34
- self._loop.stop()
26
+ def run_coroutine(self, coro):
27
+ return asyncio.run_coroutine_threadsafe(coro, self.loop)
28
36
- def is_ready(self):
37
- return self._future.done()
29
+class DeferredTask:
30
+ def __init__(self, func: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any) -> None:
31
+ self._event_loop_thread = EventLoopThread()
32
+ self._future: Future[Any] = self._event_loop_thread.run_coroutine(self._run(func, *args, **kwargs))
33
+
34
+ async def _run(self, func: Callable[..., Coroutine[Any, Any, Any]], *args: Any, **kwargs: Any) -> Any:
35
+ return await func(*args, **kwargs)
36
39
- async def result(self, timeout=None):
40
- if not self._task_initialized.wait(timeout):
41
- raise RuntimeError("Task was not initialized properly.")
37
+ def is_ready(self) -> bool:
38
+ return self._future.done()
39
40
+ async def result(self, timeout: Optional[float] = None) -> Any:
41
try:
42
return await asyncio.wait_for(asyncio.wrap_future(self._future), timeout)
43
except asyncio.TimeoutError:
44
raise TimeoutError("The task did not complete within the specified timeout.")
45
48
- def result_sync(self, timeout=None):
49
- if not self._task_initialized.wait(timeout):
50
- raise RuntimeError("Task was not initialized properly.")
51
-
46
+ def result_sync(self, timeout: Optional[float] = None) -> Any:
47
try:
48
return self._future.result(timeout)
49
except TimeoutError:
50
raise TimeoutError("The task did not complete within the specified timeout.")
51
57
- def kill(self):
58
- if self._task and not self._task.done():
59
- self._loop.call_soon_threadsafe(self._task.cancel)
52
+ def kill(self) -> None:
53
+ if not self._future.done():
54
+ self._future.cancel()
55
61
- def is_alive(self):
62
- return self._thread.is_alive() and not self._future.done()
56
+ def is_alive(self) -> bool:
57
+ return not self._future.done()
58
64
- def __del__(self):
65
- if self._loop and self._loop.is_running():
66
- self._loop.call_soon_threadsafe(self._cleanup)
67
- if self._thread and self._thread.is_alive():
68
- self._thread.join()
69
- if self._loop:
70
- self._loop.close()
\ No newline at end of file
59
+# Helper function to run async code
60
+async def run_async(func, *args, **kwargs):
61
+ return await func(*args, **kwargs)
\ No newline at end of file