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>
This commit is contained in:
Damien George 2022-04-21 13:19:34 +10:00
parent f7454f850f
commit 28e7e15c0a
3 changed files with 33 additions and 4 deletions

View File

@ -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.

View File

@ -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())

View File

@ -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