Streamline handling of async loop exit

Author: oremanj

trio_asyncio/_base.py

@@ -631,8 +631,16 @@ async def _main_loop(self, task_status=trio.TASK_STATUS_IGNORED):
         sniffio.current_async_library_cvar.set("asyncio")
 
         try:
-            while not self._stopped.is_set():
-                await self._main_loop_one()
+            # The shield here ensures that if the context surrounding
+            # the loop is cancelled, we keep processing callbacks
+            # until we reach the callback inserted by stop().
+            # There's a call to stop() in the finally block of
+            # open_loop(), and we're not shielding the body of the
+            # open_loop() context, so this should be safe against
+            # deadlocks.
+            with trio.CancelScope(shield=True):
+                while not self._stopped.is_set():
+                    await self._main_loop_one()
         except StopAsyncIteration:
             # raised by .stop_me() to interrupt the loop
             pass
@@ -693,16 +701,20 @@ async def _main_loop_exit(self):
         if self._closed:
             return
 
-        self.stop()
-        await self.wait_stopped()
-
-        while True:
-            try:
-                await self._main_loop_one(no_wait=True)
-            except trio.WouldBlock:
-                break
-            except StopAsyncIteration:
-                pass
+        with trio.CancelScope(shield=True):
+            self.stop()
+            await self.wait_stopped()
+
+            # Drain all remaining callbacks, even those after an initial
+            # call to stop(). This avoids a deadlock if stop() was called
+            # again during unwinding.
+            while True:
+                try:
+                    await self._main_loop_one(no_wait=True)
+                except trio.WouldBlock:
+                    break
+                except StopAsyncIteration:
+                    pass
 
         # Kill off unprocessed work
         self._cancel_fds()

trio_asyncio/_loop.py

@@ -390,10 +390,6 @@ async def async_main(*args):
 
     # TODO: make sure that there is no asyncio loop already running
 
-    def _main_loop_exit(self):
-        super()._main_loop_exit()
-        self._thread = None
-
     async with trio.open_nursery() as nursery:
         loop = TrioEventLoop(queue_len=queue_len)
         old_loop = current_loop.set(loop)
@@ -404,14 +400,11 @@ def _main_loop_exit(self):
             await yield_(loop)
         finally:
             try:
-                await loop.stop().wait()
+                await loop._main_loop_exit()
             finally:
-                try:
-                    await loop._main_loop_exit()
-                finally:
-                    loop.close()
-                    nursery.cancel_scope.cancel()
-                    current_loop.reset(old_loop)
+                loop.close()
+                nursery.cancel_scope.cancel()
+                current_loop.reset(old_loop)
 
 
 def run(proc, *args, queue_len=None):