pythongh-117536: Fix asyncio _asyncgen_finalizer_hook()

Make the finalization of asynchronous generators more reliable.

Store a strong reference to the asynchronous generator which is being
closed to make sure that shutdown_asyncgens() can close it even if
asyncio.run() cancels all tasks.

Author: vstinner

Lib/asyncio/base_events.py

@@ -444,6 +444,11 @@ def __init__(self):
         # A weak set of all asynchronous generators that are
         # being iterated by the loop.
         self._asyncgens = weakref.WeakSet()
+
+        # Strong references to asynchronous generators which are being being
+        # closed by the loop: see _asyncgen_finalizer_hook().
+        self._close_asyncgens = set()
+
         # Set to True when `loop.shutdown_asyncgens` is called.
         self._asyncgens_shutdown_called = False
         # Set to True when `loop.shutdown_default_executor` is called.
@@ -555,10 +560,22 @@ def _check_default_executor(self):
         if self._executor_shutdown_called:
             raise RuntimeError('Executor shutdown has been called')
 
+    async def _asyncgen_close(self, agen):
+        await agen.aclose()
+        self._close_asyncgens.discard(agen)
+
     def _asyncgen_finalizer_hook(self, agen):
+        if self.is_closed():
+            self._asyncgens.discard(agen)
+            return
+
+        # gh-117536: Store a strong reference to the asynchronous generator
+        # to make sure that shutdown_asyncgens() can close it even if
+        # asyncio.run() cancels all tasks.
+        self._close_asyncgens.add(agen)
         self._asyncgens.discard(agen)
-        if not self.is_closed():
-            self.call_soon_threadsafe(self.create_task, agen.aclose())
+
+        self.call_soon_threadsafe(self.create_task, self._asyncgen_close(agen))
 
     def _asyncgen_firstiter_hook(self, agen):
         if self._asyncgens_shutdown_called:
@@ -573,12 +590,12 @@ async def shutdown_asyncgens(self):
         """Shutdown all active asynchronous generators."""
         self._asyncgens_shutdown_called = True
 
-        if not len(self._asyncgens):
+        closing_agens = list(set(self._asyncgens) | self._close_asyncgens)
+        if not closing_agens:
             # If Python version is <3.6 or we don't have any asynchronous
             # generators alive.
             return
 
-        closing_agens = list(self._asyncgens)
         self._asyncgens.clear()
 
         results = await tasks.gather(

Lib/test/test_asyncio/test_base_events.py

@@ -1044,6 +1044,27 @@ def test_asyncgen_finalization_by_gc_in_other_thread(self):
             test_utils.run_briefly(self.loop)
             self.assertTrue(status['finalized'])
 
+    def test_shutdown_asyncgens(self):
+        # gh-117536: Test shutdown_asyncgens() using asyncio.run() which
+        # may cancel the task which closes the asynchronous generator before
+        # calling shutdown_asyncgens().
+
+        ns = {'state': 'not started'}
+        async def agen():
+            try:
+                ns['state'] = 'started'
+                yield 0
+                yield 1
+            finally:
+                ns['state'] = 'finalized'
+
+        async def reproducer():
+            async for item in agen():
+                break
+
+        asyncio.run(reproducer())
+        self.assertEqual(ns['state'], 'finalized')
+
 
 class MyProto(asyncio.Protocol):
     done = None

Misc/NEWS.d/next/Library/2024-04-11-14-11-02.gh-issue-117536.0jp7ep.rst

@@ -0,0 +1,2 @@
+:mod:`asyncio`: Make the finalization of asynchronous generators more
+reliable. Patch by Victor Stinner.