From ab8786a76c6b0a3df19d6cc2eea505433e24065e Mon Sep 17 00:00:00 2001 From: Chris33871 Date: Tue, 20 May 2025 15:06:32 +0100 Subject: [PATCH 1/4] gh-86802: Fix asyncio memory leak; shielded tasks where cancelled log once through the exception handler --- Lib/asyncio/tasks.py | 30 ++++++++++++++++++++++------- Lib/test/test_asyncio/test_tasks.py | 28 +++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 7 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 888615f8e5e1b3..39a5630d65176f 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -908,6 +908,19 @@ def _done_callback(fut, cur_task=cur_task): return outer +def _log_on_cancel_callback(inner): + if not inner.cancelled(): + exc = inner.exception() + context = { + 'message': + f'{exc.__class__.__name__} exception in shielded future', + 'exception': exc, + 'future': inner, + } + if inner._source_traceback: + context['source_traceback'] = inner._source_traceback + inner._loop.call_exception_handler(context) + def shield(arg): """Wait for a future, shielding it from cancellation. @@ -953,14 +966,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 +982,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) + # Keep only one callback to log on cancel + inner.remove_done_callback(_log_on_cancel_callback) + inner.add_done_callback(_log_on_cancel_callback) + + 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) diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 44498ef790e450..2c234d75813129 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -2116,6 +2116,34 @@ 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_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) From b500a1f86f1ef1b8bb382e7fe4f6666545af7f55 Mon Sep 17 00:00:00 2001 From: Chris33871 Date: Tue, 20 May 2025 15:14:00 +0100 Subject: [PATCH 2/4] Add blurb --- .../next/Library/2025-05-20-15-13-43.gh-issue-86802.trF7TM.rst | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 Misc/NEWS.d/next/Library/2025-05-20-15-13-43.gh-issue-86802.trF7TM.rst diff --git a/Misc/NEWS.d/next/Library/2025-05-20-15-13-43.gh-issue-86802.trF7TM.rst b/Misc/NEWS.d/next/Library/2025-05-20-15-13-43.gh-issue-86802.trF7TM.rst new file mode 100644 index 00000000000000..d3117b16f04436 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2025-05-20-15-13-43.gh-issue-86802.trF7TM.rst @@ -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. From 37837edade2aa9d664a54e3378b06d010e938545 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20Langa?= Date: Tue, 20 May 2025 16:15:03 +0200 Subject: [PATCH 3/4] Fix formatting --- Lib/asyncio/tasks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 39a5630d65176f..dbe5e5f26d154d 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -921,6 +921,7 @@ def _log_on_cancel_callback(inner): context['source_traceback'] = inner._source_traceback inner._loop.call_exception_handler(context) + def shield(arg): """Wait for a future, shielding it from cancellation. From 0d595af36930374e21e534185059a18c3efc637a Mon Sep 17 00:00:00 2001 From: Chris33871 Date: Tue, 20 May 2025 15:30:20 +0100 Subject: [PATCH 4/4] fix handle case with no exception --- Lib/asyncio/tasks.py | 33 +++++++++++++++++------------ Lib/test/test_asyncio/test_tasks.py | 12 +++++++++++ 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index dbe5e5f26d154d..fbd5c39a7c56ac 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -908,18 +908,23 @@ def _done_callback(fut, cur_task=cur_task): return outer -def _log_on_cancel_callback(inner): - if not inner.cancelled(): - exc = inner.exception() - context = { - 'message': - f'{exc.__class__.__name__} exception in shielded future', - 'exception': exc, - 'future': inner, - } - if inner._source_traceback: - context['source_traceback'] = inner._source_traceback - inner._loop.call_exception_handler(context) +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): @@ -987,8 +992,8 @@ 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_cancel_callback) - inner.add_done_callback(_log_on_cancel_callback) + 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) diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 2c234d75813129..f6f976f213ac02 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -2116,6 +2116,18 @@ 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)