From 13195a678d0155a52818beaa5ffc4f4678a4226d Mon Sep 17 00:00:00 2001
From: Damien George <damien@micropython.org>
Date: Wed, 19 Jun 2024 15:04:12 +1000
Subject: [PATCH] webassembly/asyncio: Schedule run loop when tasks are pushed
 to queue.

In the webassembly port there is no asyncio run loop running at the top
level.  Instead the Python asyncio run loop is scheduled through setTimeout
and run by the outer JavaScript event loop.  Because tasks can become
runable from an external (to Python) event (eg a JavaScript callback), the
run loop must be scheduled whenever a task is pushed to the asyncio task
queue, otherwise tasks may be waiting forever on the queue.

Signed-off-by: Damien George <damien@micropython.org>
---
 ports/webassembly/asyncio/core.py             |  5 +-
 ports/webassembly/mpconfigport.h              |  1 +
 .../webassembly/asyncio_top_level_await.mjs   | 65 +++++++++++++++++++
 .../asyncio_top_level_await.mjs.exp           | 12 ++++
 4 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/ports/webassembly/asyncio/core.py b/ports/webassembly/asyncio/core.py
index cc26e7b8d..47846fc25 100644
--- a/ports/webassembly/asyncio/core.py
+++ b/ports/webassembly/asyncio/core.py
@@ -71,7 +71,6 @@ class TopLevelCoro:
     def set(resolve, reject):
         TopLevelCoro.resolve = resolve
         TopLevelCoro.reject = reject
-        _schedule_run_iter(0)
 
     @staticmethod
     def send(value):
@@ -90,7 +89,6 @@ class ThenableEvent:
         if self.waiting:
             _task_queue.push(self.waiting)
             self.waiting = None
-            _schedule_run_iter(0)
 
     def remove(self, task):
         self.waiting = None
@@ -203,7 +201,6 @@ def create_task(coro):
         raise TypeError("coroutine expected")
     t = Task(coro, globals())
     _task_queue.push(t)
-    _schedule_run_iter(0)
     return t
 
 
@@ -253,7 +250,7 @@ def current_task():
 
 def new_event_loop():
     global _task_queue
-    _task_queue = TaskQueue()  # TaskQueue of Task instances.
+    _task_queue = TaskQueue(_schedule_run_iter)  # TaskQueue of Task instances.
     return Loop
 
 
diff --git a/ports/webassembly/mpconfigport.h b/ports/webassembly/mpconfigport.h
index ece7868fb..ab56162ca 100644
--- a/ports/webassembly/mpconfigport.h
+++ b/ports/webassembly/mpconfigport.h
@@ -56,6 +56,7 @@
 #define MICROPY_USE_INTERNAL_PRINTF (0)
 
 #define MICROPY_EPOCH_IS_1970       (1)
+#define MICROPY_PY_ASYNCIO_TASK_QUEUE_PUSH_CALLBACK (1)
 #define MICROPY_PY_RANDOM_SEED_INIT_FUNC (mp_js_random_u32())
 #define MICROPY_PY_TIME_GMTIME_LOCALTIME_MKTIME (1)
 #define MICROPY_PY_TIME_TIME_TIME_NS (1)
diff --git a/tests/ports/webassembly/asyncio_top_level_await.mjs b/tests/ports/webassembly/asyncio_top_level_await.mjs
index d8a9cad42..234b7a6ce 100644
--- a/tests/ports/webassembly/asyncio_top_level_await.mjs
+++ b/tests/ports/webassembly/asyncio_top_level_await.mjs
@@ -2,6 +2,71 @@
 
 const mp = await (await import(process.argv[2])).loadMicroPython();
 
+/**********************************************************/
+// Top-level await for an Event which is set by a JavaScript
+// callback.
+
+console.log("= TEST 1 ==========");
+
+await mp.runPythonAsync(`
+import asyncio
+import js
+
+event = asyncio.Event()
+
+def callback():
+    print("callback set event")
+    event.set()
+
+js.setTimeout(callback, 100)
+
+print("top-level wait event")
+await event.wait()
+print("top-level end")
+`);
+
+console.log("finished");
+
+/**********************************************************/
+// Top-level await for a Task which is cancelled by a
+// JavaScript callback.
+
+console.log("= TEST 2 ==========");
+
+await mp.runPythonAsync(`
+import asyncio
+import js
+import time
+
+async def task():
+    print("task start")
+    await asyncio.sleep(5)
+    print("task end")
+
+def callback():
+    print("callback cancel task")
+    t.cancel()
+
+t  = asyncio.create_task(task())
+js.setTimeout(callback, 100)
+
+print("top-level wait task")
+try:
+    t0 = time.time()
+    await t
+except asyncio.CancelledError:
+    dt = time.time() - t0
+    print("top-level task CancelledError", dt < 1)
+`);
+
+console.log("finished");
+
+/**********************************************************/
+// Top-level await for an Event and a Task, with the task
+// setting the event.
+
+console.log("= TEST 3 ==========");
+
 await mp.runPythonAsync(`
 import asyncio
 
diff --git a/tests/ports/webassembly/asyncio_top_level_await.mjs.exp b/tests/ports/webassembly/asyncio_top_level_await.mjs.exp
index 7232c5d4f..66fefd2dc 100644
--- a/tests/ports/webassembly/asyncio_top_level_await.mjs.exp
+++ b/tests/ports/webassembly/asyncio_top_level_await.mjs.exp
@@ -1,3 +1,15 @@
+= TEST 1 ==========
+top-level wait event
+callback set event
+top-level end
+finished
+= TEST 2 ==========
+top-level wait task
+task start
+callback cancel task
+top-level task CancelledError True
+finished
+= TEST 3 ==========
 top-level wait event
 task set event
 task sleep