From d5dbbf4372cd3dbf3eead1cc70ddc4261c061fd9 Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Mon, 14 Oct 2024 16:19:56 +0100 Subject: [PATCH] gh-124958: fix asyncio.TaskGroup and _PyFuture refcycles (#124959) --- Lib/asyncio/futures.py | 6 +- Lib/asyncio/taskgroups.py | 41 +++++++-- Lib/test/test_asyncio/test_futures.py | 22 +++++ Lib/test/test_asyncio/test_taskgroups.py | 92 ++++++++++++++++++- ...-10-04-08-46-00.gh-issue-124958.rea9-x.rst | 1 + 5 files changed, 147 insertions(+), 15 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2024-10-04-08-46-00.gh-issue-124958.rea9-x.rst diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py index 5f6fa2348726cf..c95fce035cd548 100644 --- a/Lib/asyncio/futures.py +++ b/Lib/asyncio/futures.py @@ -190,8 +190,7 @@ def result(self): the future is done and has an exception set, this exception is raised. """ if self._state == _CANCELLED: - exc = self._make_cancelled_error() - raise exc + raise self._make_cancelled_error() if self._state != _FINISHED: raise exceptions.InvalidStateError('Result is not ready.') self.__log_traceback = False @@ -208,8 +207,7 @@ def exception(self): InvalidStateError. """ if self._state == _CANCELLED: - exc = self._make_cancelled_error() - raise exc + raise self._make_cancelled_error() if self._state != _FINISHED: raise exceptions.InvalidStateError('Exception is not set.') self.__log_traceback = False diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py index f2ee9648c43876..9fa772ca9d02cc 100644 --- a/Lib/asyncio/taskgroups.py +++ b/Lib/asyncio/taskgroups.py @@ -66,6 +66,20 @@ async def __aenter__(self): return self async def __aexit__(self, et, exc, tb): + tb = None + try: + return await self._aexit(et, exc) + finally: + # Exceptions are heavy objects that can have object + # cycles (bad for GC); let's not keep a reference to + # a bunch of them. It would be nicer to use a try/finally + # in __aexit__ directly but that introduced some diff noise + self._parent_task = None + self._errors = None + self._base_error = None + exc = None + + async def _aexit(self, et, exc): self._exiting = True if (exc is not None and @@ -122,7 +136,10 @@ async def __aexit__(self, et, exc, tb): assert not self._tasks if self._base_error is not None: - raise self._base_error + try: + raise self._base_error + finally: + exc = None if self._parent_cancel_requested: # If this flag is set we *must* call uncancel(). @@ -133,8 +150,14 @@ async def __aexit__(self, et, exc, tb): # Propagate CancelledError if there is one, except if there # are other errors -- those have priority. - if propagate_cancellation_error is not None and not self._errors: - raise propagate_cancellation_error + try: + if propagate_cancellation_error is not None and not self._errors: + try: + raise propagate_cancellation_error + finally: + exc = None + finally: + propagate_cancellation_error = None if et is not None and not issubclass(et, exceptions.CancelledError): self._errors.append(exc) @@ -146,14 +169,14 @@ async def __aexit__(self, et, exc, tb): if self._parent_task.cancelling(): self._parent_task.uncancel() self._parent_task.cancel() - # Exceptions are heavy objects that can have object - # cycles (bad for GC); let's not keep a reference to - # a bunch of them. try: - me = BaseExceptionGroup('unhandled errors in a TaskGroup', self._errors) - raise me from None + raise BaseExceptionGroup( + 'unhandled errors in a TaskGroup', + self._errors, + ) from None finally: - self._errors = None + exc = None + def create_task(self, coro, *, name=None, context=None): """Create a new task in this group and return it. diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py index 458b70451a306a..c566b28adb2408 100644 --- a/Lib/test/test_asyncio/test_futures.py +++ b/Lib/test/test_asyncio/test_futures.py @@ -659,6 +659,28 @@ def __del__(self): fut = self._new_future(loop=self.loop) fut.set_result(Evil()) + def test_future_cancelled_result_refcycles(self): + f = self._new_future(loop=self.loop) + f.cancel() + exc = None + try: + f.result() + except asyncio.CancelledError as e: + exc = e + self.assertIsNotNone(exc) + self.assertListEqual(gc.get_referrers(exc), []) + + def test_future_cancelled_exception_refcycles(self): + f = self._new_future(loop=self.loop) + f.cancel() + exc = None + try: + f.exception() + except asyncio.CancelledError as e: + exc = e + self.assertIsNotNone(exc) + self.assertListEqual(gc.get_referrers(exc), []) + @unittest.skipUnless(hasattr(futures, '_CFuture'), 'requires the C _asyncio module') diff --git a/Lib/test/test_asyncio/test_taskgroups.py b/Lib/test/test_asyncio/test_taskgroups.py index 4852536defc93d..138f59ebf57ef7 100644 --- a/Lib/test/test_asyncio/test_taskgroups.py +++ b/Lib/test/test_asyncio/test_taskgroups.py @@ -1,7 +1,7 @@ # Adapted with permission from the EdgeDB project; # license: PSFL. - +import gc import asyncio import contextvars import contextlib @@ -11,7 +11,6 @@ from test.test_asyncio.utils import await_without_task - # To prevent a warning "test altered the execution environment" def tearDownModule(): asyncio.set_event_loop_policy(None) @@ -899,6 +898,95 @@ async def outer(): await outer() + async def test_exception_refcycles_direct(self): + """Test that TaskGroup doesn't keep a reference to the raised ExceptionGroup""" + tg = asyncio.TaskGroup() + exc = None + + class _Done(Exception): + pass + + try: + async with tg: + raise _Done + except ExceptionGroup as e: + exc = e + + self.assertIsNotNone(exc) + self.assertListEqual(gc.get_referrers(exc), []) + + + async def test_exception_refcycles_errors(self): + """Test that TaskGroup deletes self._errors, and __aexit__ args""" + tg = asyncio.TaskGroup() + exc = None + + class _Done(Exception): + pass + + try: + async with tg: + raise _Done + except* _Done as excs: + exc = excs.exceptions[0] + + self.assertIsInstance(exc, _Done) + self.assertListEqual(gc.get_referrers(exc), []) + + + async def test_exception_refcycles_parent_task(self): + """Test that TaskGroup deletes self._parent_task""" + tg = asyncio.TaskGroup() + exc = None + + class _Done(Exception): + pass + + async def coro_fn(): + async with tg: + raise _Done + + try: + async with asyncio.TaskGroup() as tg2: + tg2.create_task(coro_fn()) + except* _Done as excs: + exc = excs.exceptions[0].exceptions[0] + + self.assertIsInstance(exc, _Done) + self.assertListEqual(gc.get_referrers(exc), []) + + async def test_exception_refcycles_propagate_cancellation_error(self): + """Test that TaskGroup deletes propagate_cancellation_error""" + tg = asyncio.TaskGroup() + exc = None + + try: + async with asyncio.timeout(-1): + async with tg: + await asyncio.sleep(0) + except TimeoutError as e: + exc = e.__cause__ + + self.assertIsInstance(exc, asyncio.CancelledError) + self.assertListEqual(gc.get_referrers(exc), []) + + async def test_exception_refcycles_base_error(self): + """Test that TaskGroup deletes self._base_error""" + class MyKeyboardInterrupt(KeyboardInterrupt): + pass + + tg = asyncio.TaskGroup() + exc = None + + try: + async with tg: + raise MyKeyboardInterrupt + except MyKeyboardInterrupt as e: + exc = e + + self.assertIsNotNone(exc) + self.assertListEqual(gc.get_referrers(exc), []) + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Library/2024-10-04-08-46-00.gh-issue-124958.rea9-x.rst b/Misc/NEWS.d/next/Library/2024-10-04-08-46-00.gh-issue-124958.rea9-x.rst new file mode 100644 index 00000000000000..534d5bb8c898da --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-10-04-08-46-00.gh-issue-124958.rea9-x.rst @@ -0,0 +1 @@ +Fix refcycles in exceptions raised from :class:`asyncio.TaskGroup` and the python implementation of :class:`asyncio.Future`