Merge pull request #108 from oremanj/cancel-aio-tasks

Cancel and join asyncio tasks when exiting an open_loop() block

Author: pquentin

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/interop/test_calls.py

@@ -1,6 +1,7 @@
 import pytest
 import asyncio
 import trio
+import trio.testing
 import sniffio
 from trio_asyncio import aio_as_trio, trio_as_aio, run_aio_generator
 from tests import aiotest
@@ -331,16 +332,13 @@ async def dly_asyncio(hold, seen):
             seen.flag |= 1
             await hold.wait()
 
-        async def cancel_soon(nursery):
-            await trio.sleep(0.01)
-            nursery.cancel_scope.cancel()
-
         hold = asyncio.Event()
         seen = Seen()
 
         async with trio.open_nursery() as nursery:
             nursery.start_soon(async_gen_to_list, run_aio_generator(loop, dly_asyncio(hold, seen)))
-            nursery.start_soon(cancel_soon, nursery)
+            await trio.testing.wait_all_tasks_blocked()
+            nursery.cancel_scope.cancel()
         assert nursery.cancel_scope.cancel_called
         assert seen.flag == 1
 

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,53 @@ 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 set(record) == {"aio early", "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

@@ -13,6 +13,7 @@
     from async_generator import asynccontextmanager
 
 from ._async import TrioEventLoop
+from ._util import run_aio_future
 from ._deprecate import warn_deprecated
 
 try:
@@ -375,11 +376,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`.
+    Both asyncio-flavored tasks and Trio-flavored tasks (the latter
+    started using :meth:`~BaseTrioEventLoop.trio_as_future`,
+    :meth:`~BaseTrioEventLoop.run_trio_task`, :func:`trio_as_aio`,
+    etc) are cancelled simultaneously, and the loop waits for them to
+    exit in response to this cancellation before proceeding. 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):
@@ -406,12 +429,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
 
@@ -423,6 +445,24 @@ async def wait_for_sync():
 
                     await loop.wait_stopped()
                     sync_nursery.cancel_scope.cancel()
+
+                # Cancel and wait on all currently-running tasks.
+                # Exiting the tasks_nursery will wait for the Trio tasks
+                # automatically; we mix in the asyncio tasks by scheduling
+                # a call to run_aio_future() for each one. It's important
+                # not to wait on one kind of task before the other, so that
+                # we support Trio tasks that need to run some asyncio
+                # code during teardown as well as the opposite.
+                # Like asyncio.run(), we don't bother cancelling and waiting
+                # on any additional asyncio 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()}
+                for task in aio_tasks:
+                    tasks_nursery.start_soon(run_aio_future, task)
+                tasks_nursery.cancel_scope.cancel()
         finally:
             try:
                 await loop._main_loop_exit()