From 6c23635f2b7067ef091a550954e09f8b7c329e3f Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Sat, 21 Oct 2023 22:18:34 +0300 Subject: gh-111085: Fix invalid state handling in TaskGroup and Timeout (#111111) asyncio.TaskGroup and asyncio.Timeout classes now raise proper RuntimeError if they are improperly used. * When they are used without entering the context manager. * When they are used after finishing. * When the context manager is entered more than once (simultaneously or sequentially). * If there is no current task when entering the context manager. They now remain in a consistent state after an exception is thrown, so subsequent operations can be performed correctly (if they are allowed). Co-authored-by: James Hilton-Balfe --- Lib/asyncio/taskgroups.py | 6 +-- Lib/asyncio/timeouts.py | 12 ++++-- Lib/test/test_asyncio/test_taskgroups.py | 45 ++++++++++++++++++++ Lib/test/test_asyncio/test_timeouts.py | 48 +++++++++++++++++++++- Lib/test/test_asyncio/utils.py | 15 +++++++ .../2023-10-20-15-29-10.gh-issue-110910.u2oPwX.rst | 3 ++ 6 files changed, 120 insertions(+), 9 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2023-10-20-15-29-10.gh-issue-110910.u2oPwX.rst diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py index 24238c4..91be0de 100644 --- a/Lib/asyncio/taskgroups.py +++ b/Lib/asyncio/taskgroups.py @@ -54,16 +54,14 @@ class TaskGroup: async def __aenter__(self): if self._entered: raise RuntimeError( - f"TaskGroup {self!r} has been already entered") - self._entered = True - + f"TaskGroup {self!r} has already been entered") if self._loop is None: self._loop = events.get_running_loop() - self._parent_task = tasks.current_task(self._loop) if self._parent_task is None: raise RuntimeError( f'TaskGroup {self!r} cannot determine the parent task') + self._entered = True return self diff --git a/Lib/asyncio/timeouts.py b/Lib/asyncio/timeouts.py index 029c468..30042ab 100644 --- a/Lib/asyncio/timeouts.py +++ b/Lib/asyncio/timeouts.py @@ -49,8 +49,9 @@ class Timeout: def reschedule(self, when: Optional[float]) -> None: """Reschedule the timeout.""" - assert self._state is not _State.CREATED if self._state is not _State.ENTERED: + if self._state is _State.CREATED: + raise RuntimeError("Timeout has not been entered") raise RuntimeError( f"Cannot change state of {self._state.value} Timeout", ) @@ -82,11 +83,14 @@ class Timeout: return f"" async def __aenter__(self) -> "Timeout": + if self._state is not _State.CREATED: + raise RuntimeError("Timeout has already been entered") + task = tasks.current_task() + if task is None: + raise RuntimeError("Timeout should be used inside a task") self._state = _State.ENTERED - self._task = tasks.current_task() + self._task = task self._cancelling = self._task.cancelling() - if self._task is None: - raise RuntimeError("Timeout should be used inside a task") self.reschedule(self._when) return self diff --git a/Lib/test/test_asyncio/test_taskgroups.py b/Lib/test/test_asyncio/test_taskgroups.py index 6a0231f..7a18362 100644 --- a/Lib/test/test_asyncio/test_taskgroups.py +++ b/Lib/test/test_asyncio/test_taskgroups.py @@ -8,6 +8,8 @@ import contextlib from asyncio import taskgroups import unittest +from test.test_asyncio.utils import await_without_task + # To prevent a warning "test altered the execution environment" def tearDownModule(): @@ -779,6 +781,49 @@ class TestTaskGroup(unittest.IsolatedAsyncioTestCase): await asyncio.create_task(main()) + async def test_taskgroup_already_entered(self): + tg = taskgroups.TaskGroup() + async with tg: + with self.assertRaisesRegex(RuntimeError, "has already been entered"): + async with tg: + pass + + async def test_taskgroup_double_enter(self): + tg = taskgroups.TaskGroup() + async with tg: + pass + with self.assertRaisesRegex(RuntimeError, "has already been entered"): + async with tg: + pass + + async def test_taskgroup_finished(self): + tg = taskgroups.TaskGroup() + async with tg: + pass + coro = asyncio.sleep(0) + with self.assertRaisesRegex(RuntimeError, "is finished"): + tg.create_task(coro) + # We still have to await coro to avoid a warning + await coro + + async def test_taskgroup_not_entered(self): + tg = taskgroups.TaskGroup() + coro = asyncio.sleep(0) + with self.assertRaisesRegex(RuntimeError, "has not been entered"): + tg.create_task(coro) + # We still have to await coro to avoid a warning + await coro + + async def test_taskgroup_without_parent_task(self): + tg = taskgroups.TaskGroup() + with self.assertRaisesRegex(RuntimeError, "parent task"): + await await_without_task(tg.__aenter__()) + coro = asyncio.sleep(0) + with self.assertRaisesRegex(RuntimeError, "has not been entered"): + tg.create_task(coro) + # We still have to await coro to avoid a warning + await coro + if __name__ == "__main__": unittest.main() diff --git a/Lib/test/test_asyncio/test_timeouts.py b/Lib/test/test_asyncio/test_timeouts.py index e9b59b9..f54e79e 100644 --- a/Lib/test/test_asyncio/test_timeouts.py +++ b/Lib/test/test_asyncio/test_timeouts.py @@ -5,11 +5,12 @@ import time import asyncio +from test.test_asyncio.utils import await_without_task + def tearDownModule(): asyncio.set_event_loop_policy(None) - class TimeoutTests(unittest.IsolatedAsyncioTestCase): async def test_timeout_basic(self): @@ -257,6 +258,51 @@ class TimeoutTests(unittest.IsolatedAsyncioTestCase): cause = exc.exception.__cause__ assert isinstance(cause, asyncio.CancelledError) + async def test_timeout_already_entered(self): + async with asyncio.timeout(0.01) as cm: + with self.assertRaisesRegex(RuntimeError, "has already been entered"): + async with cm: + pass + + async def test_timeout_double_enter(self): + async with asyncio.timeout(0.01) as cm: + pass + with self.assertRaisesRegex(RuntimeError, "has already been entered"): + async with cm: + pass + + async def test_timeout_finished(self): + async with asyncio.timeout(0.01) as cm: + pass + with self.assertRaisesRegex(RuntimeError, "finished"): + cm.reschedule(0.02) + + async def test_timeout_expired(self): + with self.assertRaises(TimeoutError): + async with asyncio.timeout(0.01) as cm: + await asyncio.sleep(1) + with self.assertRaisesRegex(RuntimeError, "expired"): + cm.reschedule(0.02) + + async def test_timeout_expiring(self): + async with asyncio.timeout(0.01) as cm: + with self.assertRaises(asyncio.CancelledError): + await asyncio.sleep(1) + with self.assertRaisesRegex(RuntimeError, "expiring"): + cm.reschedule(0.02) + + async def test_timeout_not_entered(self): + cm = asyncio.timeout(0.01) + with self.assertRaisesRegex(RuntimeError, "has not been entered"): + cm.reschedule(0.02) + + async def test_timeout_without_task(self): + cm = asyncio.timeout(0.01) + with self.assertRaisesRegex(RuntimeError, "task"): + await await_without_task(cm.__aenter__()) + with self.assertRaisesRegex(RuntimeError, "has not been entered"): + cm.reschedule(0.02) + if __name__ == '__main__': unittest.main() diff --git a/Lib/test/test_asyncio/utils.py b/Lib/test/test_asyncio/utils.py index 83fac4a..18869b3 100644 --- a/Lib/test/test_asyncio/utils.py +++ b/Lib/test/test_asyncio/utils.py @@ -612,3 +612,18 @@ def mock_nonblocking_socket(proto=socket.IPPROTO_TCP, type=socket.SOCK_STREAM, sock.family = family sock.gettimeout.return_value = 0.0 return sock + + +async def await_without_task(coro): + exc = None + def func(): + try: + for _ in coro.__await__(): + pass + except BaseException as err: + nonlocal exc + exc = err + asyncio.get_running_loop().call_soon(func) + await asyncio.sleep(0) + if exc is not None: + raise exc diff --git a/Misc/NEWS.d/next/Library/2023-10-20-15-29-10.gh-issue-110910.u2oPwX.rst b/Misc/NEWS.d/next/Library/2023-10-20-15-29-10.gh-issue-110910.u2oPwX.rst new file mode 100644 index 0000000..c750447 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2023-10-20-15-29-10.gh-issue-110910.u2oPwX.rst @@ -0,0 +1,3 @@ +Fix invalid state handling in :class:`asyncio.TaskGroup` and +:class:`asyncio.Timeout`. They now raise proper RuntimeError if they are +improperly used and are left in consistent state after this. -- cgit v0.12