Finalize async generators in the correct context

Author: oremanj

newsfragments/92.feature.rst

@@ -0,0 +1,4 @@
+trio-asyncio now properly finalizes asyncio-flavored async generators
+upon closure of the event loop. Previously, Trio's async generator finalizers
+would try to finalize all async generators in Trio mode, regardless of their
+flavor, which could lead to stderr spew.

tests/test_trio_asyncio.py

@@ -3,8 +3,10 @@
 import types
 import asyncio
 import trio
+import trio.testing
 import trio_asyncio
 import contextlib
+import gc
 
 
 async def use_asyncio():
@@ -203,3 +205,69 @@ async def main():
         asyncio.run(main())
 
     assert scope.value.code == 42
+
+
+@pytest.mark.trio
+@pytest.mark.parametrize("alive_on_exit", (False, True))
+@pytest.mark.parametrize("slow_finalizer", (False, True))
+@pytest.mark.parametrize("loop_timeout", (0, 1, 20))
+async def test_asyncgens(alive_on_exit, slow_finalizer, loop_timeout, autojump_clock):
+    import sniffio
+
+    record = set()
+    holder = []
+
+    async def agen(label):
+        assert sniffio.current_async_library() == label
+        try:
+            yield 1
+        finally:
+            library = sniffio.current_async_library()
+            try:
+                await sys.modules[library].sleep(10 if slow_finalizer else 0)
+            except (trio.Cancelled, asyncio.CancelledError):
+                pass
+            record.add((label, library))
+
+    async def iterate_in_trio():
+        await agen("trio").asend(None)
+
+    async def aio_main(nursery):
+        nursery.start_soon(iterate_in_trio)
+        ag = agen("asyncio")
+        await ag.asend(None)
+        if alive_on_exit:
+            holder.append(ag)
+        else:
+            del ag
+            for _ in range(5):
+                gc.collect()
+
+    sys.unraisablehook, prev_hook = sys.__unraisablehook__, sys.unraisablehook
+    try:
+        start_time = trio.current_time()
+        with trio.move_on_after(loop_timeout) as scope:
+            if loop_timeout == 0:
+                scope.cancel()
+            async with trio_asyncio.open_loop() as loop:
+                async with trio.open_nursery() as nursery:
+                    await loop.run_aio_coroutine(aio_main(nursery))
+
+        # asyncio agen should be finalized as soon as asyncio loop ends
+        assert ("asyncio", "asyncio") in record
+
+        # asyncio agen finalizer should be able to take a cancel
+        if (slow_finalizer or loop_timeout == 0) and alive_on_exit:
+            assert trio.current_time() == start_time + min(loop_timeout, 10)
+            assert scope.cancelled_caught == (loop_timeout < 10)
+        else:
+            assert trio.current_time() == start_time
+            assert not scope.cancelled_caught
+
+        # trio asyncgen should eventually be finalized in trio mode
+        for _ in range(5):
+            gc.collect()
+        await trio.testing.wait_all_tasks_blocked()
+        assert record == {("trio", "trio"), ("asyncio", "asyncio")}
+    finally:
+        sys.unraisablehook = prev_hook

trio_asyncio/_base.py

@@ -135,6 +135,9 @@ class BaseTrioEventLoop(asyncio.SelectorEventLoop):
     # (threading) Thread this loop is running in
     _thread = None
 
+    # Trio's async generator hooks, as obtained before we added ours
+    _trio_asyncgen_hooks = None
+
     def __init__(self, queue_len=None):
         if queue_len is None:
             queue_len = math.inf
@@ -629,6 +632,30 @@ async def _main_loop_init(self, nursery):
         self._nursery = nursery
         self._task = trio.lowlevel.current_task()
         self._token = trio.lowlevel.current_trio_token()
+        self._trio_asyncgen_hooks = sys.get_asyncgen_hooks()
+        assert self._trio_asyncgen_hooks.firstiter is not None
+        assert self._trio_asyncgen_hooks.finalizer is not None
+        sys.set_asyncgen_hooks(
+            firstiter=self._dispatch_asyncgen_firstiter,
+            finalizer=self._dispatch_asyncgen_finalizer,
+        )
+
+    def _dispatch_asyncgen_firstiter(self, agen):
+        if sniffio_library.name == "asyncio":
+            return self._asyncgen_firstiter_hook(agen)
+        else:
+            agen.ag_frame.f_locals["@trio_asyncio_trio_asyncgen"] = True
+            return self._trio_asyncgen_hooks.firstiter(agen)
+
+    def _dispatch_asyncgen_finalizer(self, agen):
+        try:
+            is_ours = not agen.ag_frame.f_locals.get("@trio_asyncio_trio_asyncgen")
+        except AttributeError:  # pragma: no cover
+            is_ours = True
+        if is_ours:
+            return self._asyncgen_finalizer_hook(agen)
+        else:
+            return self._trio_asyncgen_hooks.finalizer(agen)
 
     async def _main_loop(self, task_status=trio.TASK_STATUS_IGNORED):
         """Run the loop by processing its event queue.
@@ -738,6 +765,10 @@ async def _main_loop_exit(self):
                 except TrioAsyncioExit:
                     pass
 
+        # Restore previous async generator hooks
+        sys.set_asyncgen_hooks(*self._trio_asyncgen_hooks)
+        self._trio_asyncgen_hooks = None
+
         # Kill off unprocessed work
         self._cancel_fds()
         self._cancel_timers()

trio_asyncio/_loop.py

@@ -560,6 +560,45 @@ async def wait_for_sync():
                     tasks_nursery.cancel_scope.cancel()
 
         finally:
+            # If we have any async generators left, finalize them before
+            # closing the event loop. Make sure that the finalizers have a
+            # chance to actually start before they're exposed to any
+            # external cancellation, since asyncio doesn't guarantee that
+            # cancelled tasks have a chance to start first.
+
+            asyncgens_done = trio.Event()
+            if len(loop._asyncgens) == 0:
+                asyncgens_done.set()
+            else:
+                shield_asyncgen_finalizers = trio.CancelScope(shield=True)
+
+                async def sentinel():
+                    try:
+                        yield
+                    finally:
+                        try:
+                            # Open-coded asyncio version of loop.synchronize();
+                            # since we closed the tasks_nursery, we can't do
+                            # any more asyncio-to-trio-mode conversions
+                            w = asyncio.Event()
+                            loop.call_soon(w.set)
+                            await w.wait()
+                        finally:
+                            shield_asyncgen_finalizers.shield = False
+
+                async def shutdown_asyncgens_from_aio():
+                    agen = sentinel()
+                    await agen.asend(None)
+                    try:
+                        await loop.shutdown_asyncgens()
+                    finally:
+                        asyncgens_done.set()
+
+                @loop_nursery.start_soon
+                async def shutdown_asyncgens_from_trio():
+                    with shield_asyncgen_finalizers:
+                        await loop.run_aio_coroutine(shutdown_asyncgens_from_aio())
+
             if forwarded_cancellation is not None:
                 # Now that we're outside the shielded tasks_nursery, we can
                 # add this cancellation to the set of errors propagating out
@@ -570,6 +609,8 @@ async def forward_cancellation():
                     raise forwarded_cancellation
 
             try:
+                with trio.CancelScope(shield=True):
+                    await asyncgens_done.wait()
                 await loop._main_loop_exit()
             finally:
                 loop.close()