micropython · greezybacon · Jul 28, 2024 · Aug 25, 2024
diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py
@@ -129,6 +129,15 @@ def wait_io_event(self, dt):
             else:
                 self.poller.modify(s, select.POLLIN)
 
+    def all_tasks(self):
+        tasks = set()
+        for q1, q2, _ in self.map.values():
+            if q1 is not None:
+                tasks.add(q1)
+            if q2 is not None:
+                tasks.add(q2)
+
+        return tasks
 
 ################################################################################
 # Main run loop
@@ -233,7 +242,11 @@ def run_until_complete(main_task=None):
 
 # Create a new task from a coroutine and run it until it finishes
 def run(coro):
-    return run_until_complete(create_task(coro))
+    try:
+        return run_until_complete(create_task(coro))
+    finally:
+        Loop.close()
+        run_until_complete()
 
 
 ################################################################################
@@ -271,7 +284,17 @@ def stop():
             _stop_task = None
 
     def close():
-        pass
+        for t in Loop.all_tasks():
+            if not t.done() and t is not cur_task:
+                t.cancel()
+
+    def all_tasks():
+        tasks = set(_task_queue)
+        tasks.update(_io_queue.all_tasks())
+        if cur_task is not None:
+            tasks.add(cur_task)
+
+        return tasks
 
     def set_exception_handler(handler):
         Loop._exc_handler = handler
@@ -299,6 +322,13 @@ def current_task():
     return cur_task
 
 
+def all_tasks():
+    if cur_task is None:
+        raise RuntimeError("no running event loop")
+
+    return Loop.all_tasks()
+
+
 def new_event_loop():
     global _task_queue, _io_queue
     # TaskQueue of Task instances

diff --git a/extmod/asyncio/task.py b/extmod/asyncio/task.py
@@ -116,6 +116,12 @@ def pop(self):
     def remove(self, v):
         self.heap = ph_delete(self.heap, v)
 
+    def __iter__(self):
+        heap = self.heap
+        while heap:
+            yield heap
+            heap = heap.ph_child
+
 
 # Task class representing a coroutine, can be waited on and cancelled.
 class Task:

diff --git a/extmod/modasyncio.c b/extmod/modasyncio.c
@@ -151,6 +151,38 @@ static mp_obj_t task_queue_remove(mp_obj_t self_in, mp_obj_t task_in) {
 }
 static MP_DEFINE_CONST_FUN_OBJ_2(task_queue_remove_obj, task_queue_remove);
 
+typedef struct _mp_obj_task_queue_it_t {
+    mp_obj_base_t base;
+    mp_fun_1_t iternext;
+    mp_obj_task_t *cur;
+} mp_obj_task_queue_it_t;
+
+static mp_obj_t task_queue_iternext(mp_obj_t self_in) {
+    mp_obj_task_queue_it_t *self = MP_OBJ_TO_PTR(self_in);
+
+    if (self->cur != NULL) {
+        mp_obj_task_t *current = self->cur;
+        self->cur = (mp_obj_task_t *)self->cur->pairheap.child;
+        return current;
+    }
+    return MP_OBJ_STOP_ITERATION;
+}
+
+static mp_obj_t task_queue_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
+    assert(sizeof(mp_obj_task_queue_it_t) <= sizeof(mp_obj_iter_buf_t));
+    mp_obj_task_queue_t *self = MP_OBJ_TO_PTR(self_in);
+    mp_obj_task_queue_it_t *o = (mp_obj_task_queue_it_t *)iter_buf;
+    o->base.type = &mp_type_polymorph_iter;
+    o->iternext = task_queue_iternext;
+    o->cur = self->heap;
+    return MP_OBJ_FROM_PTR(o);
+}
+
+static const mp_getiter_iternext_custom_t task_queue_getiter_iternext = {
+    .getiter = task_queue_getiter,
+    .iternext = task_queue_iternext,
+};
+
 static const mp_rom_map_elem_t task_queue_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_peek), MP_ROM_PTR(&task_queue_peek_obj) },
     { MP_ROM_QSTR(MP_QSTR_push), MP_ROM_PTR(&task_queue_push_obj) },
