gh-117536: Fix asyncio _asyncgen_finalizer_hook()

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_asyncgen.py

@@ -1570,13 +1570,8 @@ async def main():
 
         message, = messages
         self.assertIsInstance(message['exception'], ZeroDivisionError)
-        self.assertIn('unhandled exception during asyncio.run() shutdown',
+        self.assertIn('an error occurred during closing of asynchronous generator ',
                       message['message'])
-        with self.assertWarnsRegex(RuntimeWarning,
-                f"coroutine method 'aclose' of '{async_iterate.__qualname__}' "
-                f"was never awaited"):
-            del message, messages
-            gc_collect()
 
     def test_async_gen_expression_01(self):
         async def arange(n):
@@ -1630,10 +1625,6 @@ async def main():
         asyncio.run(main())
 
         self.assertEqual([], messages)
-        with self.assertWarnsRegex(RuntimeWarning,
-                f"coroutine method 'aclose' of '{async_iterate.__qualname__}' "
-                f"was never awaited"):
-            gc_collect()
 
     def test_async_gen_await_same_anext_coro_twice(self):
         async def async_iterate():

Lib/test/test_asyncio/test_base_events.py

@@ -1044,6 +1044,43 @@ 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_basic():
+            try:
+                ns['state'] = 'started'
+                yield 0
+                yield 1
+            finally:
+                ns['state'] = 'finalized'
+
+        async def reproducer(agen):
+            async for item in agen():
+                break
+
+        asyncio.run(reproducer(agen_basic))
+        self.assertEqual(ns['state'], 'finalized')
+
+        # Similar than previous test, but the generator uses 'await' in the
+        # finally block.
+        ns['state'] = 'not started'
+        async def agen_await():
+            try:
+                ns['state'] = 'started'
+                yield 0
+                yield 1
+            finally:
+                ns['state'] = 'await'
+                await asyncio.sleep(0)
+                ns['state'] = 'finalized'
+
+        asyncio.run(reproducer(agen_await))
+        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.