diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py index f0e690b61a73dd..2072772124a922 100644 --- a/Lib/asyncio/base_events.py +++ b/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( diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py index 39605dca3886c8..923d7f8eb1af0d 100644 --- a/Lib/test/test_asyncgen.py +++ b/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(): diff --git a/Lib/test/test_asyncio/test_base_events.py b/Lib/test/test_asyncio/test_base_events.py index c14a0bb180d79b..bbcd0efcf01d2f 100644 --- a/Lib/test/test_asyncio/test_base_events.py +++ b/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 diff --git a/Misc/NEWS.d/next/Library/2024-04-11-14-11-02.gh-issue-117536.0jp7ep.rst b/Misc/NEWS.d/next/Library/2024-04-11-14-11-02.gh-issue-117536.0jp7ep.rst new file mode 100644 index 00000000000000..f2e28b12a5c17c --- /dev/null +++ b/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.