extmod/uasyncio/task.py: Fix crash when non-awaited task is awaited.

A task that has been sent to the loop's exception handler due to being
re-scheduled twice will then subsequently cause a `raise None` if it is
subsequently awaited. In the C version of task.py, this causes a segfault.

This makes the await succeed (via raising StopIteration instead).

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

Author: jimmo

extmod/moduasyncio.c

@@ -290,8 +290,13 @@ STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
 STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     if (TASK_IS_DONE(self)) {
-        // Task finished, raise return value to caller so it can continue.
-        nlr_raise(self->data);
+        if (self->data == mp_const_none) {
+            // Task finished but has already been sent to the loop's exception handler.
+            mp_raise_StopIteration(MP_OBJ_NULL);
+        } else {
+            // Task finished, raise return value to caller so it can continue.
+            nlr_raise(self->data);
+        }
     } else {
         // Put calling task on waiting queue.
         mp_obj_t cur_task = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_cur_task));

extmod/uasyncio/task.py

@@ -141,8 +141,12 @@ def __await__(self):
 
     def __next__(self):
         if not self.state:
-            # Task finished, raise return value to caller so it can continue.
-            raise self.data
+            if self.data is None:
+                # Task finished but has already been sent to the loop's exception handler.
+                raise StopIteration
+            else:
+                # Task finished, raise return value to caller so it can continue.
+                raise self.data
         else:
             # Put calling task on waiting queue.
             self.state.push_head(core.cur_task)

tests/extmod/uasyncio_task_exception.py

@@ -0,0 +1,42 @@
+# In MicroPython, a non-awaited task with a pending exception will raise to
+# the loop's exception handler the second time it is scheduled. This is
+# because without reference counting we have no way to know when the task is
+# truly "non awaited" -- i.e. we only know that it wasn't awaited in the time
+# it took to be re-scheduled.
+
+# If the task _is_ subsequently awaited, then the await should succeed without
+# raising.
+
+try:
+    import uasyncio as asyncio
+except ImportError:
+    try:
+        import asyncio
+    except ImportError:
+        print("SKIP")
+        raise SystemExit
+
+
+def custom_handler(loop, context):
+    print("exception handler", type(context["exception"]).__name__)
+
+
+async def main():
+    loop = asyncio.get_event_loop()
+    loop.set_exception_handler(custom_handler)
+
+    async def task():
+        print("raise")
+        raise OSError
+
+    print("create")
+    t = asyncio.create_task(task())
+    print("sleep 1")
+    await asyncio.sleep(0)
+    print("sleep 2")
+    await asyncio.sleep(0)
+    print("await")
+    await t  # should not raise.
+
+
+asyncio.run(main())

tests/extmod/uasyncio_task_exception.py.exp

@@ -0,0 +1,6 @@
+create
+sleep 1
+raise
+sleep 2
+exception handler OSError
+await