diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index fa853283c0c5e4..74b0bf0713ad66 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -468,15 +468,13 @@ async def wait_for(fut, timeout): try: await waiter except exceptions.CancelledError: - if fut.done(): - return fut.result() - else: - fut.remove_done_callback(cb) + if not fut.done(): # We must ensure that the task is not running # after wait_for() returns. # See https://bugs.python.org/issue32751 + fut.remove_done_callback(cb) await _cancel_and_wait(fut, loop=loop) - raise + raise if fut.done(): return fut.result() diff --git a/Lib/test/test_asyncio/test_waitfor.py b/Lib/test/test_asyncio/test_waitfor.py index 45498fa097f6bc..a428ec2d9a0fd6 100644 --- a/Lib/test/test_asyncio/test_waitfor.py +++ b/Lib/test/test_asyncio/test_waitfor.py @@ -289,6 +289,60 @@ async def test_cancel_blocking_wait_for(self): async def test_cancel_wait_for(self): await self._test_cancel_wait_for(60.0) + async def simultaneous_self_cancel_and_inner_result( + self, + waitfor_timeout, + inner_action, + ): + """Construct scenario where external cancellation and + awaitable becoming done happen simultaneously. + inner_acion is one of 'cancel', 'exception' or 'result'. + Make sure waitfor_timeout > 0.1. Trying to make it == 0.1 is not + a reliable way to make timeout happen at the same time as the above. + """ + loop = asyncio.get_running_loop() + inner = loop.create_future() + waitfor_task = asyncio.create_task( + asyncio.wait_for(inner, timeout=waitfor_timeout)) + await asyncio.sleep(0.1) + # Even if waitfor_timeout == 0.1, there's still no guarantee whether the + # timer handler (or similar) in wait_for() or code below in this + # coroutine executes first. If the timer handler executes first, then + # inner will be cancelled(), and code below will raise + # InvalidStateError. + if inner_action == 'cancel': + inner.cancel() + elif inner_action == 'exception': + inner.set_exception(RuntimeError('inner exception')) + else: + assert inner_action == 'result' + inner.set_result('inner result') + waitfor_task.cancel() + with self.assertRaises(asyncio.CancelledError): + return await waitfor_task + # Consume inner's exception, to avoid "Future exception was never + # retrieved" messages + if inner_action == 'exception': + self.assertIsInstance(inner.exception(), RuntimeError) + + async def test_simultaneous_self_cancel_and_inner_result(self): + for timeout in (10, None): + with self.subTest(waitfor_timeout=timeout): + await self.simultaneous_self_cancel_and_inner_result( + timeout, 'result') + + async def test_simultaneous_self_cancel_and_inner_exc(self): + for timeout in (10, None): + with self.subTest(waitfor_timeout=timeout): + await self.simultaneous_self_cancel_and_inner_result( + timeout, 'exception') + + async def test_simultaneous_self_cancel_and_inner_cancel(self): + for timeout in (10, None): + with self.subTest(waitfor_timeout=timeout): + await self.simultaneous_self_cancel_and_inner_result( + timeout, 'cancel') + if __name__ == '__main__': unittest.main() diff --git a/Misc/NEWS.d/next/Library/2022-10-26-06-51-59.gh-issue-86296.l1MdXh.rst b/Misc/NEWS.d/next/Library/2022-10-26-06-51-59.gh-issue-86296.l1MdXh.rst new file mode 100644 index 00000000000000..b37e5ec9dca050 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2022-10-26-06-51-59.gh-issue-86296.l1MdXh.rst @@ -0,0 +1 @@ +Make :func:`asyncio.wait_for` not ignore external cancellation if the inner awaitable becomes done at the same time. Patch by twisteroid ambassador.