@@ -162,9 +194,10 @@ static MP_DEFINE_CONST_DICT(task_queue_locals_dict, task_queue_locals_dict_table
 static MP_DEFINE_CONST_OBJ_TYPE(
     task_queue_type,
     MP_QSTR_TaskQueue,
-    MP_TYPE_FLAG_NONE,
+    MP_TYPE_FLAG_ITER_IS_CUSTOM,
     make_new, task_queue_make_new,
-    locals_dict, &task_queue_locals_dict
+    locals_dict, &task_queue_locals_dict,
+    iter, &task_queue_getiter_iternext
     );
 
 /******************************************************************************/

diff --git a/tests/extmod/asyncio_all_tasks.py b/tests/extmod/asyncio_all_tasks.py
@@ -0,0 +1,50 @@
+try:
+    import asyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+import sys
+
+async def waiter():
+    await asyncio.sleep(1)
+
+async def reader():
+    await asyncio.StreamReader().read()
+
+
+async def print_all_tasks():
+    print(list('T' for t in asyncio.all_tasks()))
+
+async def short_wait():
+    await asyncio.sleep(0.01)
+
+async def waiters():
+    t1 = asyncio.create_task(waiter())
+    # NOTE: The singleton generator for ::sleep expects the task to be scheduled
+    # before using it again.
+    await asyncio.sleep(0)
+    t2 = asyncio.create_task(waiter())
+    await asyncio.sleep(0)
+    t3 = asyncio.create_task(short_wait())
+    # Tasks are in the list before they start to run
+    assert t1 in asyncio.all_tasks()
+    assert t2 in asyncio.all_tasks()
+    await print_all_tasks()
+    await t3
+    assert t1 in asyncio.all_tasks()
+    assert t2 in asyncio.all_tasks()
+    assert t3 not in asyncio.all_tasks()
+
+async def readers():
+    t1 = asyncio.create_task(reader())
+    t2 = asyncio.create_task(reader())
+    await print_all_tasks()
+    assert t1 in asyncio.all_tasks()
+    assert t2 in asyncio.all_tasks()
+
+asyncio.run(print_all_tasks())
+asyncio.new_event_loop()
+asyncio.run(readers())
+asyncio.new_event_loop()
+asyncio.run(waiters())
diff --git a/tests/extmod/asyncio_shutdown.py b/tests/extmod/asyncio_shutdown.py
@@ -0,0 +1,135 @@
+try:
+    import asyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+import sys
+
+# CPython and Micropython disagree on the sorting of tasks in the ::all_tasks
+# set, so sort the output of the tasks as they shut down rather than expecting
+# them to shut down in the same order.
+class SortedList(list):
+    def print(self):
+        print(*sorted(self))
+
+finished = SortedList()
+
+async def task1():
+    try:
+        await asyncio.sleep(1)
+    finally:
+        finished.append('task1 finished')
+
+async def task2():
+    try:
+        await asyncio.sleep(1)
+    finally:
+        finished.append("task2 finished")
+
+async def reader():
+    try:
+        await asyncio.StreamReader(sys.stdin).read()
+    finally:
+        finished.append("reader finished")
+
+# KeyboardInterrupt in main task should be propagated
+async def ctrlc():
+    asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    asyncio.create_task(task2())
+    await asyncio.sleep(0.01)
+    raise KeyboardInterrupt
+
+
+# If the main task is canceled, the loop should be closed
+async def cancel_main():
+    async def canceler(delay, task):
+        await asyncio.sleep(delay)
+        task.cancel()
+
+    asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    asyncio.create_task(task2())
+    asyncio.create_task(canceler(0.01, asyncio.current_task()))
+    await asyncio.sleep(0.02)
+
+# SystemExit should close the loop
+async def sys_exit():
+    asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    asyncio.create_task(task2())
+    await asyncio.sleep(0.01)
+    raise SystemExit
+
+# Cancelling a background task should have no effect
+async def cancel_bg():
+    t1 = asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    t2 = asyncio.create_task(task2())
+    await asyncio.sleep(0.01)
+    t1.cancel()
+    t2.cancel()
+    await asyncio.sleep(0.01)
+
+# Reader tasks should also be cancelled when the loop is shutdown
+async def cancel_reader():
+    t1 = asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    t2 = asyncio.create_task(reader())
+    await asyncio.sleep(0.01)
+    t1.cancel()
+    t2.cancel()
+    await asyncio.sleep(0.01)
+
+async def close_loop():
+    asyncio.create_task(task1())
+    await asyncio.sleep(0)
+    asyncio.create_task(reader())
+    await asyncio.sleep(0.01)
+
+try:
+    asyncio.new_event_loop()
+    finished.clear()
+    asyncio.run(ctrlc())
+except KeyboardInterrupt:
+    print("KeyboardInterrupt raised")
+finally:
+    finished.print()
+
+try:
+    asyncio.new_event_loop()
+    finished.clear()
+    asyncio.run(cancel_main())
+except asyncio.CancelledError:
+    print("CancelledError raised")
+finally:
+    finished.print()
+
+# Cancel a background waiting task doesn't raise an exception
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_bg())
+finished.print()
+
+# Cancel a background IO task doesn't raise an exception
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_reader())
+finished.print()
+
+print('---')
+asyncio.new_event_loop()
+finished.clear()
+asyncio.run(cancel_reader())
+finished.print()
+
+try:
+    asyncio.new_event_loop()
+    finished.clear()
+    asyncio.run(sys_exit())
+except SystemExit:
+    print("SystemExit raised")
+