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

… once through the exception handler
python · ambv · May 21, 2025 · May 20, 2025 · May 20, 2025 · May 20, 2025
commit ab8786a76c6b0a3df19d6cc2eea505433e24065e
diff --git 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
@@ -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)