diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 9a9d0d6e3cc269..6e5c0eec12b5df 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -433,15 +433,12 @@ async def wait_for(fut, timeout): try: await waiter except exceptions.CancelledError: - if fut.done(): - return fut.result() - else: - fut.remove_done_callback(cb) - # We must ensure that the task is not running - # after wait_for() returns. - # See https://bugs.python.org/issue32751 - await _cancel_and_wait(fut, loop=loop) - raise + fut.remove_done_callback(cb) + # We must ensure that the task is not running + # after wait_for() returns. + # See https://bugs.python.org/issue32751 + await _cancel_and_wait(fut, loop=loop) + raise if fut.done(): return fut.result() diff --git a/Lib/test/test_asyncio/test_asyncio_waitfor.py b/Lib/test/test_asyncio/test_asyncio_waitfor.py index 2ca64abbeb527c..fc5a4218ac6999 100644 --- a/Lib/test/test_asyncio/test_asyncio_waitfor.py +++ b/Lib/test/test_asyncio/test_asyncio_waitfor.py @@ -28,6 +28,7 @@ async def run(self): self.exited = True + class AsyncioWaitForTest(unittest.TestCase): async def atest_asyncio_wait_for_cancelled(self): @@ -56,6 +57,38 @@ async def atest_asyncio_wait_for_timeout(self): def test_asyncio_wait_for_timeout(self): asyncio.run(self.atest_asyncio_wait_for_timeout()) + async def atest_asyncio_wait_for_hang(self, inner_steps, outer_steps): + # bpo-42130: wait_for can swallow cancellation causing task to hang + async def inner(steps): + for _ in range(steps): + await asyncio.sleep(0) + + finished = False + + async def wait_for_coro(steps): + nonlocal finished + await asyncio.wait_for(inner(steps), timeout=1) + await asyncio.sleep(0.1) + finished = True + + task = asyncio.create_task(wait_for_coro(inner_steps)) + for _ in range(outer_steps): + await asyncio.sleep(0) + assert not task.done() + + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertFalse(finished) + + def test_asyncio_wait_for_hang(self): + # Test with different number of inner/outer steps to weaken dependence + # on implementation details + for inner_steps in range(3): + for outer_steps in range(3): + with self.subTest(inner_steps=inner_steps, outer_steps=outer_steps): + asyncio.run(self.atest_asyncio_wait_for_hang(inner_steps, outer_steps)) + if __name__ == '__main__': unittest.main() diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index a9e4cf53566ca9..f84b9c97a559ea 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -1147,22 +1147,6 @@ def gen(): 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/2021-05-13-11-16-06.bpo-42130.ez0TO5.rst b/Misc/NEWS.d/next/Library/2021-05-13-11-16-06.bpo-42130.ez0TO5.rst new file mode 100644 index 00000000000000..f632aa99b923d2 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2021-05-13-11-16-06.bpo-42130.ez0TO5.rst @@ -0,0 +1 @@ +:meth:`asyncio.wait_for` fix swallowing of cancellation \ No newline at end of file