Annoying fixes for the sync loop

Problem: $THING calls loop.stop() but then doesn't immediately return.

* split the processing part of the mainloop
  so that we can cobble up a non-waiting single-step mode
* call the single-step mode while there's something to do
  (whenever we wait for run_until_complete)

* return a RuntimeError when the loop has been cancelled
* improve error handling by copying the __cause__ of that
  cancellation (let's hope that it exists) to that error

Author: smurfix

trio_asyncio/async_.py

@@ -68,7 +68,7 @@ def stop(self, waiter=None):
 
         def stop_me():
             waiter.set()
-            raise StopIteration
+            raise StopAsyncIteration
 
         if self._stopped.is_set():
             waiter.set()

trio_asyncio/base.py

@@ -636,59 +636,66 @@ async def _main_loop(self, task_status=trio.TASK_STATUS_IGNORED):
 
         try:
             while True:
-                obj = None
-                if self._timers:
-                    timeout = self._timers[0]._when - self.time()
-                    if timeout <= 0:
-                        # If the timer ran out, process the object now.
-                        obj = heapq.heappop(self._timers)
-                else:
-                    timeout = math.inf
-
-                if obj is None:
-                    with trio.move_on_after(timeout):
-                        obj = await self._q.get()
-                    if obj is None:
-                        # Timeout reached. Presumably now a timer is ready,
-                        # so restart from the beginning.
-                        continue
-
-                    if isinstance(obj, trio.Event):
-                        # Events are used for synchronization.
-                        # Simply set them.
-                        if obj is self._stopped:
-                            break
-                        obj.set()
-                        continue
-
-                    if isinstance(obj, TimerHandle):
-                        # A TimerHandle is added to the list of timers.
-                        heapq.heappush(self._timers, obj)
-                        continue
-
-                assert isinstance(obj, asyncio.Handle)
-                # Hopefully one of ours
-                # but it might be a standard asyncio handle
-
-                if obj._cancelled:
-                    # simply skip cancelled handlers
-                    continue
-
-                # Don't go through the expensive nursery dance
-                # if this is a sync function.
-                if getattr(obj, '_is_sync', True):
-                    obj._callback(*obj._args)
-                else:
-                    await self._nursery.start(obj._call_async)
-
-        except StopIteration:
+                await self._main_loop_one()
+        except StopAsyncIteration:
             # raised by .stop_me() to interrupt the loop
             pass
-
+        except trio.Cancelled:
+            import pdb;pdb.set_trace()
+            raise
         finally:
             # Signal that the loop is no longer running
             self._stopped.set()
 
+    async def _main_loop_one(self, no_wait=False):
+        obj = None
+        if self._timers:
+            timeout = self._timers[0]._when - self.time()
+            if timeout <= 0:
+                # If the timer ran out, process the object now.
+                obj = heapq.heappop(self._timers)
+        else:
+            timeout = math.inf
+
+        if obj is None:
+            if no_wait:
+                obj = self._q.get_nowait()
+            else:
+                with trio.move_on_after(timeout):
+                    obj = await self._q.get()
+                if obj is None:
+                    # Timeout reached. Presumably now a timer is ready,
+                    # so restart from the beginning.
+                    return
+
+            if isinstance(obj, trio.Event):
+                # Events are used for synchronization.
+                # Simply set them.
+                if obj is self._stopped:
+                    raise StopAsyncIteration
+                obj.set()
+                return
+
+            if isinstance(obj, TimerHandle):
+                # A TimerHandle is added to the list of timers.
+                heapq.heappush(self._timers, obj)
+                return
+
+        assert isinstance(obj, asyncio.Handle)
+        # Hopefully one of ours
+        # but it might be a standard asyncio handle
+
+        if obj._cancelled:
+            # simply skip cancelled handlers
+            return
+
+        # Don't go through the expensive nursery dance
+        # if this is a sync function.
+        if getattr(obj, '_is_sync', True):
+            obj._callback(*obj._args)
+        else:
+            await self._nursery.start(obj._call_async)
+
     async def _main_loop_exit(self):
         """Finalize the loop. It may not be re-entered."""
         if self._closed:

trio_asyncio/sync.py

@@ -1,7 +1,9 @@
+import sys
 import trio
 import queue
 import asyncio
 import threading
+import traceback
 
 from .base import BaseTrioEventLoop
 from .handles import Handle
@@ -48,12 +50,12 @@ def stop(self):
         """
 
         def do_stop():
-            raise StopIteration
+            raise StopAsyncIteration
 
 
 #        async def stop_me():
 #            def kick_():
-#                raise StopIteration
+#                raise StopAsyncIteration
 #            self._queue_handle(Handle(kick_, (), self, context=None, is_sync=True))
 #            await self._main_loop()
 #        if threading.current_thread() != self._thread:
@@ -141,10 +143,19 @@ def is_done(_):
         future.add_done_callback(is_done)
         try:
             await self._main_loop()
+            try:
+                while result is None:
+                    try:
+                        await self._main_loop_one(no_wait=True)
+                    except StopAsyncIteration:
+                        pass
+            except trio.WouldBlock:
+                pass
         finally:
             future.remove_done_callback(is_done)
+
         if result is None:
-            result = trio.hazmat.Error(RuntimeError('Event loop stopped before Future completed.'))
+            raise RuntimeError('Event loop stopped before Future completed.')
         return result.unwrap()
 
     def __run_in_thread(self, async_fn, *args):
@@ -155,6 +166,8 @@ def __run_in_thread(self, async_fn, *args):
             raise RuntimeError("The Trio thread is not running")
         self.__blocking_job_queue.put((async_fn, args))
         res = self.__blocking_result_queue.get()
+        if res is None:
+            raise RuntimeError("Loop has died / terminated")
         return res.unwrap()
 
     def _start_loop(self):
@@ -188,8 +201,13 @@ async def __trio_thread_main(self):
                 async_fn, args = req
 
                 result = await trio.hazmat.Result.acapture(async_fn, *args)
+                if type(result) == trio.hazmat.Error and type(result.error) == trio.Cancelled:
+                    res = RuntimeError("Main loop cancelled")
+                    res.__cause__ = result.error.__cause__
+                    result = trio.hazmat.Error(res)
                 self.__blocking_result_queue.put(result)
-            await self._main_loop_exit()
+            with trio.open_cancel_scope(shield=True):
+                await self._main_loop_exit()
             self.__blocking_result_queue.put(None)
             nursery.cancel_scope.cancel()