diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 53eef84427be10..a4d8b3f24fe287 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -439,7 +439,7 @@ async def wait_for(fut, timeout): # after wait_for() returns. # See https://bugs.python.org/issue32751 await _cancel_and_wait(fut, loop=loop) - raise + return fut.result() 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 a88cb89e4e6eef..5879dd128d1016 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1009,16 +1009,63 @@ def gen(): 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) + # Test that if task is cancelled at the same time the future + # completes, the result is still returned. + 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_suppress_cancellation(self): + def gen(): + yield + yield 0.1 + yield 0.1 + yield 0.1 + async def inner(): - with contextlib.suppress(asyncio.CancelledError): + try: await asyncio.sleep(1) - return 1 + except asyncio.CancelledError: + return "ok" - async def main(): - result = await asyncio.wait_for(inner(), timeout=.01) - assert result == 1 + loop = self.new_test_loop(gen) + + fut = self.new_future(loop) + task = loop.create_task(asyncio.wait_for(inner(), timeout=1)) + loop.call_later(0.1, task.cancel) + # Cancellation is suppressed in inner(), so should still return here. + res = loop.run_until_complete(task) + self.assertEqual(res, "ok") + + def test_wait_for_cancellation_inner_race_condition(self): + def gen(): + yield + yield 0.1 + yield 0 + yield 0 + + async def inner(): + # Test that if fut completes at same time as timeout, the result + # is still returned. + return await asyncio.wait_for(fut, timeout=1) - asyncio.run(main()) + loop = self.new_test_loop(gen) + + fut = self.new_future(loop) + loop.call_later(0.1, fut.set_result, "ok") + res = loop.run_until_complete(asyncio.wait_for(inner(), timeout=0.1)) + self.assertEqual(res, "ok") def test_wait_for_waits_for_task_cancellation(self): loop = asyncio.new_event_loop() diff --git a/Misc/NEWS.d/next/Library/2021-11-29-17-17-45.bpo-42130.Ho8dUh.rst b/Misc/NEWS.d/next/Library/2021-11-29-17-17-45.bpo-42130.Ho8dUh.rst new file mode 100644 index 00000000000000..ee837d5636cfcf --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-11-29-17-17-45.bpo-42130.Ho8dUh.rst @@ -0,0 +1,2 @@ +Fix issue where ``asyncio.wait_for()`` is cancelled when the coro +actually suppresses the cancellation.