gh-158239: Fix asyncio.gather performance regression

Lib/asyncio/tasks.py

@@ -786,7 +786,7 @@ def cancel(self, msg=None):
         return ret
 
 
-def _discard_awaited_by(children, waiter, outer):
+def _discard_awaited_by(children, waiter):
     for fut in children:
         futures.future_discard_from_awaited_by(fut, waiter)
 
@@ -852,11 +852,14 @@ def _done_callback(fut, cur_task=cur_task):
                 # 'fut.exception()' will *raise* a CancelledError
                 # instead of returning it.
                 exc = fut._make_cancelled_error()
+                # gh-157213: children outliving gather() must lose the edge
+                _discard_awaited_by(children, cur_task)
                 outer.set_exception(exc)
                 return
             else:
                 exc = fut.exception()
                 if exc is not None:
+                    _discard_awaited_by(children, cur_task)
                     outer.set_exception(exc)
                     return
 
@@ -924,10 +927,6 @@ def _done_callback(fut, cur_task=cur_task):
         children.append(fut)
 
     outer = _GatheringFuture(children, loop=loop)
-    if cur_task is not None:
-        # gh-157213: a child outliving gather() must lose the awaited-by edge
-        outer.add_done_callback(
-            functools.partial(_discard_awaited_by, children, cur_task))
     # Run done callbacks after GatheringFuture created so any post-processing
     # can be performed at this point
     # optimization: in the special case that *all* futures finished eagerly,

Lib/test/test_asyncio/test_tasks.py

@@ -1250,6 +1250,24 @@ async def coro():
 
         self.loop.run_until_complete(self.new_task(self.loop, coro()))
 
+    def test_gather_discards_awaited_by_for_cancelled_sibling(self):
+        # gh-157213: same, when gather() is ended by a cancelled child
+        async def survivor():
+            await asyncio.Future()
+
+        async def coro():
+            t = self.new_task(self.loop, survivor())
+            victim = self.new_task(self.loop, asyncio.sleep(10))
+            victim.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await asyncio.gather(t, victim)
+            self.assertFalse(t._asyncio_awaited_by)
+            t.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await t
+
+        self.loop.run_until_complete(self.new_task(self.loop, coro()))
+
     def test_wait_really_done(self):
         # there is possibility that some tasks in the pending list
         # became done but their callbacks haven't all been called yet