gh-124958: Revert "gh-125472: Revert "gh-124958: fix asyncio.TaskGroup and _PyFuture refcycles ... (#125486)

graingert · web-flow · commit 0b28ea4a35dc · 2024-10-14T15:45:58.000-07:00
* Revert "gh-125472: Revert "gh-124958: fix asyncio.TaskGroup and _PyFuture refcycles (#12… (#125476)" This reverts commit e99650b. * fix incompatability with gh-124392
diff --git 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
@@ -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
@@ -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
@@ -1,7 +1,8 @@
 # Adapted with permission from the EdgeDB project;
 # license: PSFL.
 
-
+import sys
+import gc
 import asyncio
 import contextvars
 import contextlib
@@ -11,7 +12,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)
@@ -29,6 +29,15 @@ def get_error_types(eg):
     return {type(exc) for exc in eg.exceptions}
 
 
+def no_other_refs():
+    # due to gh-124392 coroutines now refer to their locals
+    coro = asyncio.current_task().get_coro()
+    frame = sys._getframe(1)
+    while coro.cr_frame != frame:
+        coro = coro.cr_await
+    return [coro]
+
+
 class TestTaskGroup(unittest.IsolatedAsyncioTestCase):
 
     async def test_taskgroup_01(self):
@@ -899,6 +908,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), no_other_refs())
+
+
+    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), no_other_refs())
+
+
+    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), no_other_refs())
+
+    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), no_other_refs())
+
+    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), no_other_refs())
+
 
 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
@@ -0,0 +1 @@
+Fix refcycles in exceptions raised from :class:`asyncio.TaskGroup` and the python implementation of :class:`asyncio.Future`

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Fix refcycles in exceptions raised from :class:`asyncio.TaskGroup` and the python implementation of :class:`asyncio.Future`