extmod/uasyncio: Fix edge case for cancellation of wait_for.

This fixes the cases where the task being waited on finishes just before or
just after the wait_for itself is cancelled.

Fixes issue #8717.

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

1 files changed, 24 insertions, 22 deletions

diff --git a/extmod/uasyncio/funcs.py b/extmod/uasyncio/funcs.py
index 258948f73..a1d38fbcb 100644
--- a/extmod/uasyncio/funcs.py
+++ b/extmod/uasyncio/funcs.py
@@ -1,49 +1,51 @@
 # MicroPython uasyncio module
-# MIT license; Copyright (c) 2019-2020 Damien P. George
+# MIT license; Copyright (c) 2019-2022 Damien P. George
 
 from . import core
 
 
+def _run(waiter, aw):
+    try:
+        result = await aw
+        status = True
+    except BaseException as er:
+        result = None
+        status = er
+    if waiter.data is None:
+        # The waiter is still waiting, cancel it.
+        if waiter.cancel():
+            # Waiter was cancelled by us, change its CancelledError to an instance of
+            # CancelledError that contains the status and result of waiting on aw.
+            # If the wait_for task subsequently gets cancelled externally then this
+            # instance will be reset to a CancelledError instance without arguments.
+            waiter.data = core.CancelledError(status, result)
+
+
 async def wait_for(aw, timeout, sleep=core.sleep):
     aw = core._promote_to_task(aw)
     if timeout is None:
         return await aw
 
-    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))
+    runner_task = core.create_task(_run(core.cur_task, aw))
 
     try:
         # 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:
+        status = er.value
+        if status is None:
             # This wait_for was cancelled externally, so cancel aw and re-raise.
-            status = True
             runner_task.cancel()
             raise er
+        elif status is True:
+            # aw completed successfully and cancelled the sleep, so return aw's result.
+            return er.args[1]
         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