diff --git a/extmod/uasyncio/funcs.py b/extmod/uasyncio/funcs.py index d306752312..93f4fd256c 100644 --- a/extmod/uasyncio/funcs.py +++ b/extmod/uasyncio/funcs.py @@ -9,24 +9,44 @@ async def wait_for(aw, timeout, sleep=core.sleep): if timeout is None: return await aw - def cancel(aw, timeout, sleep): - await sleep(timeout) - aw.cancel() + def runner(waiter, aw): + nonlocal status, result + try: + result = await aw + s = True + except BaseException as er: + s = er + if status is None: + # The waiter is still waiting, set status for it and cancel it. + status = s + waiter.cancel() + + # Run aw in a separate runner task that manages its exceptions. + status = None + result = None + runner_task = core.create_task(runner(core.cur_task, aw)) - cancel_task = core.create_task(cancel(aw, timeout, sleep)) try: - ret = await aw - except core.CancelledError: - # Ignore CancelledError from aw, it's probably due to timeout - pass - finally: - # Cancel the "cancel" task if it's still active (optimisation instead of cancel_task.cancel()) - if cancel_task.coro is not cancel_task: - core._task_queue.remove(cancel_task) - if cancel_task.coro is cancel_task: - # Cancel task ran to completion, ie there was a timeout - raise core.TimeoutError - return ret + # Wait for the timeout to elapse. + await sleep(timeout) + except core.CancelledError as er: + if status is True: + # aw completed successfully and cancelled the sleep, so return aw's result. + return result + elif status is None: + # This wait_for was cancelled externally, so cancel aw and re-raise. + status = True + runner_task.cancel() + raise er + else: + # aw raised an exception, propagate it out to the caller. + raise status + + # The sleep finished before aw, so cancel aw and raise TimeoutError. + status = True + runner_task.cancel() + await runner_task + raise core.TimeoutError def wait_for_ms(aw, timeout): diff --git a/tests/extmod/uasyncio_wait_for.py b/tests/extmod/uasyncio_wait_for.py index 92fd174b84..9612d16204 100644 --- a/tests/extmod/uasyncio_wait_for.py +++ b/tests/extmod/uasyncio_wait_for.py @@ -31,30 +31,85 @@ async def task_raise(): raise ValueError +async def task_cancel_other(t, other): + print("task_cancel_other start") + await asyncio.sleep(t) + print("task_cancel_other cancel") + other.cancel() + + +async def task_wait_for_cancel(id, t, t_wait): + print("task_wait_for_cancel start") + try: + await asyncio.wait_for(task(id, t), t_wait) + except asyncio.CancelledError as er: + print("task_wait_for_cancel cancelled") + raise er + + +async def task_wait_for_cancel_ignore(t_wait): + print("task_wait_for_cancel_ignore start") + try: + await asyncio.wait_for(task_catch(), t_wait) + except asyncio.CancelledError as er: + print("task_wait_for_cancel_ignore cancelled") + raise er + + async def main(): + sep = "-" * 10 + # When task finished before the timeout print(await asyncio.wait_for(task(1, 0.01), 10)) + print(sep) # When timeout passes and task is cancelled try: print(await asyncio.wait_for(task(2, 10), 0.01)) except asyncio.TimeoutError: print("timeout") + print(sep) # When timeout passes and task is cancelled, but task ignores the cancellation request try: print(await asyncio.wait_for(task_catch(), 0.1)) except asyncio.TimeoutError: print("TimeoutError") + print(sep) # When task raises an exception try: print(await asyncio.wait_for(task_raise(), 1)) except ValueError: print("ValueError") + print(sep) # Timeout of None means wait forever print(await asyncio.wait_for(task(3, 0.1), None)) + print(sep) + + # When task is cancelled by another task + t = asyncio.create_task(task(4, 10)) + asyncio.create_task(task_cancel_other(0.01, t)) + try: + print(await asyncio.wait_for(t, 1)) + except asyncio.CancelledError as er: + print(repr(er)) + print(sep) + + # When wait_for gets cancelled + t = asyncio.create_task(task_wait_for_cancel(4, 1, 2)) + await asyncio.sleep(0.01) + t.cancel() + await asyncio.sleep(0.01) + print(sep) + + # When wait_for gets cancelled and awaited task ignores the cancellation request + t = asyncio.create_task(task_wait_for_cancel_ignore(2)) + await asyncio.sleep(0.01) + t.cancel() + await asyncio.sleep(0.01) + print(sep) print("finish") diff --git a/tests/extmod/uasyncio_wait_for.py.exp b/tests/extmod/uasyncio_wait_for.py.exp index 41a5f2481e..a4201d31ff 100644 --- a/tests/extmod/uasyncio_wait_for.py.exp +++ b/tests/extmod/uasyncio_wait_for.py.exp @@ -1,15 +1,35 @@ task start 1 task end 1 2 +---------- task start 2 timeout +---------- task_catch start ignore cancel task_catch done TimeoutError +---------- task start ValueError +---------- task start 3 task end 3 6 +---------- +task start 4 +task_cancel_other start +task_cancel_other cancel +CancelledError() +---------- +task_wait_for_cancel start +task start 4 +task_wait_for_cancel cancelled +---------- +task_wait_for_cancel_ignore start +task_catch start +task_wait_for_cancel_ignore cancelled +ignore cancel +task_catch done +---------- finish diff --git a/tests/extmod/uasyncio_wait_for_fwd.py b/tests/extmod/uasyncio_wait_for_fwd.py new file mode 100644 index 0000000000..33738085ce --- /dev/null +++ b/tests/extmod/uasyncio_wait_for_fwd.py @@ -0,0 +1,60 @@ +# Test asyncio.wait_for, with forwarding cancellation + +try: + import uasyncio as asyncio +except ImportError: + try: + import asyncio + except ImportError: + print("SKIP") + raise SystemExit + + +async def awaiting(t, return_if_fail): + try: + print("awaiting started") + await asyncio.sleep(t) + except asyncio.CancelledError as er: + # CPython wait_for raises CancelledError inside task but TimeoutError in wait_for + print("awaiting canceled") + if return_if_fail: + return False # return has no effect if Cancelled + else: + raise er + except Exception as er: + print("caught exception", er) + raise er + + +async def test_cancellation_forwarded(catch, catch_inside): + print("----------") + + async def wait(): + try: + await asyncio.wait_for(awaiting(2, catch_inside), 1) + except asyncio.TimeoutError as er: + print("Got timeout error") + raise er + except asyncio.CancelledError as er: + print("Got canceled") + if not catch: + raise er + + async def cancel(t): + print("cancel started") + await asyncio.sleep(0.01) + print("cancel wait()") + t.cancel() + + t = asyncio.create_task(wait()) + k = asyncio.create_task(cancel(t)) + try: + await t + except asyncio.CancelledError: + print("waiting got cancelled") + + +asyncio.run(test_cancellation_forwarded(False, False)) +asyncio.run(test_cancellation_forwarded(False, True)) +asyncio.run(test_cancellation_forwarded(True, True)) +asyncio.run(test_cancellation_forwarded(True, False)) diff --git a/tests/extmod/uasyncio_wait_for_fwd.py.exp b/tests/extmod/uasyncio_wait_for_fwd.py.exp new file mode 100644 index 0000000000..9f22f1a7d1 --- /dev/null +++ b/tests/extmod/uasyncio_wait_for_fwd.py.exp @@ -0,0 +1,26 @@ +---------- +cancel started +awaiting started +cancel wait() +Got canceled +awaiting canceled +waiting got cancelled +---------- +cancel started +awaiting started +cancel wait() +Got canceled +awaiting canceled +waiting got cancelled +---------- +cancel started +awaiting started +cancel wait() +Got canceled +awaiting canceled +---------- +cancel started +awaiting started +cancel wait() +Got canceled +awaiting canceled diff --git a/tests/run-tests b/tests/run-tests index 23c30d3c3c..cb5b5cd0a5 100755 --- a/tests/run-tests +++ b/tests/run-tests @@ -434,7 +434,10 @@ def run_tests(pyb, tests, args, result_dir): skip_tests.add('basics/scope_implicit.py') # requires checking for unbound local skip_tests.add('basics/try_finally_return2.py') # requires raise_varargs skip_tests.add('basics/unboundlocal.py') # requires checking for unbound local + skip_tests.add('extmod/uasyncio_event.py') # unknown issue skip_tests.add('extmod/uasyncio_lock.py') # requires async with + skip_tests.add('extmod/uasyncio_micropython.py') # unknown issue + skip_tests.add('extmod/uasyncio_wait_for.py') # unknown issue skip_tests.add('misc/features.py') # requires raise_varargs skip_tests.add('misc/print_exception.py') # because native doesn't have proper traceback info skip_tests.add('misc/sys_exc_info.py') # sys.exc_info() is not supported for native