run_in_fork: allow starting the function immediately

Doing so skips creating the EventFD.

1 files changed, 46 insertions, 16 deletions

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index 8adf45b..5b163df 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -26,27 +26,39 @@ class run_in_fork:
     """Decorator for running the decorated function once in a separate process.
     """
 
-    def __init__(self, function: typing.Callable[[], None]):
-        """Fork a new process that will eventually run the given function and
-        then exit.
+    def __init__(
+        self, function: typing.Callable[[], None], start: bool = False
+    ):
+        """Fork a new process that will run the given function and then exit.
+        If start is true, run it immediately, otherwise the start or __call__
+        method should be used.
         """
-        self.efd = EventFD()
+        self.efd = None if start else EventFD()
         self.pid = os.fork()
         if self.pid == 0:
             try:
-                self.efd.read()
-                self.efd.close()
+                if self.efd is not None:
+                    self.efd.read()
+                    self.efd.close()
+                    self.efd = None
                 function()
             except:
                 os._exit(1)
             os._exit(0)
 
+    @classmethod
+    def now(cls, function: typing.Callable[[], None]) -> typing.Self:
+        """Fork a new process that will immediately run the given function and
+        then exit."""
+        return cls(function, start=True)
+
     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()
+        self.efd = None
 
     def wait(self) -> None:
         """Wait for the process running the decorated function to finish."""
@@ -57,8 +69,11 @@ class run_in_fork:
             raise ValueError("something failed")
 
     def __call__(self) -> None:
-        """Start the decorated function and wait for its process to finish."""
-        self.start()
+        """Start the decorated function if needed and wait for its process to
+        finish.
+        """
+        if self.efd:
+            self.start()
         self.wait()
 
 
@@ -69,9 +84,12 @@ class async_run_in_fork:
     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.
+    def __init__(
+        self, function: typing.Callable[[], None], start: bool = False
+    ):
+        """Fork a new process that will run the given function and then exit.
+        If start is true, run it immediately, otherwise the start or __call__
+        method should be used.
         """
         loop = asyncio.get_running_loop()
         with asyncio.get_child_watcher() as watcher:
@@ -80,18 +98,26 @@ class async_run_in_fork:
                     "active child watcher required for creating a process"
                 )
             self.future = loop.create_future()
-            self.efd = EventFD()
+            self.efd = None if start else EventFD()
             self.pid = os.fork()
             if self.pid == 0:
                 try:
-                    self.efd.read()
-                    self.efd.close()
+                    if self.efd:
+                        self.efd.read()
+                        self.efd.close()
+                        self.efd = None
                     function()
                 except:
                     os._exit(1)
                 os._exit(0)
             watcher.add_child_handler(self.pid, self._child_callback)
 
+    @classmethod
+    def now(cls, function: typing.Callable[[], None]) -> typing.Self:
+        """Fork a new process that will immediately run the given function and
+        then exit."""
+        return cls(function, start=True)
+
     def _child_callback(self, pid: int, returncode: int) -> None:
         if self.pid != pid:
             return
@@ -103,6 +129,7 @@ class async_run_in_fork:
             raise ValueError("this function can only be called once")
         self.efd.write(1)
         self.efd.close()
+        self.efd = None
 
     async def wait(self) -> None:
         """Wait for the process running the decorated function to finish."""
@@ -113,8 +140,11 @@ class async_run_in_fork:
             raise ValueError("something failed")
 
     async def __call__(self) -> None:
-        """Start the decorated function and wait for its process to finish."""
-        self.start()
+        """Start the decorated function if needed and wait for its process to
+        finish.
+        """
+        if self.efd:
+            self.start()
         await self.wait()