Nested / parallel loops

Monkey-patch asyncio's get_event_loop{,_policy} calls to return the
currently-running loop / Trio event policy whenever we're within a
Trio task.

Author: smurfix

tests/conftest.py

@@ -48,4 +48,15 @@ def pytest_pyfunc_call(pyfuncitem):
         pyfuncitem.obj = pytest.mark.trio(pyfuncitem.obj)
 
 
-asyncio.set_event_loop_policy(trio_asyncio.TrioPolicy())
+_old_policy = asyncio.get_event_loop_policy()
+_new_policy = trio_asyncio.TrioPolicy()
+asyncio.set_event_loop_policy(_new_policy)
+
+@pytest.fixture
+def old_policy():
+    asyncio.set_event_loop_policy(_old_policy)
+    try:
+        yield _old_policy
+    finally:
+        asyncio.set_event_loop_policy(_new_policy)
+

tests/test_concurrent.py

@@ -0,0 +1,66 @@
+import trio
+import trio_asyncio
+import asyncio
+import pytest
+
+# Tests for concurrent or nested loops
+
+@pytest.mark.trio
+async def test_parallel():
+	loops = [None, None]
+	async with trio.open_nursery() as n:
+		async def gen_loop(i, task_status=trio.TASK_STATUS_IGNORED):
+			task_status.started()
+			async with trio_asyncio.open_loop() as loop:
+				loops[i] = loop
+
+		assert not isinstance(asyncio._get_running_loop(), trio_asyncio.TrioEventLoop)
+		await n.start(gen_loop, 0)
+		await n.start(gen_loop, 1)
+
+	assert isinstance(loops[0], trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[1], trio_asyncio.TrioEventLoop)
+	assert loops[0] is not loops[1]
+
+@pytest.mark.trio
+async def test_nested():
+	loops = [None, None]
+	async with trio.open_nursery() as n:
+		async def gen_loop(i, task_status=trio.TASK_STATUS_IGNORED):
+			task_status.started()
+			async with trio_asyncio.open_loop() as loop:
+				loops[i] = loop
+				if i > 0:
+					await n.start(gen_loop, i-1)
+
+		assert not isinstance(asyncio._get_running_loop(), trio_asyncio.TrioEventLoop)
+		await n.start(gen_loop, 1)
+	assert not isinstance(asyncio._get_running_loop(), trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[0], trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[1], trio_asyncio.TrioEventLoop)
+	assert loops[0] is not loops[1]
+
+async def _test_same_task():
+	loops = [None, None]
+	assert isinstance(asyncio.get_event_loop_policy(), trio_asyncio.TrioPolicy)
+	def get_loop(i):
+		loops[i] = (asyncio.get_event_loop(), asyncio.get_event_loop_policy())
+	async with trio.open_nursery() as n:
+		async with trio_asyncio.open_loop() as loop1:
+			async with trio_asyncio.open_loop() as loop2:
+				loop1.call_later(0.1, get_loop, 0)
+				loop2.call_later(0.1, get_loop, 1)
+				await trio.sleep(0.2)
+
+	assert isinstance(asyncio.get_event_loop_policy(), trio_asyncio.TrioPolicy)
+	assert not isinstance(asyncio._get_running_loop(), trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[0][0], trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[1][0], trio_asyncio.TrioEventLoop)
+	assert isinstance(loops[1][1], trio_asyncio.TrioPolicy)
+	assert loops[0][0] is not loops[1][0]
+	assert loops[0][1] is loops[1][1]
+	assert loops[0][1] is asyncio.get_event_loop_policy()
+
+def test_same_task(old_policy):
+	assert not isinstance(asyncio.get_event_loop_policy(), trio_asyncio.TrioPolicy)
+	trio.run(_test_same_task)

trio_asyncio/async_.py

@@ -105,8 +105,12 @@ def _main_loop_exit(self):
         super()._main_loop_exit()
         self._thread = None
 
+    from .loop import _current_loop
+    outer_loop = _current_loop.loop
+
     async with trio.open_nursery() as nursery:
         loop = TrioEventLoop(queue_len=queue_len)
+        _current_loop.loop = loop
         try:
             loop._closed = False
             await loop._main_loop_init(nursery)
