Don't allow hijacking loop

aiudirog · Oct 17, 2023 · 51221fd · 51221fd
1 parent 73d7c01
commit 51221fd
Showing 1 changed file with 10 additions and 5 deletions.
diff --git 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