diff --git a/extmod/moduasyncio.c b/extmod/moduasyncio.c index 0b15c9e079..e8822c0697 100644 --- a/extmod/moduasyncio.c +++ b/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 { diff --git a/extmod/uasyncio/core.py b/extmod/uasyncio/core.py index 045b4cd139..6a84b0982c 100644 --- a/extmod/uasyncio/core.py +++ b/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 diff --git a/extmod/uasyncio/funcs.py b/extmod/uasyncio/funcs.py index 6e1305c94f..d306752312 100644 --- a/extmod/uasyncio/funcs.py +++ b/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 diff --git a/extmod/uasyncio/task.py b/extmod/uasyncio/task.py index 1788cf0ed0..2420ab7193 100644 --- a/extmod/uasyncio/task.py +++ b/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) diff --git a/tests/extmod/uasyncio_set_exception_handler.py b/tests/extmod/uasyncio_set_exception_handler.py index ad62a79b7b..fe7b83eb4d 100644 --- a/tests/extmod/uasyncio_set_exception_handler.py +++ b/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") diff --git a/tests/extmod/uasyncio_set_exception_handler.py.exp b/tests/extmod/uasyncio_set_exception_handler.py.exp index 4744641e54..fb4711469d 100644 --- a/tests/extmod/uasyncio_set_exception_handler.py.exp +++ b/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