diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 7ecec963848..8b05434f273 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -465,9 +465,12 @@ async def wait_for(fut, timeout, *, loop=None): try: await waiter except exceptions.CancelledError: - fut.remove_done_callback(cb) - fut.cancel() - raise + if fut.done(): + return fut.result() + else: + fut.remove_done_callback(cb) + fut.cancel() + raise if fut.done(): return fut.result() diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 511961c3200..74fc1e4a421 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1120,6 +1120,22 @@ class BaseTaskTests: res = loop.run_until_complete(task) self.assertEqual(res, "ok") + def test_wait_for_cancellation_race_condition(self): + def gen(): + yield 0.1 + yield 0.1 + yield 0.1 + yield 0.1 + + loop = self.new_test_loop(gen) + + fut = self.new_future(loop) + loop.call_later(0.1, fut.set_result, "ok") + task = loop.create_task(asyncio.wait_for(fut, timeout=1)) + loop.call_later(0.1, task.cancel) + res = loop.run_until_complete(task) + self.assertEqual(res, "ok") + def test_wait_for_waits_for_task_cancellation(self): loop = asyncio.new_event_loop() self.addCleanup(loop.close) diff --git a/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst b/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst new file mode 100644 index 00000000000..694fbbbf346 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst @@ -0,0 +1,2 @@ +:meth:`asyncio.wait_for` now properly handles races between cancellation of +itself and the completion of the wrapped awaitable.