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>

1 files changed, 6 insertions, 2 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.