From 2ecbad4e91192c88f831d8690dbad31ddba72135 Mon Sep 17 00:00:00 2001 From: Damien George Date: Fri, 19 Jan 2024 10:07:51 +1100 Subject: [PATCH] extmod/asyncio: Support gather of tasks that finish early. Adds support to asyncio.gather() for the case that one or more (or all) sub-tasks finish and/or raise an exception before the gather starts. Signed-off-by: Damien George --- extmod/asyncio/core.py | 5 ++ extmod/asyncio/funcs.py | 49 +++++++++----- tests/extmod/asyncio_gather_finished_early.py | 65 +++++++++++++++++++ 3 files changed, 102 insertions(+), 17 deletions(-) create mode 100644 tests/extmod/asyncio_gather_finished_early.py diff --git a/extmod/asyncio/core.py b/extmod/asyncio/core.py index 214cc52f45f53..e5af3038f77bb 100644 --- a/extmod/asyncio/core.py +++ b/extmod/asyncio/core.py @@ -219,6 +219,11 @@ def run_until_complete(main_task=None): elif t.state is None: # Task is already finished and nothing await'ed on the task, # so call the exception handler. + + # Save exception raised by the coro for later use. + t.data = exc + + # Create exception context and call the exception handler. _exc_context["exception"] = exc _exc_context["future"] = t Loop.call_exception_handler(_exc_context) diff --git a/extmod/asyncio/funcs.py b/extmod/asyncio/funcs.py index 599091dfbdc2b..3ef8a76b1d151 100644 --- a/extmod/asyncio/funcs.py +++ b/extmod/asyncio/funcs.py @@ -63,9 +63,6 @@ def remove(t): # async def gather(*aws, return_exceptions=False): - if not aws: - return [] - def done(t, er): # Sub-task "t" has finished, with exception "er". nonlocal state @@ -86,26 +83,39 @@ def done(t, er): # Gather waiting is done, schedule the main gather task. core._task_queue.push(gather_task) + # Prepare the sub-tasks for the gather. + # The `state` variable counts the number of tasks to wait for, and can be negative + # if the gather should not run at all (because a task already had an exception). ts = [core._promote_to_task(aw) for aw in aws] + state = 0 for i in range(len(ts)): - if ts[i].state is not True: - # Task is not running, gather not currently supported for this case. + if ts[i].state is True: + # Task is running, register the callback to call when the task is done. + ts[i].state = done + state += 1 + elif not ts[i].state: + # Task finished already. + if not isinstance(ts[i].data, StopIteration): + # Task finished by raising an exception. + if not return_exceptions: + # Do not run this gather at all. + state = -len(ts) + else: + # Task being waited on, gather not currently supported for this case. raise RuntimeError("can't gather") - # Register the callback to call when the task is done. - ts[i].state = done # Set the state for execution of the gather. gather_task = core.cur_task - state = len(ts) cancel_all = False - # Wait for the a sub-task to need attention. - gather_task.data = _Remove - try: - yield - except core.CancelledError as er: - cancel_all = True - state = er + # Wait for a sub-task to need attention (if there are any to wait for). + if state > 0: + gather_task.data = _Remove + try: + yield + except core.CancelledError as er: + cancel_all = True + state = er # Clean up tasks. for i in range(len(ts)): @@ -118,8 +128,13 @@ def done(t, er): # Sub-task ran to completion, get its return value. ts[i] = ts[i].data.value else: - # Sub-task had an exception with return_exceptions==True, so get its exception. - ts[i] = ts[i].data + # Sub-task had an exception. + if return_exceptions: + # Get the sub-task exception to return in the list of return values. + ts[i] = ts[i].data + elif isinstance(state, int): + # Raise the sub-task exception, if there is not already an exception to raise. + state = ts[i].data # Either this gather was cancelled, or one of the sub-tasks raised an exception with # return_exceptions==False, so reraise the exception here. diff --git a/tests/extmod/asyncio_gather_finished_early.py b/tests/extmod/asyncio_gather_finished_early.py new file mode 100644 index 0000000000000..030e79e3571b5 --- /dev/null +++ b/tests/extmod/asyncio_gather_finished_early.py @@ -0,0 +1,65 @@ +# Test asyncio.gather() when a task is already finished before the gather starts. + +try: + import asyncio +except ImportError: + print("SKIP") + raise SystemExit + + +# CPython and MicroPython differ in when they signal (and print) that a task raised an +# uncaught exception. So define an empty custom_handler() to suppress this output. +def custom_handler(loop, context): + pass + + +async def task_that_finishes_early(id, event, fail): + print("task_that_finishes_early", id) + event.set() + if fail: + raise ValueError("intentional exception", id) + + +async def task_that_runs(): + for i in range(5): + print("task_that_runs", i) + await asyncio.sleep(0) + + +async def main(start_task_that_runs, task_fail, return_exceptions): + print("== start", start_task_that_runs, task_fail, return_exceptions) + + # Set exception handler to suppress exception output. + loop = asyncio.get_event_loop() + loop.set_exception_handler(custom_handler) + + # Create tasks. + event_a = asyncio.Event() + event_b = asyncio.Event() + tasks = [] + if start_task_that_runs: + tasks.append(asyncio.create_task(task_that_runs())) + tasks.append(asyncio.create_task(task_that_finishes_early("a", event_a, task_fail))) + tasks.append(asyncio.create_task(task_that_finishes_early("b", event_b, task_fail))) + + # Make sure task_that_finishes_early() are both done, before calling gather(). + await event_a.wait() + await event_b.wait() + + # Gather the tasks. + try: + result = "complete", await asyncio.gather(*tasks, return_exceptions=return_exceptions) + except Exception as er: + result = "exception", er, start_task_that_runs and tasks[0].done() + + # Wait for the final task to finish (if it was started). + if start_task_that_runs: + await tasks[0] + + # Print results. + print(result) + + +# Run the test in the 8 different combinations of its arguments. +for i in range(8): + asyncio.run(main(bool(i & 4), bool(i & 2), bool(i & 1))) pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy