Cancel asyncio tasks when exiting the open_loop() block

Author: oremanj

newsfragments/91.feature.rst

@@ -0,0 +1,8 @@
+Exiting an ``async with trio_asyncio.open_loop():`` block now cancels
+any asyncio tasks that are still running in the background, like
+:func:`asyncio.run` does, so that they have a chance to clean up
+resources by running async context managers and ``finally``
+blocks. Previously such tasks would simply be abandoned to the garbage
+collector, resulting in potential deadlocks and stderr spew. Note that,
+like :func:`asyncio.run`, we *do* still abandon any tasks that are
+started during this finalization phase and outlive the existing tasks.

tests/test_trio_asyncio.py

@@ -1,6 +1,8 @@
 import pytest
 import sys
+import types
 import asyncio
+import trio
 from async_generator import async_generator, yield_
 import trio_asyncio
 
@@ -47,3 +49,54 @@ async def test_get_running_loop():
             pass  # Python 3.6
         else:
             assert get_running_loop() == loop
+
+
+@pytest.mark.trio
+async def test_tasks_get_cancelled():
+    record = []
+    tasks = []
+
+    @types.coroutine
+    def aio_yield():
+        yield
+
+    async def aio_sleeper(key):
+        try:
+            await asyncio.sleep(10)
+            record.append("expired")
+        finally:
+            try:
+                # Prove that we're still running in the aio loop, not
+                # some GC pass
+                await aio_yield()
+            finally:
+                record.append(key)
+                if "early" in key:
+                    tasks.append(asyncio.ensure_future(aio_sleeper("aio late")))
+                    asyncio.get_event_loop().run_trio_task(trio_sleeper, "trio late")
+
+    async def trio_sleeper(key):
+        try:
+            await trio.sleep_forever()
+        finally:
+            await trio.lowlevel.cancel_shielded_checkpoint()
+            record.append(key)
+
+    async with trio_asyncio.open_loop() as loop:
+        tasks.append(asyncio.ensure_future(aio_sleeper("aio early")))
+        loop.run_trio_task(trio_sleeper, "trio early")
+
+    assert record[0] == "aio early"
+    assert set(record[1:]) == {"trio early", "trio late"}
+    assert len(tasks) == 2 and tasks[0].done() and not tasks[1].done()
+
+    # Suppress "Task was destroyed but it was pending!" message
+    tasks[1]._log_traceback = False
+    tasks[1]._log_destroy_pending = False
+
+    # Suppress the "coroutine ignored GeneratorExit" message
+    while True:
+        try:
+            tasks[1]._coro.throw(SystemExit)
+        except SystemExit:
+            break

trio_asyncio/_loop.py

@@ -377,11 +377,33 @@ async def open_loop(queue_len=None):
     """Returns a Trio-flavored async context manager which provides
     an asyncio event loop running on top of Trio.
 
+    The context manager evaluates to a new `TrioEventLoop` object.
+
     Entering the context manager is not enough on its own to immediately
     run asyncio code; it just provides the context that makes running that
     code possible. You additionally need to wrap any asyncio functions
     that you want to run in :func:`aio_as_trio`.
 
+    Exiting the context manager will attempt to do an orderly shutdown
+    of the tasks it contains, analogously to :func:`asyncio.run`.
+    asyncio-flavored tasks are cancelled and awaited first, then
+    Trio-flavored tasks that were started using
+    :meth:`~BaseTrioEventLoop.trio_as_future` or
+    :meth:`~BaseTrioEventLoop.run_trio_task`. All
+    :meth:`~asyncio.loop.call_soon` callbacks that are submitted
+    before exiting the context manager will run before starting
+    this shutdown sequence, and all callbacks that are submitted
+    before the last task exits will run before the loop closes.
+    The exact point at which the loop stops running callbacks is
+    not specified.
+
+    .. warning:: As with :func:`asyncio.run`, asyncio-flavored tasks
+       that are started *after* exiting the context manager (such as by
+       another task as it unwinds) may or may not be cancelled, and will
+       be abandoned if they survive the shutdown sequence. This may lead
+       to unclosed resources, stderr spew about "coroutine ignored
+       GeneratorExit", etc. Trio-flavored tasks do not have this hazard.
+
     Example usage::
 
             async def async_main(*args):
@@ -408,12 +430,11 @@ async def async_main(*args):
                 await loop._main_loop_init(tasks_nursery)
                 await loop_nursery.start(loop._main_loop)
                 yield loop
-                tasks_nursery.cancel_scope.cancel()
 
-                # Allow all submitted run_trio() tasks calls a chance
-                # to start before the tasks_nursery closes, unless the
-                # loop stops (due to someone else calling stop())
-                # before that:
+                # Allow all already-submitted tasks a chance to start
+                # (and then immediately be cancelled), unless the loop
+                # stops (due to someone else calling stop()) before
+                # that.
                 async with trio.open_nursery() as sync_nursery:
                     sync_nursery.cancel_scope.shield = True
 
@@ -425,6 +446,29 @@ async def wait_for_sync():
 
                     await loop.wait_stopped()
                     sync_nursery.cancel_scope.cancel()
+
+                # Cancel and wait on all currently-running asyncio tasks.
+                # Like asyncio.run(), we don't bother cancelling and waiting
+                # on any additional tasks that these tasks start as they
+                # unwind.
+                if sys.version_info >= (3, 7):
+                    aio_tasks = asyncio.all_tasks(loop)
+                else:
+                    aio_tasks = {t for t in asyncio.Task.all_tasks(loop) if not t.done()}
+                if aio_tasks:
+                    # Start one Trio task to wait for each still-running
+                    # asyncio task. This provides better exception
+                    # propagation than using asyncio.gather().
+                    async with trio.open_nursery() as aio_cancel_nursery:
+                        for task in aio_tasks:
+                            aio_cancel_nursery.start_soon(run_aio_future, task)
+                        aio_cancel_nursery.cancel_scope.cancel()
+
+                # If there are any trio_as_aio tasks still going after
+                # the cancellation of asyncio tasks above, this will
+                # cancel them, and exiting the tasks_nursery block
+                # will wait for them to exit.
+                tasks_nursery.cancel_scope.cancel()
         finally:
             try:
                 await loop._main_loop_exit()