gh-158239: Fix asyncio.gather() performance regression (GH-158240)

(cherry picked from commit b06dd22)

Co-authored-by: Timofei Ivankov <128279579+deadlovelll@users.noreply.github.com>

Author: deadlovelll

Lib/asyncio/tasks.py

@@ -784,7 +784,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)
 
@@ -850,11 +850,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
 
@@ -922,10 +925,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,37 @@ 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_gather_does_not_add_callback_to_outer(self):
+        # gh-158239: gather() must not add an internal done callback to
+        # the outer future just to maintain the await graph.
+        async def child():
+            await asyncio.sleep(0)
+
+        async def coro():
+            outer = asyncio.gather(child(), child())
+            self.assertFalse(outer._callbacks)
+            await outer
+
+        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

Misc/ACKS

@@ -881,6 +881,7 @@ Dean Inwood
 Bob Ippolito
 Roger Irwin
 Atsuo Ishimoto
+Timofei Ivankov
 Alexey Izbyshev
 Kasia Jachim
 Adam Jackson