gh-86802: Fix asyncio memory leak; shielded task exceptions log once through the...

Co-authored-by: Łukasz Langa <lukasz@langa.pl>

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 888615f8e5e1b34dbbb52fce8068f1b5f6258838..fbd5c39a7c56ac40dd4143f7d86119e0480935d7 100644 (file)
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -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)

diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 44498ef790e450dab3bc5db46d753234b25e2d03..f6f976f213ac02e14a7f2b70435bfcd37569edbe 100644 (file)
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -2116,6 +2116,46 @@ class BaseTaskTests:
         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)

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 (file)
index 0000000..d3117b1
--- /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.