From 51221fdf27642ffc8bcdb83adee77a12ec01493f Mon Sep 17 00:00:00 2001 From: Roger Aiudi Date: Tue, 17 Oct 2023 10:41:47 -0400 Subject: [PATCH] Don't allow hijacking loop --- aiuti/asyncio.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/aiuti/asyncio.py b/aiuti/asyncio.py index 9d8b801..b6239dd 100644 --- a/aiuti/asyncio.py +++ b/aiuti/asyncio.py @@ -423,11 +423,16 @@ async def _wrapper(*args: Any, **kwargs: Any) -> Any: else: waiting = True # Need to wait for other loop - if waiting: # Wait for other loop, maybe across threads - try: - await ensure_aw(event.wait(), loop) - except RuntimeError: # Target loop most likely closed - pass + if waiting: # Wait for other task, maybe across threads + if aio.get_running_loop() is loop: + await event.wait() + else: + try: + await aio.wrap_future( + aio.run_coroutine_threadsafe(event.wait(), loop), + ) + except RuntimeError: # Target loop most likely closed + pass continue # First to arrive, cache the value