Add loop.shutdown_asyncgens() method.

Author: Yury Selivanov

tests/test_base.py

@@ -498,6 +498,47 @@ def coro():
         self.assertFalse(isinstance(task, MyTask))
         self.loop.run_until_complete(task)
 
+    def test_shutdown_asyncgens_01(self):
+        finalized = list()
+
+        if not hasattr(self.loop, 'shutdown_asyncgens'):
+            raise unittest.SkipTest()
+
+        waiter_src = '''async def waiter(timeout, finalized, loop):
+            try:
+                await asyncio.sleep(timeout, loop=loop)
+                yield 1
+            finally:
+                await asyncio.sleep(0, loop=loop)
+                finalized.append(1)
+        '''
+
+        try:
+            g = {}
+            exec(waiter_src, globals(), g)
+        except SyntaxError:
+            # Python < 3.6
+            raise unittest.SkipTest()
+        else:
+            waiter = g['waiter']
+
+        async def wait():
+            async for _ in waiter(1, finalized, self.loop):
+                pass
+
+        t1 = self.loop.create_task(wait())
+        t2 = self.loop.create_task(wait())
+
+        self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
+
+        self.loop.run_until_complete(self.loop.shutdown_asyncgens())
+        self.assertEqual(finalized, [1, 1])
+
+        # Silence warnings
+        t1.cancel()
+        t2.cancel()
+        self.loop.run_until_complete(asyncio.sleep(0.1, loop=self.loop))
+
 
 class TestBaseUV(_TestBase, UVTestCase):
 

uvloop/includes/stdlib.pxi

@@ -125,6 +125,7 @@ cdef tb_format_list = traceback.format_list
 cdef warnings_warn = warnings.warn
 
 cdef weakref_WeakValueDictionary = weakref.WeakValueDictionary
+cdef weakref_WeakSet = weakref.WeakSet
 
 
 # Cython doesn't clean-up imported objects properly in Py3 mode,

uvloop/loop.pxd

@@ -71,6 +71,9 @@ cdef class Loop:
 
         cdef object __weakref__
 
+        object _asyncgens
+        bint _asyncgens_shutdown_called
+
         char _recv_buffer[UV_STREAM_RECV_BUF_SIZE]
         bint _recv_buffer_in_use
 

uvloop/loop.pyx

@@ -131,6 +131,17 @@ cdef class Loop:
 
         self._coroutine_wrapper_set = False
 
+        if hasattr(sys, 'get_asyncgen_hooks'):
+            # Python >= 3.6
+            # A weak set of all asynchronous generators that are
+            # being iterated by the loop.
+            self._asyncgens = weakref_WeakSet()
+        else:
+            self._asyncgens = None
+
+        # Set to True when `loop.shutdown_asyncgens` is called.
+        self._asyncgens_shutdown_called = False
+
     def __init__(self):
         self.set_debug((not sys_ignore_environment
                         and bool(os_environ.get('PYTHONASYNCIODEBUG'))))
@@ -1081,10 +1092,16 @@ cdef class Loop:
             # This is how asyncio loop behaves.
             mode = uv.UV_RUN_NOWAIT
         self._set_coroutine_wrapper(self._debug)
+        if self._asyncgens is not None:
+            old_agen_hooks = sys.get_asyncgen_hooks()
+            sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
+                                   finalizer=self._asyncgen_finalizer_hook)
         try:
             self._run(mode)
         finally:
-            self._set_coroutine_wrapper(0)
+            self._set_coroutine_wrapper(False)
+            if self._asyncgens is not None:
+                sys.set_asyncgen_hooks(*old_agen_hooks)
 
     def close(self):
         """Close the event loop.
@@ -2471,6 +2488,50 @@ cdef class Loop:
         await waiter
         return udp, protocol
 
+    def _asyncgen_finalizer_hook(self, agen):
+        self._asyncgens.discard(agen)
+        if not self.is_closed():
+            self.create_task(agen.aclose())
+            # Wake up the loop if the finalizer was called from
+            # a different thread.
+            self._write_to_self()
+
+    def _asyncgen_firstiter_hook(self, agen):
+        if self._asyncgens_shutdown_called:
+            warnings_warn(
+                "asynchronous generator {!r} was scheduled after "
+                "loop.shutdown_asyncgens() call".format(agen),
+                ResourceWarning, source=self)
+
+        self._asyncgens.add(agen)
+
+    async def shutdown_asyncgens(self):
+        """Shutdown all active asynchronous generators."""
+        self._asyncgens_shutdown_called = True
+
+        if self._asyncgens is None or not len(self._asyncgens):
+            # If Python version is <3.6 or we don't have any asynchronous
+            # generators alive.
+            return
+
+        closing_agens = list(self._asyncgens)
+        self._asyncgens.clear()
+
+        shutdown_coro = aio_gather(
+            *[ag.aclose() for ag in closing_agens],
+            return_exceptions=True,
+            loop=self)
+
+        results = await shutdown_coro
+        for result, agen in zip(results, closing_agens):
+            if isinstance(result, Exception):
+                self.call_exception_handler({
+                    'message': 'an error occurred during closing of '
+                               'asynchronous generator {!r}'.format(agen),
+                    'exception': result,
+                    'asyncgen': agen
+                })
+
 
 cdef void __loop_alloc_buffer(uv.uv_handle_t* uvhandle,
                               size_t suggested_size,