asyncio: Properly cancel the main task on exception.

greezybacon · greezybacon · commit 4a8aa483e7cd · 2024-07-27T11:50:46.000-05: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.

Additionally, if an error is encountered in the poll loop, it will be
injected into the main task.

Signed-off-by: Jared Hancock &lt;jared@greezybacon.me&gt;
diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py
@@ -3,6 +3,7 @@
 
 from time import ticks_ms as ticks, ticks_diff, ticks_add
 import sys, select
+from select import POLLIN, POLLOUT
 
 # Import TaskQueue and Task, preferring built-in C code over Python code
 try:
@@ -54,7 +55,8 @@ def __next__(self):
 # Use a SingletonGenerator to do it without allocating on the heap
 def sleep_ms(t, sgen=SingletonGenerator()):
     assert sgen.state is None
-    sgen.state = ticks_add(ticks(), max(0, t))
+    now = ticks()
+    sgen.state = ticks_add(now, t) if t > 0 else now
     return sgen
 
 
@@ -66,7 +68,6 @@ def sleep(t):
 ################################################################################
 # Queue and poller for stream IO
 
-
 class IOQueue:
     def __init__(self):
         self.poller = select.poll()
@@ -77,13 +78,13 @@ def _enqueue(self, s, idx):
             entry = [None, None, s]
             entry[idx] = cur_task
             self.map[id(s)] = entry
-            self.poller.register(s, select.POLLIN if idx == 0 else select.POLLOUT)
+            self.poller.register(s, POLLIN if idx == 0 else POLLOUT)
         else:
             sm = self.map[id(s)]
             assert sm[idx] is None
             assert sm[1 - idx] is not None
             sm[idx] = cur_task
-            self.poller.modify(s, select.POLLIN | select.POLLOUT)
+            self.poller.modify(s, POLLIN | POLLOUT)
         # Link task to this IOQueue so it can be removed if needed
         cur_task.data = self
 
@@ -114,20 +115,20 @@ def wait_io_event(self, dt):
         for s, ev in self.poller.ipoll(dt):
             sm = self.map[id(s)]
             # print('poll', s, sm, ev)
-            if ev & ~select.POLLOUT and sm[0] is not None:
+            if ev & ~POLLOUT and sm[0] is not None:
                 # POLLIN or error
                 _task_queue.push(sm[0])
                 sm[0] = None
-            if ev & ~select.POLLIN and sm[1] is not None:
+            if ev & ~POLLIN and sm[1] is not None:
                 # POLLOUT or error
                 _task_queue.push(sm[1])
                 sm[1] = None
             if sm[0] is None and sm[1] is None:
                 self._dequeue(s)
             elif sm[0] is None:
-                self.poller.modify(s, select.POLLOUT)
+                self.poller.modify(s, POLLOUT)
             else:
-                self.poller.modify(s, select.POLLIN)
+                self.poller.modify(s, POLLIN)
 
 
 ################################################################################
@@ -153,24 +154,36 @@ 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
+    queue_peek = _task_queue.peek
+    queue_pop = _task_queue.pop
+    wait_io_event = _io_queue.wait_io_event
     while True:
-        # Wait until the head of _task_queue is ready to run
-        dt = 1
-        while dt > 0:
-            dt = -1
-            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()))
-            elif not _io_queue.map:
-                # No tasks can be woken so finished running
-                cur_task = None
-                return
-            # print('(poll {})'.format(dt), len(_io_queue.map))
-            _io_queue.wait_io_event(dt)
+        try:
+            # Wait until the head of _task_queue is ready to run
+            dt = 1
+            while dt > 0:
+                dt = -1
+                t = queue_peek()
+                if t:
+                    # A task waiting on _task_queue; "ph_key" is time to schedule task at
+                    dt = ticks_diff(t.ph_key, ticks())
+                    dt = dt if dt > 0 else 0
+                elif not _io_queue.map:
+                    # No tasks can be woken so finished running
+                    cur_task = None
+                    return
+                # print('(poll {})'.format(dt), len(_io_queue.map))
+                wait_io_event(dt)
+        except BaseException as exc:
+            try:
+                if main_task:
+                    main_task.coro.throw(exc)
+            except StopIteration:
+                pass
+            raise
 
         # Get next task to run and continue it
-        t = _task_queue.pop()
+        t = queue_pop()
         cur_task = t
         try:
             # Continue running the coroutine, it's responsible for rescheduling itself
