asyncio: Properly cancel the main task on exception

greezybacon · greezybacon · commit 615bf4ac2682 · 2022-12-02T21:33:54.000-06:00
If the main task is interrupted by e.g. a KeyboardInterrupt, then the
main task needs to have the exception injected into it so it will run
the exception handlers and contextmanager __aexit__ methods.
diff --git a/extmod/uasyncio/core.py b/extmod/uasyncio/core.py
@@ -145,81 +145,88 @@ def create_task(coro):
     _task_queue.push(t)
     return t
 
-
 # Keep scheduling tasks until there are none left to schedule
 def run_until_complete(main_task=None):
     global cur_task
     excs_all = (CancelledError, Exception)  # To prevent heap allocation in loop
     excs_stop = (CancelledError, StopIteration)  # To prevent heap allocation in loop
-    while True:
-        # Wait until the head of _task_queue is ready to run
-        dt = 1
-        while dt > 0:
-            dt = -1
+    try:
+        _io_queue.wait_io_event(0)
+        while True:
+            # Wait until the head of _task_queue is ready to run
             t = _task_queue.peek()
             if t:
                 # A task waiting on _task_queue; "ph_key" is time to schedule task at
-                dt = max(0, ticks_diff(t.ph_key, ticks()))
+                dt = ticks_diff(t.ph_key, ticks())
+                if dt > 0:
+                    _io_queue.wait_io_event(dt)
             elif not _io_queue.map:
                 # No tasks can be woken so finished running
                 return
-            # print('(poll {})'.format(dt), len(_io_queue.map))
-            _io_queue.wait_io_event(dt)
-
-        # Get next task to run and continue it
-        t = _task_queue.pop()
-        cur_task = t
-        try:
-            # Continue running the coroutine, it's responsible for rescheduling itself
-            exc = t.data
-            if not exc:
-                t.coro.send(None)
             else:
-                # If the task is finished and on the run queue and gets here, then it
-                # had an exception and was not await'ed on.  Throwing into it now will
-                # raise StopIteration and the code below will catch this and run the
-                # call_exception_handler function.
-                t.data = None
-                t.coro.throw(exc)
-        except excs_all as er:
-            # Check the task is not on any event queue
-            assert t.data is None
-            # This task is done, check if it's the main task and then loop should stop
-            if t is main_task:
-                if isinstance(er, StopIteration):
-                    return er.value
-                raise er
-            if t.state:
-                # Task was running but is now finished.
-                waiting = False
-                if t.state is True:
-                    # "None" indicates that the task is complete and not await'ed on (yet).
-                    t.state = None
-                elif callable(t.state):
-                    # The task has a callback registered to be called on completion.
-                    t.state(t, er)
-                    t.state = False
-                    waiting = True
+                _io_queue.wait_io_event(-1)
+
+            # Get next task to run and continue it
+            t = _task_queue.pop()
+            cur_task = t
+            try:
+                # Continue running the coroutine, it's responsible for rescheduling itself
+                exc = t.data
+                if not exc:
+                    t.coro.send(None)
                 else:
-                    # Schedule any other tasks waiting on the completion of this task.
-                    while t.state.peek():
-                        _task_queue.push(t.state.pop())
+                    # If the task is finished and on the run queue and gets here, then it
+                    # had an exception and was not await'ed on.  Throwing into it now will
+                    # raise StopIteration and the code below will catch this and run the
+                    # call_exception_handler function.
+                    t.data = None
+                    t.coro.throw(exc)
+            except excs_all as er:
+                # Check the task is not on any event queue
+                assert t.data is None
+                # This task is done, check if it's the main task and then loop should stop
+                if t is main_task:
+                    if isinstance(er, StopIteration):
+                        return er.value
+                    raise er
+                if t.state:
+                    # Task was running but is now finished.
+                    waiting = False
+                    if t.state is True:
+                        # "None" indicates that the task is complete and not await'ed on (yet).
+                        t.state = None
+                    elif callable(t.state):
+                        # The task has a callback registered to be called on completion.
+                        t.state(t, er)
+                        t.state = False
                         waiting = True
-                    # "False" indicates that the task is complete and has been await'ed on.
-                    t.state = False
-                if not waiting and not isinstance(er, excs_stop):
-                    # An exception ended this detached task, so queue it for later
-                    # execution to handle the uncaught exception if no other task retrieves
-                    # the exception in the meantime (this is handled by Task.throw).
-                    _task_queue.push(t)
-                # Save return value of coro to pass up to caller.
-                t.data = er
-            elif t.state is None:
-                # Task is already finished and nothing await'ed on the task,
-                # so call the exception handler.
-                _exc_context["exception"] = exc
-                _exc_context["future"] = t
-                Loop.call_exception_handler(_exc_context)
+                    else:
+                        # Schedule any other tasks waiting on the completion of this task.
+                        while t.state.peek():
+                            _task_queue.push(t.state.pop())
+                            waiting = True
+                        # "False" indicates that the task is complete and has been await'ed on.
+                        t.state = False
+                    if not waiting and not isinstance(er, excs_stop):
+                        # An exception ended this detached task, so queue it for later
+                        # execution to handle the uncaught exception if no other task retrieves
+                        # the exception in the meantime (this is handled by Task.throw).
+                        _task_queue.push(t)
+                    # Save return value of coro to pass up to caller.
+                    t.data = er
+                elif t.state is None:
+                    # Task is already finished and nothing await'ed on the task,
+                    # so call the exception handler.
+                    _exc_context["exception"] = exc
+                    _exc_context["future"] = t
+                    Loop.call_exception_handler(_exc_context)
+    except BaseException as exc:
+        try:
+            if main_task:
+                main_task.coro.throw(exc)
+        except StopIteration:
+            pass
+        raise
 
 
 # Create a new task from a coroutine and run it until it finishes
