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>

2 files changed, 27 insertions, 2 deletions

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