extmod/uasyncio: Fix race with cancelled task waiting on finished task.

This commit fixes a problem with a race between cancellation of task A and
completion of task B, when A waits on B.  If task B completes just before
task A is cancelled then the cancellation of A does not work.  Instead,
the CancelledError meant to cancel A gets passed through to B (that's
expected behaviour) but B handles it as a "Task exception wasn't retrieved"
scenario, printing out such a message (this is because finished tasks point
their "coro" attribute to themselves to indicate they are done, and
implement the throw() method, but that method inadvertently catches the
CancelledError).  The correct behaviour is for B to bounce that
CancelledError back out.

This bug is mainly seen when wait_for() is used, and in that context the
symptoms are:
- occurs when using wait_for(T, S), if the task T being waited on finishes
  at exactly the same time as the wait-for timeout S expires
- task T will have run to completion
- the "Task exception wasn't retrieved message" is printed with
  "<class 'CancelledError'>" as the error (ie no traceback)
- the wait_for(T, S) call never returns (it's never put back on the
  uasyncio run queue) and all tasks waiting on this are blocked forever
  from running
- uasyncio otherwise continues to function and other tasks continue to be
  scheduled as normal

The fix here reworks the "waiting" attribute of Task to be called "state"
and uses it to indicate whether a task is: running and not awaited on,
running and awaited on, finished and not awaited on, or finished and
awaited on.  This means the task does not need to point "coro" to itself to
indicate finished, and also allows removal of the throw() method.

A benefit of this is that "Task exception wasn't retrieved" messages can go
back to being able to print the name of the coroutine function.

Fixes issue #7386.

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

1 files changed, 22 insertions, 47 deletions

diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c
index fe0f748ca..9717e3856 100644
--- a/extmod/moduasyncio.c
+++ b/extmod/moduasyncio.c
@@ -31,12 +31,19 @@
 
 #if MICROPY_PY_UASYNCIO
 
+#define TASK_STATE_RUNNING_NOT_WAITED_ON (mp_const_true)
+#define TASK_STATE_DONE_NOT_WAITED_ON (mp_const_none)
+#define TASK_STATE_DONE_WAS_WAITED_ON (mp_const_false)
+
+#define TASK_IS_DONE(task) ( \
+    (task)->state == TASK_STATE_DONE_NOT_WAITED_ON \
+    || (task)->state == TASK_STATE_DONE_WAS_WAITED_ON)
+
 typedef struct _mp_obj_task_t {
     mp_pairheap_t pairheap;
     mp_obj_t coro;
     mp_obj_t data;
-    mp_obj_t waiting;
-
+    mp_obj_t state;
     mp_obj_t ph_key;
 } mp_obj_task_t;
 
@@ -146,9 +153,6 @@ 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;
 
@@ -159,7 +163,7 @@ STATIC mp_obj_t task_make_new(const mp_obj_type_t *type, size_t n_args, size_t n
     mp_pairheap_init_node(task_lt, &self->pairheap);
     self->coro = args[0];
     self->data = mp_const_none;
-    self->waiting = mp_const_none;
+    self->state = TASK_STATE_RUNNING_NOT_WAITED_ON;
     self->ph_key = MP_OBJ_NEW_SMALL_INT(0);
     if (n_args == 2) {
         uasyncio_context = args[1];
@@ -218,24 +222,6 @@ 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) {
@@ -244,32 +230,24 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
             dest[0] = self->coro;
         } else if (attr == MP_QSTR_data) {
             dest[0] = self->data;
-        } else if (attr == MP_QSTR_waiting) {
-            if (self->waiting != mp_const_none && self->waiting != mp_const_false) {
-                dest[0] = self->waiting;
-            }
+        } else if (attr == MP_QSTR_state) {
+            dest[0] = self->state;
         } else if (attr == MP_QSTR_done) {
             dest[0] = MP_OBJ_FROM_PTR(&task_done_obj);
             dest[1] = self_in;
         } 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;
         }
     } else if (dest[1] != MP_OBJ_NULL) {
         // Store
-        if (attr == MP_QSTR_coro) {
-            self->coro = dest[1];
-            dest[0] = MP_OBJ_NULL;
-        } else if (attr == MP_QSTR_data) {
+        if (attr == MP_QSTR_data) {
             self->data = dest[1];
             dest[0] = MP_OBJ_NULL;
-        } else if (attr == MP_QSTR_waiting) {
-            self->waiting = dest[1];
+        } else if (attr == MP_QSTR_state) {
+            self->state = dest[1];
             dest[0] = MP_OBJ_NULL;
         }
     }
@@ -278,15 +256,12 @@ STATIC void task_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
 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) {
-        // 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);
-        }
+    if (TASK_IS_DONE(self)) {
+        // Signal that the completed-task has been await'ed on.
+        self->state = TASK_STATE_DONE_WAS_WAITED_ON;
+    } else if (self->state == TASK_STATE_RUNNING_NOT_WAITED_ON) {
+        // Allocate the waiting queue.
+        self->state = task_queue_make_new(&task_queue_type, 0, 0, NULL);
     }
     return self_in;
 }
@@ -299,7 +274,7 @@ STATIC mp_obj_t task_iternext(mp_obj_t self_in) {
     } 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));
-        mp_obj_t args[2] = { self->waiting, cur_task };
+        mp_obj_t args[2] = { self->state, cur_task };
         task_queue_push_sorted(2, args);
         // Set calling task's data to this task that it waits on, to double-link it.
         ((mp_obj_task_t *)MP_OBJ_TO_PTR(cur_task))->data = self_in;