[stdlib] Adds functionality to spawn and manage processes from exec. file

mojo/stdlib/stdlib/os/__init__.mojo

@@ -44,3 +44,4 @@ from .os import (
     unlink,
 )
 from .pathlike import PathLike
+from .process import Process, Pipe

mojo/stdlib/stdlib/os/process.mojo

@@ -0,0 +1,377 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2025, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+"""Implements os methods for dealing with processes.
+
+Example:
+
+```mojo
+from os import Process
+```
+"""
+from collections import List, Optional
+from collections.string import StringSlice
+
+from memory import LegacyUnsafePointer
+
+from sys import CompilationTarget
+from sys._libc import (
+    waitpid,
+    vfork,
+    execvp,
+    posix_spawnp,
+    exit,
+    kill,
+    SignalCodes,
+    pipe,
+    fcntl,
+    FcntlCommands,
+    FcntlFDFlags,
+    close,
+)
+from sys.ffi import c_char, c_int
+from sys.os import sep
+
+
+# ===----------------------------------------------------------------------=== #
+# Process comm.
+# ===----------------------------------------------------------------------=== #
+struct Pipe:
+    """Create a pipe for interprocess communication.
+
+    Example usage:
+    ```
+    pipe().write_bytes("TEST".as_bytes())
+    ```
+    """
+
+    var fd_in: Optional[FileDescriptor]
+    """File descriptor for pipe input."""
+    var fd_out: Optional[FileDescriptor]
+    """File descriptor for pipe output."""
+
+    fn __init__(
+        out self,
+        in_close_on_exec: Bool = False,
+        out_close_on_exec: Bool = False,
+    ) raises:
+        """Struct to manage interprocess pipe comms.
+
+        Args:
+            in_close_on_exec: Close the read side of pipe if `exec` sys. call is issued in process.
+            out_close_on_exec: Close the write side of pipe if `exec` sys. call is issued in process.
+        """
+        var pipe_fds = alloc[c_int](2)
+        if pipe(pipe_fds) < 0:
+            pipe_fds.free()
+            raise Error("Failed to create pipe")
+
+        if in_close_on_exec:
+            if not self._set_close_on_exec(pipe_fds[0]):
+                pipe_fds.free()
+                raise Error("Failed to configure input pipe close on exec")
+
+        if out_close_on_exec:
+            if not self._set_close_on_exec(pipe_fds[1]):
+                pipe_fds.free()
+                raise Error("Failed to configure output pipe close on exec")
+
+        self.fd_in = FileDescriptor(Int(pipe_fds[0]))
+        self.fd_out = FileDescriptor(Int(pipe_fds[1]))
+        pipe_fds.free()
+
+    fn __del__(deinit self):
+        """Ensures pipes input and output file descriptors are closed, when the object is destroyed.
+        """
+        self.set_input_only()
+        self.set_output_only()
+
+    @staticmethod
+    fn _set_close_on_exec(fd: c_int) -> Bool:
+        return (
+            fcntl(
+                fd,
+                FcntlCommands.F_SETFD,
+                fcntl(fd, FcntlCommands.F_GETFD, 0) | FcntlFDFlags.FD_CLOEXEC,
+            )
+            == 0
+        )
+
+    @always_inline
+    fn set_input_only(mut self):
+        """Close the output descriptor/ channel for this side of the pipe."""
+        if self.fd_out:
+            _ = close(rebind[Int](self.fd_out.value()))
+            self.fd_out = None
+
+    @always_inline
+    fn set_output_only(mut self):
+        """Close the input descriptor/ channel for this side of the pipe."""
+        if self.fd_in:
+            _ = close(rebind[Int](self.fd_in.value()))
+            self.fd_in = None
+
+    @always_inline
+    fn write_bytes(mut self, bytes: Span[Byte, _]) raises:
+        """
+        Write a span of bytes to the pipe.
+
+        Args:
+            bytes: The byte span to write to this pipe.
+
+        """
+        if self.fd_out:
+            self.fd_out.value().write_bytes(bytes)
+        else:
+            raise Error("Can not write from read only side of pipe")
+
+    @always_inline
+    fn read_bytes(mut self, mut buffer: Span[mut=True, Byte]) raises -> UInt:
+        """
+        Read a number of bytes from this pipe.
+
+        Args:
+            buffer: Span[Byte] of length n where to store read bytes. n = number of bytes to read.
+
+        Returns:
+            Actual number of bytes read.
+        """
+        if self.fd_in:
+            return self.fd_in.value().read_bytes(buffer)
+
+        raise Error("Can not read from write only side of pipe")
+
+
+# ===----------------------------------------------------------------------=== #
+# Process execution
+# ===----------------------------------------------------------------------=== #
+
+alias ERR_STR_LEN = 8
+
+
+struct Process:
+    """Create and manage child processes from file executables.
+
+    Example usage:
+    ```
+    child_process = Process.run("ls", List[String]("-lha"))
+    if child_process.interrupt():
+        print("Successfully interrupted.")
+    ```
+    """
+
+    var child_pid: Int
+    """Child process id."""
+
+    fn __init__(out self, child_pid: Int):
+        """Struct to manage metadata about child process.
+        Use the `run` static method to create new process.
+
+        Args:
+          child_pid: The pid of child processed returned by `vfork` that the struct will manage.
+        """
+
+        self.child_pid = child_pid
+
+    fn __del__(deinit self):
+        """ """
+        print("IN DEINIT")
+        if not self.wait_process_status():
+            print("ERROR when waiting on subprocess")
+
+    fn _kill(self, signal: Int) -> Bool:
+        # `kill` returns 0 on success and -1 on failure
+        return kill(self.child_pid, signal) > -1
+
+    fn hangup(self) -> Bool:
+        """Send the Hang up signal to the managed child process.
+
+        Returns:
+          Upon successful completion, True is returned else False.
+        """
+        return self._kill(SignalCodes.HUP)
+
+    fn interrupt(self) -> Bool:
+        """Send the Interrupt signal to the managed child process.
+
+        Returns:
+          Upon successful completion, True is returned else False.
+        """
+        return self._kill(SignalCodes.INT)
+
+    fn kill(self) -> Bool:
+        """Send the Kill signal to the managed child process.
+
+        Returns:
+          Upon successful completion, True is returned else False.
+        """
+        return self._kill(SignalCodes.KILL)
+
+    fn wait_process_status(self) -> Bool:
+        """Wait on stuff.
+
+        Returns:
+          Upon successful completion, True is returned else False.
+        """
+        print("Start wait...")
+        var status: c_int = 0
+        var chk_pid = waitpid(self.child_pid, UnsafePointer(to=status), 0)
+        print("Done wait ...", chk_pid, "==?", self.child_pid)
+        return self.child_pid == chk_pid
+
+    @staticmethod
+    fn run(var path: String, argv: List[String]) raises -> Process:
+        """Spawn new process from file executable.
+
+        Args:
+          path: The path to the file.
+          argv: A list of string arguments to be passed to executable.
+
+        Returns:
+          An instance of `Process` struct.
+        """
+
+        print("Called run")
+
+        @parameter
+        if CompilationTarget.is_linux() or CompilationTarget.is_macos():
+            var file_name = String(path.split(sep)[-1])
+
+            var arg_count = len(argv)
+            var argv_array_ptr_cstr_ptr = LegacyUnsafePointer[
+                LegacyUnsafePointer[c_char, mut=False]
+            ].alloc(arg_count + 2)
+            var offset = 0
+            # Arg 0 in `argv` ptr array should be the file name
+            argv_array_ptr_cstr_ptr[offset] = file_name.unsafe_cstr_ptr()
+            offset += 1
+
+            for var arg in argv:
+                argv_array_ptr_cstr_ptr[offset] = arg.unsafe_cstr_ptr()
+                offset += 1
+
+            # `argv` ptr array terminates with NULL PTR
+            argv_array_ptr_cstr_ptr[offset] = LegacyUnsafePointer[c_char]()
+
+            var path_cptr = path.unsafe_cstr_ptr()
+
+            var pid: Int = 0
+
+            print("Before s")
+
+            var has_error_code = posix_spawnp(
+                UnsafePointer(to=pid),
+                path_cptr,
+                argv_array_ptr_cstr_ptr,
+                LegacyUnsafePointer[LegacyUnsafePointer[Int8, mut=False]](),
+            )
+            print(has_error_code)
+
+            print("After s")
+
+            if has_error_code > 0:
+                raise Error(
+                    "Failed to execute "
+                    + path
+                    + ", EINT error code: "
+                    + String(has_error_code)
+                )
+
+            return Process(child_pid=pid)
+        else:
+            constrained[
+                False, "Unknown platform process execution not implemented"
+            ]()
+            return abort[Process]()
+
+    # @staticmethod
+    # fn run_e(var path: String, argv: List[String]) raises -> Process:
+    #     """Spawn new process from file executable.
+    #
+    #     Args:
+    #       path: The path to the file.
+    #       argv: A list of string arguments to be passed to executable.
+    #
+    #     Returns:
+    #       An instance of `Process` struct.
+    #     """
+    #
+    #     @parameter
+    #     if CompilationTarget.is_linux() or CompilationTarget.is_macos():
+    #         var file_name = String(path.split(sep)[-1])
+    #         var pipe = Pipe(out_close_on_exec=True)
+    #         var exec_err_code = StaticString("EXEC_ERR")
+    #
+    #         var pid = vfork()
+    #
+    #         if pid == 0:
+    #             # Child process.
+    #             pipe.set_output_only()
+    #
+    #             var arg_count = len(argv)
+    #             var argv_array_ptr_cstr_ptr = LegacyUnsafePointer[
+    #                 LegacyUnsafePointer[c_char, mut=False]
+    #             ].alloc(arg_count + 2)
+    #             var offset = 0
+    #             # Arg 0 in `argv` ptr array should be the file name
+    #             argv_array_ptr_cstr_ptr[offset] = file_name.unsafe_cstr_ptr()
+    #             offset += 1
+    #
+    #             for var arg in argv:
+    #                 argv_array_ptr_cstr_ptr[offset] = arg.unsafe_cstr_ptr()
+    #                 offset += 1
+    #
+    #             # `argv` ptr array terminates with NULL PTR
+    #             argv_array_ptr_cstr_ptr[offset] = LegacyUnsafePointer[c_char]()
+    #
+    #             var path_cptr = path.unsafe_cstr_ptr()
+    #
+    #             _ = execvp(path_cptr, argv_array_ptr_cstr_ptr)
+    #
+    #             # This will only get reached if exec call fails to replace currently executing code
+    #             argv_array_ptr_cstr_ptr.free()
+    #
+    #             # Canonical fork/ exec error handling pattern of using a pipe that closes on exec is
+    #             # used to signal error to parent process `https://cr.yp.to/docs/selfpipe.html`
+    #             pipe.write_bytes(exec_err_code.as_bytes())
+    #
+    #             exit(1)
+    #
+    #         elif pid < 0:
+    #             raise Error("Unable to fork parent")
+    #
+    #         var err: Optional[StringSlice[MutAnyOrigin]] = None
+    #         var err_buff_data = InlineArray[Byte, ERR_STR_LEN](fill=0)
+    #
+    #         try:
+    #             pipe.set_input_only()
+    #             var buf = Span[Byte, MutAnyOrigin](
+    #                 ptr=err_buff_data.unsafe_ptr(), length=ERR_STR_LEN
+    #             )
+    #             var bytes_read = pipe.read_bytes(buf)
+    #             err = StringSlice(unsafe_from_utf8=buf)
+    #         except e:
+    #             raise Error(
+    #                 "Failed to read child process response from pipe, exception"
+    #                 " was: "
+    #                 + String(e)
+    #             )
+    #
+    #         if err and len(err.value()) > 0 and err.value() == exec_err_code:
+    #             raise Error("Failed to execute " + path)
+    #
+    #         return Process(child_pid=pid)
+    #     else:
+    #         constrained[
+    #             False, "Unknown platform process execution not implemented"
+    #         ]()
+    #         return abort[Process]()