Support asyncio actors with the trio spawner backend

Author: goodboy

tractor/_child.py

@@ -2,5 +2,8 @@
 import trio
 import cloudpickle
 
+from ._entry import _trio_main
+
 if __name__ == "__main__":
-    trio.run(cloudpickle.load(sys.stdin.buffer))
+    subactor, async_fn, kwargs = cloudpickle.load(sys.stdin.buffer)
+    _trio_main(subactor, async_fn, **kwargs)

tractor/_entry.py

@@ -2,7 +2,7 @@
 Process entry points.
 """
 from functools import partial
-from typing import Tuple, Any
+from typing import Tuple, Any, Callable
 
 import trio  # type: ignore
 
@@ -21,6 +21,7 @@ def _mp_main(
     forkserver_info: Tuple[Any, Any, Any, Any, Any],
     start_method: str,
     parent_addr: Tuple[str, int] = None,
+    infect_asyncio: bool = False,
 ) -> None:
     """The routine called *after fork* which invokes a fresh ``trio.run``
     """
@@ -56,24 +57,30 @@ def _mp_main(
     log.info(f"Actor {actor.uid} terminated")
 
 
-async def _trio_main(
+def _trio_main(
     actor: 'Actor',
-    accept_addr: Tuple[str, int],
-    parent_addr: Tuple[str, int] = None
+    trio_main: Callable,
+    *,
+    infect_asyncio: bool = False,
 ) -> None:
     """Entry point for a `trio_run_in_process` subactor.
 
     Here we don't need to call `trio.run()` since trip does that as
     part of its subprocess startup sequence.
     """
+    log.info(f"Started new trio process for {actor.uid}")
+
     if actor.loglevel is not None:
         log.info(
             f"Setting loglevel for {actor.uid} to {actor.loglevel}")
         get_console_log(actor.loglevel)
 
-    log.info(f"Started new trio process for {actor.uid}")
-
     _state._current_actor = actor
 
-    await actor._async_main(accept_addr, parent_addr=parent_addr)
+    if infect_asyncio:
+        actor._infected_aio = True
+        run_as_asyncio_guest(trio_main)
+    else:
+        trio.run(trio_main)
+
     log.info(f"Actor {actor.uid} terminated")

tractor/_spawn.py

@@ -6,7 +6,7 @@
 import subprocess
 import multiprocessing as mp
 import platform
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Callable
 from functools import partial
 
 import trio
@@ -30,7 +30,7 @@
 from .log import get_logger
 from ._portal import Portal
 from ._actor import Actor, ActorFailure
-from ._entry import _mp_main, _trio_main
+from ._entry import _mp_main
 
 
 log = get_logger('tractor')
@@ -158,8 +158,14 @@ async def cancel_on_completion(
 
 
 @asynccontextmanager
-async def run_in_process(subactor, async_fn, *args, **kwargs):
-    encoded_job = cloudpickle.dumps(partial(async_fn, *args, **kwargs))
+async def run_in_process(
+    subactor: 'Actor',
+    async_fn: Callable,
+    **kwargs
+):
+    encoded_blob = cloudpickle.dumps(
+        (subactor, async_fn, kwargs,)
+    )
 
     async with await trio.open_process(
         [
@@ -176,7 +182,7 @@ async def run_in_process(subactor, async_fn, *args, **kwargs):
     ) as proc:
 
         # send func object to call in child
-        await proc.stdin.send_all(encoded_job)
+        await proc.stdin.send_all(encoded_blob)
         yield proc
 
 
@@ -189,6 +195,7 @@ async def new_proc(
     bind_addr: Tuple[str, int],
     parent_addr: Tuple[str, int],
     use_trio_run_in_process: bool = False,
+    infect_asyncio: bool = False,
     task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED
 ) -> None:
     """Create a new ``multiprocessing.Process`` using the
@@ -203,10 +210,13 @@ async def new_proc(
         if use_trio_run_in_process or _spawn_method == 'trio':
             async with run_in_process(
                 subactor,
-                _trio_main,
-                subactor,
-                bind_addr,
-                parent_addr,
+                partial(
+                    subactor._async_main,
+                    accept_addr=bind_addr,
+                    parent_addr=parent_addr,
+                ),
+                # kwargs bundeled as last pickled tuple element
+                infect_asyncio=infect_asyncio
             ) as proc:
                 log.info(f"Started {proc}")
 
@@ -274,6 +284,7 @@ async def new_proc(
                     fs_info,
                     start_method,
                     parent_addr,
+                    infect_asyncio,
                 ),
                 # daemon=True,
                 name=name,