gh-86802: Fix asyncio memory leak; shielded task exceptions log once through the exception handler (gh-134331)

Co-authored-by: Łukasz Langa <lukasz@langa.pl>
This commit is contained in:
Christian Harries 2025-05-20 16:14:27 +01:00 committed by GitHub
parent f3acbb72ff
commit f695eca60c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 72 additions and 7 deletions

View file

@ -908,6 +908,25 @@ def gather(*coros_or_futures, return_exceptions=False):
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.
@ -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():
@ -972,10 +988,16 @@ def shield(arg):
else:
outer.set_result(inner.result())
def _outer_done_callback(outer):
if not inner.done():
inner.remove_done_callback(_inner_done_callback)
# 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)
inner.add_done_callback(_inner_done_callback)
outer.add_done_callback(_outer_done_callback)