Skip to content

Commit 32d4579

Browse files
authored
[3.12] gh-124958: fix asyncio.TaskGroup and _PyFuture refcycles (#124959) (#125466)
gh-124958: fix asyncio.TaskGroup and _PyFuture refcycles (#124959)
1 parent 42b8e52 commit 32d4579

File tree

5 files changed

+147
-15
lines changed

5 files changed

+147
-15
lines changed

Lib/asyncio/futures.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -194,8 +194,7 @@ def result(self):
194194
the future is done and has an exception set, this exception is raised.
195195
"""
196196
if self._state == _CANCELLED:
197-
exc = self._make_cancelled_error()
198-
raise exc
197+
raise self._make_cancelled_error()
199198
if self._state != _FINISHED:
200199
raise exceptions.InvalidStateError('Result is not ready.')
201200
self.__log_traceback = False
@@ -212,8 +211,7 @@ def exception(self):
212211
InvalidStateError.
213212
"""
214213
if self._state == _CANCELLED:
215-
exc = self._make_cancelled_error()
216-
raise exc
214+
raise self._make_cancelled_error()
217215
if self._state != _FINISHED:
218216
raise exceptions.InvalidStateError('Exception is not set.')
219217
self.__log_traceback = False

Lib/asyncio/taskgroups.py

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,20 @@ async def __aenter__(self):
6666
return self
6767

6868
async def __aexit__(self, et, exc, tb):
69+
tb = None
70+
try:
71+
return await self._aexit(et, exc)
72+
finally:
73+
# Exceptions are heavy objects that can have object
74+
# cycles (bad for GC); let's not keep a reference to
75+
# a bunch of them. It would be nicer to use a try/finally
76+
# in __aexit__ directly but that introduced some diff noise
77+
self._parent_task = None
78+
self._errors = None
79+
self._base_error = None
80+
exc = None
81+
82+
async def _aexit(self, et, exc):
6983
self._exiting = True
7084

7185
if (exc is not None and
@@ -126,25 +140,34 @@ async def __aexit__(self, et, exc, tb):
126140
assert not self._tasks
127141

128142
if self._base_error is not None:
129-
raise self._base_error
143+
try:
144+
raise self._base_error
145+
finally:
146+
exc = None
130147

131148
# Propagate CancelledError if there is one, except if there
132149
# are other errors -- those have priority.
133-
if propagate_cancellation_error and not self._errors:
134-
raise propagate_cancellation_error
150+
try:
151+
if propagate_cancellation_error and not self._errors:
152+
try:
153+
raise propagate_cancellation_error
154+
finally:
155+
exc = None
156+
finally:
157+
propagate_cancellation_error = None
135158

136159
if et is not None and et is not exceptions.CancelledError:
137160
self._errors.append(exc)
138161

139162
if self._errors:
140-
# Exceptions are heavy objects that can have object
141-
# cycles (bad for GC); let's not keep a reference to
142-
# a bunch of them.
143163
try:
144-
me = BaseExceptionGroup('unhandled errors in a TaskGroup', self._errors)
145-
raise me from None
164+
raise BaseExceptionGroup(
165+
'unhandled errors in a TaskGroup',
166+
self._errors,
167+
) from None
146168
finally:
147-
self._errors = None
169+
exc = None
170+
148171

149172
def create_task(self, coro, *, name=None, context=None):
150173
"""Create a new task in this group and return it.

Lib/test/test_asyncio/test_futures.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -640,6 +640,28 @@ def __del__(self):
640640
fut = self._new_future(loop=self.loop)
641641
fut.set_result(Evil())
642642

643+
def test_future_cancelled_result_refcycles(self):
644+
f = self._new_future(loop=self.loop)
645+
f.cancel()
646+
exc = None
647+
try:
648+
f.result()
649+
except asyncio.CancelledError as e:
650+
exc = e
651+
self.assertIsNotNone(exc)
652+
self.assertListEqual(gc.get_referrers(exc), [])
653+
654+
def test_future_cancelled_exception_refcycles(self):
655+
f = self._new_future(loop=self.loop)
656+
f.cancel()
657+
exc = None
658+
try:
659+
f.exception()
660+
except asyncio.CancelledError as e:
661+
exc = e
662+
self.assertIsNotNone(exc)
663+
self.assertListEqual(gc.get_referrers(exc), [])
664+
643665

644666
@unittest.skipUnless(hasattr(futures, '_CFuture'),
645667
'requires the C _asyncio module')

Lib/test/test_asyncio/test_taskgroups.py

Lines changed: 90 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# Adapted with permission from the EdgeDB project;
22
# license: PSFL.
33

4-
4+
import gc
55
import asyncio
66
import contextvars
77
import contextlib
@@ -10,7 +10,6 @@
1010

1111
from test.test_asyncio.utils import await_without_task
1212

13-
1413
# To prevent a warning "test altered the execution environment"
1514
def tearDownModule():
1615
asyncio.set_event_loop_policy(None)
@@ -824,6 +823,95 @@ async def test_taskgroup_without_parent_task(self):
824823
# We still have to await coro to avoid a warning
825824
await coro
826825

826+
async def test_exception_refcycles_direct(self):
827+
"""Test that TaskGroup doesn't keep a reference to the raised ExceptionGroup"""
828+
tg = asyncio.TaskGroup()
829+
exc = None
830+
831+
class _Done(Exception):
832+
pass
833+
834+
try:
835+
async with tg:
836+
raise _Done
837+
except ExceptionGroup as e:
838+
exc = e
839+
840+
self.assertIsNotNone(exc)
841+
self.assertListEqual(gc.get_referrers(exc), [])
842+
843+
844+
async def test_exception_refcycles_errors(self):
845+
"""Test that TaskGroup deletes self._errors, and __aexit__ args"""
846+
tg = asyncio.TaskGroup()
847+
exc = None
848+
849+
class _Done(Exception):
850+
pass
851+
852+
try:
853+
async with tg:
854+
raise _Done
855+
except* _Done as excs:
856+
exc = excs.exceptions[0]
857+
858+
self.assertIsInstance(exc, _Done)
859+
self.assertListEqual(gc.get_referrers(exc), [])
860+
861+
862+
async def test_exception_refcycles_parent_task(self):
863+
"""Test that TaskGroup deletes self._parent_task"""
864+
tg = asyncio.TaskGroup()
865+
exc = None
866+
867+
class _Done(Exception):
868+
pass
869+
870+
async def coro_fn():
871+
async with tg:
872+
raise _Done
873+
874+
try:
875+
async with asyncio.TaskGroup() as tg2:
876+
tg2.create_task(coro_fn())
877+
except* _Done as excs:
878+
exc = excs.exceptions[0].exceptions[0]
879+
880+
self.assertIsInstance(exc, _Done)
881+
self.assertListEqual(gc.get_referrers(exc), [])
882+
883+
async def test_exception_refcycles_propagate_cancellation_error(self):
884+
"""Test that TaskGroup deletes propagate_cancellation_error"""
885+
tg = asyncio.TaskGroup()
886+
exc = None
887+
888+
try:
889+
async with asyncio.timeout(-1):
890+
async with tg:
891+
await asyncio.sleep(0)
892+
except TimeoutError as e:
893+
exc = e.__cause__
894+
895+
self.assertIsInstance(exc, asyncio.CancelledError)
896+
self.assertListEqual(gc.get_referrers(exc), [])
897+
898+
async def test_exception_refcycles_base_error(self):
899+
"""Test that TaskGroup deletes self._base_error"""
900+
class MyKeyboardInterrupt(KeyboardInterrupt):
901+
pass
902+
903+
tg = asyncio.TaskGroup()
904+
exc = None
905+
906+
try:
907+
async with tg:
908+
raise MyKeyboardInterrupt
909+
except MyKeyboardInterrupt as e:
910+
exc = e
911+
912+
self.assertIsNotNone(exc)
913+
self.assertListEqual(gc.get_referrers(exc), [])
914+
827915

828916
if __name__ == "__main__":
829917
unittest.main()
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix refcycles in exceptions raised from :class:`asyncio.TaskGroup` and the python implementation of :class:`asyncio.Future`

0 commit comments

Comments
 (0)