gh-86802: Fix asyncio memory leak; shielded tasks where cancelled log once through the exception handler

Author: ChristianHrs

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)

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)