extmod/uasyncio: Delay calling Loop.call_exception_handler by 1 loop.

When a tasks raises an exception which is uncaught, and no other task
await's on that task, then an error message is printed (or a user function
called) via a call to Loop.call_exception_handler.  In CPython this call is
made when the Task object is freed (eg via reference counting) because it's
at that point that it is known that the exception that was raised will
never be handled.

MicroPython does not have reference counting and the current behaviour is
to deal with uncaught exceptions as early as possible, ie as soon as they
terminate the task.  But this can be undesirable because in certain cases
a task can start and raise an exception immediately (before any await is
executed in that task's coro) and before any other task gets a chance to
await on it to catch the exception.

This commit changes the behaviour so that tasks which end due to an
uncaught exception are scheduled one more time for execution, and if they
are not await'ed on by the next scheduling loop, then the exception handler
is called (eg the exception is printed out).

Signed-off-by: Damien George <damien@micropython.org>

extmod/moduasyncio.c

@@ -146,6 +146,9 @@ STATIC const mp_obj_type_t task_queue_type = {
 /******************************************************************************/
 // Task class
 
+// For efficiency, the task object is stored to the coro entry when the task is done.
+#define TASK_IS_DONE(task) ((task)->coro == MP_OBJ_FROM_PTR(task))
+
 // This is the core uasyncio context with cur_task, _task_queue and CancelledError.
 STATIC mp_obj_t uasyncio_context = MP_OBJ_NULL;
 
@@ -167,7 +170,7 @@ STATIC mp_obj_t task_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
 STATIC mp_obj_t task_cancel(mp_obj_t self_in) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     // Check if task is already finished.
-    if (self->coro == mp_const_none) {
+    if (TASK_IS_DONE(self)) {
         return mp_const_false;
     }
     // Can't cancel self (not supported yet).
@@ -209,6 +212,24 @@ STATIC mp_obj_t task_cancel(mp_obj_t self_in) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(task_cancel_obj, task_cancel);
 
+STATIC mp_obj_t task_throw(mp_obj_t self_in, mp_obj_t value_in) {
+    // This task raised an exception which was uncaught; handle that now.
+    mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
+    // Set the data because it was cleared by the main scheduling loop.
+    self->data = value_in;
+    if (self->waiting == mp_const_none) {
+        // Nothing await'ed on the task so call the exception handler.
+        mp_obj_t _exc_context = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR__exc_context));
+        mp_obj_dict_store(_exc_context, MP_OBJ_NEW_QSTR(MP_QSTR_exception), value_in);
+        mp_obj_dict_store(_exc_context, MP_OBJ_NEW_QSTR(MP_QSTR_future), self_in);
+        mp_obj_t Loop = mp_obj_dict_get(uasyncio_context, MP_OBJ_NEW_QSTR(MP_QSTR_Loop));
+        mp_obj_t call_exception_handler = mp_load_attr(Loop, MP_QSTR_call_exception_handler);
+        mp_call_function_1(call_exception_handler, _exc_context);
+    }
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_2(task_throw_obj, task_throw);
+
 STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     if (dest[0] == MP_OBJ_NULL) {
@@ -218,12 +239,15 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
         } else if (attr == MP_QSTR_data) {
             dest[0] = self->data;
         } else if (attr == MP_QSTR_waiting) {
-            if (self->waiting != mp_const_none) {
+            if (self->waiting != mp_const_none && self->waiting != mp_const_false) {
                 dest[0] = self->waiting;
             }
         } else if (attr == MP_QSTR_cancel) {
             dest[0] = MP_OBJ_FROM_PTR(&task_cancel_obj);
             dest[1] = self_in;
+        } else if (attr == MP_QSTR_throw) {
+            dest[0] = MP_OBJ_FROM_PTR(&task_throw_obj);
+            dest[1] = self_in;
         } else if (attr == MP_QSTR_ph_key) {
             dest[0] = self->ph_key;
         }
@@ -246,14 +270,21 @@ STATIC mp_obj_t task_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
     (void)iter_buf;
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
     if (self->waiting == mp_const_none) {
-        self->waiting = task_queue_make_new(&task_queue_type, 0, 0, NULL);
+        // The is the first access of the "waiting" entry.
+        if (TASK_IS_DONE(self)) {
+            // Signal that the completed-task has been await'ed on.
+            self->waiting = mp_const_false;
+        } else {
+            // Lazily allocate the waiting queue.
+            self->waiting = task_queue_make_new(&task_queue_type, 0, 0, NULL);
+        }
     }
     return self_in;
 }
 
 STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
     mp_obj_task_t *self = MP_OBJ_TO_PTR(self_in);
-    if (self->coro == mp_const_none) {
+    if (TASK_IS_DONE(self)) {
         // Task finished, raise return value to caller so it can continue.
         nlr_raise(self->data);
     } else {

extmod/uasyncio/core.py

@@ -185,8 +185,6 @@ def run_until_complete(main_task=None):
                 if isinstance(er, StopIteration):
                     return er.value
                 raise er
-            # Save return value of coro to pass up to caller
-            t.data = er
             # Schedule any other tasks waiting on the completion of this task
             waiting = False
             if hasattr(t, "waiting"):
@@ -194,13 +192,15 @@ def run_until_complete(main_task=None):
                     _task_queue.push_head(t.waiting.pop_head())
                     waiting = True
                 t.waiting = None  # Free waiting queue head
-            # Print out exception for detached tasks
             if not waiting and not isinstance(er, excs_stop):
-                _exc_context["exception"] = er
-                _exc_context["future"] = t
-                Loop.call_exception_handler(_exc_context)
-            # Indicate task is done
-            t.coro = None
+                # 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_head(t)
+            # Indicate task is done by setting coro to the task object itself
+            t.coro = t
+            # Save return value of coro to pass up to caller
+            t.data = er
 
 
 # Create a new task from a coroutine and run it until it finishes

extmod/uasyncio/funcs.py

@@ -21,9 +21,9 @@ async def wait_for(aw, timeout, sleep=core.sleep):
         pass
     finally:
         # Cancel the "cancel" task if it's still active (optimisation instead of cancel_task.cancel())
-        if cancel_task.coro is not None:
+        if cancel_task.coro is not cancel_task:
             core._task_queue.remove(cancel_task)
-    if cancel_task.coro is None:
+    if cancel_task.coro is cancel_task:
         # Cancel task ran to completion, ie there was a timeout
         raise core.TimeoutError
     return ret

extmod/uasyncio/task.py

@@ -130,13 +130,16 @@ class Task:
         self.ph_rightmost_parent = None  # Paring heap
 
     def __iter__(self):
-        if not hasattr(self, "waiting"):
+        if self.coro is self:
+            # Signal that the completed-task has been await'ed on.
+            self.waiting = None
+        elif not hasattr(self, "waiting"):
             # Lazily allocated head of linked list of Tasks waiting on completion of this task.
             self.waiting = TaskQueue()
         return self
 
     def __next__(self):
-        if not self.coro:
+        if self.coro is self:
             # Task finished, raise return value to caller so it can continue.
             raise self.data
         else:
@@ -147,7 +150,7 @@ class Task:
 
     def cancel(self):
         # Check if task is already finished.
-        if self.coro is None:
+        if self.coro is self:
             return False
         # Can't cancel self (not supported yet).
         if self is core.cur_task:
@@ -166,3 +169,13 @@ class Task:
             core._task_queue.push_head(self)
         self.data = core.CancelledError
         return True
+
+    def throw(self, value):
+        # This task raised an exception which was uncaught; handle that now.
+        # Set the data because it was cleared by the main scheduling loop.
+        self.data = value
+        if not hasattr(self, "waiting"):
+            # Nothing await'ed on the task so call the exception handler.
+            core._exc_context["exception"] = value
+            core._exc_context["future"] = self
+            core.Loop.call_exception_handler(core._exc_context)

tests/extmod/uasyncio_set_exception_handler.py

@@ -32,13 +32,23 @@ async def main():
     # Create a task that raises and uses the custom exception handler
     asyncio.create_task(task(0))
     print("sleep")
-    await asyncio.sleep(0)
+    for _ in range(2):
+        await asyncio.sleep(0)
 
     # Create 2 tasks to test order of printing exception
     asyncio.create_task(task(1))
     asyncio.create_task(task(2))
     print("sleep")
+    for _ in range(2):
+        await asyncio.sleep(0)
+
+    # Create a task, let it run, then await it (no exception should be printed)
+    t = asyncio.create_task(task(3))
     await asyncio.sleep(0)
+    try:
+        await t
+    except ValueError as er:
+        print(repr(er))
 
     print("done")
 

tests/extmod/uasyncio_set_exception_handler.py.exp

@@ -5,4 +5,5 @@ custom_handler ValueError(0, 1)
 sleep
 custom_handler ValueError(1, 2)
 custom_handler ValueError(2, 3)
+ValueError(3, 4)
 done