extmod/uasyncio: Fix bug with task ending just after gather is cancel'd.

This fixes a bug where the gather is cancelled externally and then one of
its sub-tasks (that the gather was waiting on) finishes right between the
cancellation being queued and being executed.

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

3 files changed, 33 insertions, 4 deletions

diff --git a/extmod/uasyncio/funcs.py b/extmod/uasyncio/funcs.py
index b19edeb6e..a76ab0d32 100644
--- a/extmod/uasyncio/funcs.py
+++ b/extmod/uasyncio/funcs.py
@@ -61,9 +61,13 @@ class _Remove:
 
 async def gather(*aws, return_exceptions=False):
     def done(t, er):
+        # Sub-task "t" has finished, with exception "er".
         nonlocal state
-        if type(state) is not int:
-            # A sub-task already raised an exception, so do nothing.
+        if gather_task.data is not _Remove:
+            # The main gather task has already been scheduled, so do nothing.
+            # This happens if another sub-task already raised an exception and
+            # woke the main gather task (via this done function), or if the main
+            # gather task was cancelled externally.
             return
         elif not return_exceptions and not isinstance(er, StopIteration):
             # A sub-task raised an exception, indicate that to the gather task.
diff --git a/tests/extmod/uasyncio_gather.py b/tests/extmod/uasyncio_gather.py
index 366c113a2..fb47e753f 100644
--- a/tests/extmod/uasyncio_gather.py
+++ b/tests/extmod/uasyncio_gather.py
@@ -20,9 +20,9 @@ async def factorial(name, number):
     return f
 
 
-async def task(id):
+async def task(id, t=0.02):
     print("start", id)
-    await asyncio.sleep(0.02)
+    await asyncio.sleep(t)
     print("end", id)
     return id
 
@@ -96,5 +96,14 @@ async def main():
     t.cancel()
     await asyncio.sleep(0.04)
 
+    # Test edge cases where the gather is cancelled just as tasks are created and ending.
+    for i in range(1, 4):
+        print("====")
+        t = asyncio.create_task(gather_task(task(1, t=0), task(2, t=0)))
+        for _ in range(i):
+            await asyncio.sleep(0)
+        t.cancel()
+        await asyncio.sleep(0.2)
+
 
 asyncio.run(main())
diff --git a/tests/extmod/uasyncio_gather.py.exp b/tests/extmod/uasyncio_gather.py.exp
index 74a2ecf75..a5ea47ab5 100644
--- a/tests/extmod/uasyncio_gather.py.exp
+++ b/tests/extmod/uasyncio_gather.py.exp
@@ -36,3 +36,19 @@ True True
 gather_task
 start 1
 start 2
+====
+gather_task
+start 1
+start 2
+====
+gather_task
+start 1
+start 2
+end 1
+end 2
+====
+gather_task
+start 1
+start 2
+end 1
+end 2