@@ -119,3 +123,5 @@ def _main_loop_exit(self):
                 await loop._main_loop_exit()
                 loop.close()
                 nursery.cancel_scope.cancel()
+                _current_loop.loop = outer_loop
+

trio_asyncio/base.py

@@ -646,7 +646,6 @@ async def _main_loop_init(self, nursery):
         self._nursery = nursery
         self._task = trio.hazmat.current_task()
         self._token = trio.hazmat.current_trio_token()
-        asyncio.events._set_running_loop(self)
 
     async def _main_loop(self, task_status=trio.TASK_STATUS_IGNORED):
         """Run the loop by processing its event queue.
@@ -726,7 +725,6 @@ async def _main_loop_exit(self):
         # clean core fields
         self._nursery = None
         self._task = None
-        asyncio.events._set_running_loop(None)
 
     def is_running(self):
         if self._stopped is None:

trio_asyncio/loop.py

@@ -30,13 +30,14 @@
     'TrioPolicy',
 ]
 
+_current_loop = trio.TaskLocal(loop=None, policy=None)
 
 class _TrioPolicy(asyncio.events.BaseDefaultEventLoopPolicy):
     _loop_factory = TrioEventLoop
 
     def __init__(self):
         super().__init__()
-        self._trio_local = trio.TaskLocal(_loop=None, _task=False)
+        current_loop = trio.TaskLocal(_loop=None, _task=False)
 
     def new_event_loop(self):
         try:
@@ -75,45 +76,68 @@ def get_event_loop(self):
             # this creates a new loop in the main task
             return super().get_event_loop()
         else:
-            return self._trio_local._loop
+            return _current_loop.loop
 
     @property
     def current_event_loop(self):
         """The currently-running event loop, if one exists."""
         try:
-            return self._trio_local._loop
+            return _current_loop.loop
         except RuntimeError:
             # in the main thread this would create a new loop
             # return super().get_event_loop()
-            return self._local._loop
+            return super().get_event_loop()
 
     def set_event_loop(self, loop):
         """Set the current event loop."""
         try:
-            task = trio.hazmat.current_task()
+            _current_loop.loop = loop
         except RuntimeError:
             return super().set_event_loop(loop)
 
-        # This test will not trigger if you create a new asyncio event loop
-        # in a sub-task, which is exactly what we intend to be possible
-        if self._trio_local._loop is not None and loop is not None and \
-                self._trio_local._task == task:
-            raise RuntimeError('You cannot replace an event loop.', self._trio_local._loop, loop)
-        self._trio_local._loop = loop
-        self._trio_local._task = task
+
+
+# We need to monkey-patch asyncio's policy+loop getters to return our
+# TrioPolicy+loop whenever we are within Trio.
+
+from asyncio import events as _aio_event
+
+_orig_policy_get = _aio_event.get_event_loop_policy
+def _new_policy_get():
+    try:
+        policy = _current_loop.policy
+    except RuntimeError:
+        return _orig_policy_get()
+
+    if policy is None:
+        policy = TrioPolicy()
+        _current_loop.policy = policy
+    return policy
+_aio_event.get_event_loop_policy = _new_policy_get
+asyncio.get_event_loop_policy = _new_policy_get
+
+_orig_loop_get = _aio_event._get_running_loop
+def _new_loop_get():
+    try:
+        return _current_loop.loop
+    except RuntimeError:
+        return _orig_loop_get()
+_aio_event._get_running_loop = _new_loop_get
 
 
 class TrioPolicy(_TrioPolicy, asyncio.DefaultEventLoopPolicy):
+    """This is the loop policy that's active whenever we're in a Trio context."""
+
     def _init_watcher(self):
         with asyncio.events._lock:
             if self._watcher is None:  # pragma: no branch
                 self._watcher = TrioChildWatcher()
                 if isinstance(threading.current_thread(), threading._MainThread):
-                    self._watcher.attach_loop(self._trio_local._loop)
+                    self._watcher.attach_loop(_current_loop.loop)
 
         if self._watcher is not None and \
                 isinstance(threading.current_thread(), threading._MainThread):
-            self._watcher.attach_loop(self._trio_local._loop)
+            self._watcher.attach_loop(_current_loop.loop)
 
     def set_child_watcher(self, watcher):
         if watcher is not None: