Skip to content

gh-86802: Fix asyncio memory leak; shielded tasks where cancelled log… #134331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 20, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 29 additions & 7 deletions Lib/asyncio/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -908,6 +908,25 @@ def _done_callback(fut, cur_task=cur_task):
return outer


def _log_on_exception(fut):
if fut.cancelled():
return

exc = fut.exception()
if exc is None:
return

context = {
'message':
f'{exc.__class__.__name__} exception in shielded future',
'exception': exc,
'future': fut,
}
if fut._source_traceback:
context['source_traceback'] = fut._source_traceback
fut._loop.call_exception_handler(context)


def shield(arg):
"""Wait for a future, shielding it from cancellation.

Expand Down Expand Up @@ -953,14 +972,11 @@ def shield(arg):
else:
cur_task = None

def _inner_done_callback(inner, cur_task=cur_task):
if cur_task is not None:
futures.future_discard_from_awaited_by(inner, cur_task)
def _clear_awaited_by_callback(inner):
futures.future_discard_from_awaited_by(inner, cur_task)

def _inner_done_callback(inner):
if outer.cancelled():
if not inner.cancelled():
# Mark inner's result as retrieved.
inner.exception()
return

if inner.cancelled():
Expand All @@ -972,10 +988,16 @@ def _inner_done_callback(inner, cur_task=cur_task):
else:
outer.set_result(inner.result())


def _outer_done_callback(outer):
if not inner.done():
inner.remove_done_callback(_inner_done_callback)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before this PR, this line made the future_discard_from_awaited_by never execute for cases when the outer shield was cancelled. That was a memory leak, which is now fixed.

# Keep only one callback to log on cancel
inner.remove_done_callback(_log_on_exception)
inner.add_done_callback(_log_on_exception)

if cur_task is not None:
inner.add_done_callback(_clear_awaited_by_callback)
Comment on lines +998 to +999
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now we're only adding the callback to future_discard_from_awaited_by when there is a current task. Nice.



inner.add_done_callback(_inner_done_callback)
outer.add_done_callback(_outer_done_callback)
Expand Down
40 changes: 40 additions & 0 deletions Lib/test/test_asyncio/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -2116,6 +2116,46 @@ def test_shield_cancel_outer(self):
self.assertTrue(outer.cancelled())
self.assertEqual(0, 0 if outer._callbacks is None else len(outer._callbacks))

def test_shield_cancel_outer_result(self):
mock_handler = mock.Mock()
self.loop.set_exception_handler(mock_handler)
inner = self.new_future(self.loop)
outer = asyncio.shield(inner)
test_utils.run_briefly(self.loop)
outer.cancel()
test_utils.run_briefly(self.loop)
inner.set_result(1)
test_utils.run_briefly(self.loop)
mock_handler.assert_not_called()

def test_shield_cancel_outer_exception(self):
mock_handler = mock.Mock()
self.loop.set_exception_handler(mock_handler)
inner = self.new_future(self.loop)
outer = asyncio.shield(inner)
test_utils.run_briefly(self.loop)
outer.cancel()
test_utils.run_briefly(self.loop)
inner.set_exception(Exception('foo'))
test_utils.run_briefly(self.loop)
mock_handler.assert_called_once()

def test_shield_duplicate_log_once(self):
mock_handler = mock.Mock()
self.loop.set_exception_handler(mock_handler)
inner = self.new_future(self.loop)
outer = asyncio.shield(inner)
test_utils.run_briefly(self.loop)
outer.cancel()
test_utils.run_briefly(self.loop)
outer = asyncio.shield(inner)
test_utils.run_briefly(self.loop)
outer.cancel()
test_utils.run_briefly(self.loop)
inner.set_exception(Exception('foo'))
test_utils.run_briefly(self.loop)
mock_handler.assert_called_once()

def test_shield_shortcut(self):
fut = self.new_future(self.loop)
fut.set_result(42)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fixed asyncio memory leak in cancelled shield tasks. For shielded tasks
where the shield was cancelled, log potential exceptions through the
exception handler. Contributed by Christian Harries.
Loading