add an asyncio variant of run_in_fork

1 files changed, 53 insertions, 0 deletions

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index 02d1feb..d27fd7b 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -5,6 +5,7 @@
 Python.
 """
 
+import asyncio
 import bisect
 import contextlib
 import dataclasses
@@ -233,6 +234,58 @@ class run_in_fork:
         self.wait()
 
 
+class async_run_in_fork:
+    """Decorator for running the decorated function once in a separate process.
+    Note that the decorator can only be used inside asynchronous code as it
+    uses the running event loop. The decorated function insetad must be
+    synchronous and it must not access the event loop of the main process.
+    """
+    def __init__(self, function: typing.Callable[[], None]):
+        """Fork a new process that will eventually run the given function and
+        then exit.
+        """
+        loop = asyncio.get_running_loop()
+        with asyncio.get_child_watcher() as watcher:
+            if not watcher.is_active():
+                raise RuntimeError(
+                    "active child watcher required for creating a process"
+                )
+            self.future = loop.create_future()
+            self.efd = EventFD()
+            self.pid = os.fork()
+            if self.pid == 0:
+                self.efd.read()
+                self.efd.close()
+                function()
+                os._exit(0)
+            watcher.add_child_handler(self.pid, self._child_callback)
+
+    def _child_callback(self, pid: int, returncode: int) -> None:
+        if self.pid != pid:
+            return
+        self.future.set_result(returncode)
+
+    def start(self) -> None:
+        """Start the decorated function. It can only be started once."""
+        if not self.efd:
+            raise ValueError("this function can only be called once")
+        self.efd.write(1)
+        self.efd.close()
+
+    async def wait(self) -> None:
+        """Wait for the process running the decorated function to finish."""
+        if self.efd:
+            raise ValueError("start must be called before wait")
+        ret = await self.future
+        if ret != 0:
+            raise ValueError("something failed")
+
+    async def __call__(self) -> None:
+        """Start the decorated function and wait for its process to finish."""
+        self.start()
+        await self.wait()
+
+
 def bind_mount(
     source: AtLocationLike,
     target: AtLocationLike,