8000 gh-86802: Fix asyncio memory leak; shielded tasks where cancelled log… · ChristianHrs/cpython@ab8786a · GitHub
[go: up one dir, main page]
Skip to content

Commit ab8786a

Browse files
ChristianHrsChristianHrs
committed
pythongh-86802: Fix asyncio memory leak; shielded tasks where cancelled log once through the exception handler
1 parent fc7f4c3 commit ab8786a

File tree

2 files changed

+51
-7
lines changed

2 files changed

+51
-7
lines changed

‎Lib/asyncio/tasks.py

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -908,6 +908,19 @@ def _done_callback(fut, cur_task=cur_task):
908908
return outer
909909

910910

911+
def _log_on_cancel_callback(inner):
912+
if not inner.cancelled():
913+
exc = inner.exception()
914+
context = {
915+
'message':
916+
f'{exc.__class__.__name__} exception in shielded future',
917+
'exception': exc,
918+
'future': inner,
919+
}
920+
if inner._source_traceback:
921+
context['source_traceback'] = inner._source_traceback
922+
inner._loop.call_exception_handler(context)
923+
911924
def shield(arg):
912925
"""Wait for a future, shielding it from cancellation.
913926
@@ -953,14 +966,11 @@ def shield(arg):
953966
else:
954967
cur_task = None
955968

956-
def _inner_done_callback(inner, cur_task=cur_task):
957-
if cur_task is not None:
958-
futures.future_discard_from_awaited_by(inner, cur_task)
969+
def _clear_awaited_by_callback(inner):
970+
futures.future_discard_from_awaited_by(inner, cur_task)
959971

972+
def _inner_done_callback(inner):
960973
if outer.cancelled():
961-
if not inner.cancelled():
962-
# Mark inner's result as retrieved.
963-
inner.exception()
964974
return
965975

966976
if inner.cancelled():
@@ -972,10 +982,16 @@ def _inner_done_callback(inner, cur_task=cur_task):
972982
else:
973983
outer.set_result(inner.result())
974984

975-
976985
def _outer_done_callback(outer):
977986
if not inner.done():
978987
inner.remove_done_callback(_inner_done_callback)
988+
# Keep only one callback to log on cancel
989+
inner.remove_done_callback(_log_on_cancel_callback)
990+
inner.add_done_callback(_log_on_cancel_callback)
991+
992+
if cur_task is not None:
993+
inner.add_done_callback(_clear_awaited_by_callback)
994+
979995

980996
inner.add_done_callback(_inner_done_callback)
981997
outer.add_done_callback(_outer_done_callback)

‎Lib/test/test_asyncio/test_tasks.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2116,6 +2116,34 @@ def test_shield_cancel_outer(self):
21162116
self.assertTrue(outer.cancelled())
21172117
self.assertEqual(0, 0 if outer._callbacks is None else len(outer._callbacks))
21182118

2119+
def test_shield_cancel_outer_exception(self):
2120+
mock_handler = mock.Mock()
2121+
self.loop.set_exception_handler(mock_handler)
2122+
inner = self.new_future(self.loop)
2123+
outer = asyncio.shield(inner)
2124+
test_utils.run_briefly(self.loop)
2125+
outer.cancel()
2126+
test_utils.run_briefly(self.loop)
2127+
inner.set_exception(Exception('foo'))
2128+
test_utils.run_briefly(self.loop)
2129+
mock_handler.assert_called_once()
2130+
2131+
def test_shield_duplicate_log_once(self):
2132+
mock_handler = mock.Mock()
2133+
self.loop.set_exception_handler(mock_handler)
2134+
inner = self.new_future(self.loop)
2135+
outer = asyncio.shield(inner)
2136+
test_utils.run_briefly(self.loop)
2137+
outer.cancel()
2138+
test_utils.run_briefly(self.loop)
2139+
outer = asyncio.shield(inner)
2140+
test_utils.run_briefly(self.loop)
2141+
outer.cancel()
2142+
test_utils.run_briefly(self.loop)
2143+
inner.set_exception(Exception('foo'))
2144+
test_utils.run_briefly(self.loop)
2145+
mock_handler.assert_called_once()
2146+
21192147
def test_shield_shortcut(self):
21202148
fut = self.new_future(self.loop)
21212149
fut.set_result(42)

0 commit comments

Comments
 (0)
0