RustPython · yossi-k · Jul 24, 2021 · Aug 6, 2022
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -2976,8 +2976,6 @@ def check_waitpid(self, code, exitcode, callback=None):
         self.assertEqual(os.waitstatus_to_exitcode(status), exitcode)
         self.assertEqual(pid2, pid)
 
-    # TODO: RUSTPYTHON (AttributeError: module 'os' has no attribute 'spawnv')
-    @unittest.expectedFailure
     def test_waitpid(self):
         self.check_waitpid(code='pass', exitcode=0)
 

diff --git a/Lib/test/test_platform.py b/Lib/test/test_platform.py
@@ -254,8 +254,6 @@ def test_mac_ver(self):
                 self.assertEqual(res[2], 'PowerPC')
 
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     @unittest.skipUnless(sys.platform == 'darwin', "OSX only test")
     def test_mac_ver_with_fork(self):
         # Issue7895: platform.mac_ver() crashes when using fork without exec

diff --git a/Lib/test/test_pty.py b/Lib/test/test_pty.py
@@ -139,8 +139,6 @@ def test_basic(self):
         # to ignore this signal.
         os.close(master_fd)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_fork(self):
         debug("calling pty.fork()")
         pid, master_fd = pty.fork()

diff --git a/Lib/test/test_threading.py b/Lib/test/test_threading.py
@@ -473,6 +473,7 @@ def background_thread(evt):
         self.assertEqual(out, b'')
         self.assertEqual(err, b'')
 
+    @unittest.skip("TODO: RUSTPYTHON, requires sys.getswitchinterval")
     @unittest.skipUnless(hasattr(os, 'fork'), "needs os.fork()")
     def test_is_alive_after_fork(self):
         # Try hard to trigger #18418: is_alive() could sometimes be True on
@@ -802,6 +803,7 @@ def test_1_join_on_shutdown(self):
             """
         self._run_and_join(script)
 
+    @unittest.skip("TODO: RUSTPYTHON")
     @unittest.skipUnless(hasattr(os, 'fork'), "needs os.fork()")
     @unittest.skipIf(sys.platform in platforms_to_skip, "due to known OS bug")
     def test_2_join_in_forked_process(self):
@@ -885,6 +887,7 @@ def main():
         rc, out, err = assert_python_ok('-c', script)
         self.assertFalse(err)
 
+    @unittest.skipIf(sys.platform == "linux", "TODO: RUSTPYTHON, deadlock (?) on linux")
     @unittest.skipUnless(hasattr(os, 'fork'), "needs os.fork()")
     @unittest.skipIf(sys.platform in platforms_to_skip, "due to known OS bug")
     def test_reinit_tls_after_fork(self):
@@ -909,6 +912,7 @@ def do_fork_and_wait():
         for t in threads:
             t.join()
 
+    @unittest.skip("TODO: RUSTPYTHON, requires sys._current_frames")
     @unittest.skipUnless(hasattr(os, 'fork'), "needs os.fork()")
     def test_clear_threads_states_after_fork(self):
         # Issue #17094: check that threads states are cleared after fork()

diff --git a/Lib/test/tf_inherit_check.py b/Lib/test/tf_inherit_check.py
@@ -0,0 +1,27 @@
+# Helper script for test_tempfile.py.  argv[2] is the number of a file
+# descriptor which should _not_ be open.  Check this by attempting to
+# write to it -- if we succeed, something is wrong.
+
+import sys
+import os
+from test.support import SuppressCrashReport
+
+with SuppressCrashReport():
+    verbose = (sys.argv[1] == 'v')
+    try:
+        fd = int(sys.argv[2])
+
+        try:
+            os.write(fd, b"blat")
+        except OSError:
+            # Success -- could not write to fd.
+            sys.exit(0)
+        else:
+            if verbose:
+                sys.stderr.write("fd %d is open in child" % fd)
+            sys.exit(1)
+
+    except Exception:
+        if verbose:
+            raise
+        sys.exit(1)
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -97,6 +97,7 @@ exitcode = "1.1.2"
 uname = "0.1.1"
 strum = "0.24.0"
 strum_macros = "0.24.0"
+fork = "0.1.18"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 rustyline = "9"

diff --git a/vm/src/function/buffer.rs b/vm/src/function/buffer.rs
@@ -1,7 +1,9 @@
 use crate::{
     builtins::{PyStr, PyStrRef},
     common::borrow::{BorrowedValue, BorrowedValueMut},
+    convert::ToPyException,
     protocol::PyBuffer,
+    utils::ToCString,
     PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, TryFromObject, VirtualMachine,
 };
 
@@ -134,6 +136,12 @@ impl TryFromObject for ArgStrOrBytesLike {
     }
 }
 
+impl ToCString for ArgStrOrBytesLike {
+    fn to_cstring(&self, vm: &VirtualMachine) -> PyResult<std::ffi::CString> {
+        std::ffi::CString::new(self.borrow_bytes().to_vec()).map_err(|err| err.to_pyexception(vm))
+    }
+}
+
 impl ArgStrOrBytesLike {
     pub fn borrow_bytes(&self) -> BorrowedValue<'_, [u8]> {
         match self {

diff --git a/vm/src/stdlib/posix.rs b/vm/src/stdlib/posix.rs
@@ -31,7 +31,7 @@ pub mod module {
     use crate::{
         builtins::{PyDictRef, PyInt, PyIntRef, PyListRef, PyStrRef, PyTupleRef, PyTypeRef},
         convert::{IntoPyException, ToPyObject, TryFromObject},
-        function::{Either, OptionalArg},
+        function::{ArgStrOrBytesLike, Either, FromArgs, OptionalArg},
         stdlib::os::{
             errno_err, DirFd, FollowSymlinks, PathOrFd, PyPathLike, SupportFunc, TargetIsDirectory,
             _os, fs_metadata, IOErrorBuilder,
@@ -41,6 +41,7 @@ pub mod module {
         AsObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
     };
     use bitflags::bitflags;
+    use fork::{fork, Fork};
     use nix::{
         fcntl,
         unistd::{self, Gid, Pid, Uid},
@@ -212,6 +213,76 @@ pub mod module {
         }
     }
 
+    #[pyfunction]
+    fn _exit(status: i32) -> PyResult {
+        std::process::exit(status);
+    }
+
+    #[derive(FromArgs)]
+    struct RegisterAtForkArgs {
+        #[pyarg(named, optional)]
+        before: OptionalArg<PyObjectRef>,
+        #[pyarg(named, optional)]
+        after_in_parent: OptionalArg<PyObjectRef>,
+        #[pyarg(named, optional)]
+        after_in_child: OptionalArg<PyObjectRef>,
+    }
+
+    #[pyfunction]
+    fn register_at_fork(args: RegisterAtForkArgs, vm: &VirtualMachine) -> PyResult<()> {
+        for (arg, list) in &mut [
+            (args.before, vm.state.fork_before.lock()),
+            (args.after_in_parent, vm.state.fork_after_parent.lock()),
+            (args.after_in_child, vm.state.fork_after_child.lock()),
+        ] {
+            if let OptionalArg::Present(arg) = arg {
+                if !vm.is_callable(arg) {
+                    return Err(vm.new_type_error("argument must be callable".to_owned()));
+                }
+
+                list.push(arg.clone());
+            }
+        }
+
+        Ok(())
+    }
+
+    fn _before_fork(vm: &VirtualMachine) {
+        // Functions registered for execution before forking are called in reverse order.
+        for func in vm.state.fork_before.lock().drain(..).rev() {
+            vm.invoke(&func, ()).ok();
+        }
+    }
+
+    fn _after_fork_child(vm: &VirtualMachine) {
+        for func in vm.state.fork_after_child.lock().drain(..) {
+            vm.invoke(&func, ()).ok();
+        }
+    }
+
+    fn _after_fork_parent(vm: &VirtualMachine) {
+        for func in vm.state.fork_after_parent.lock().drain(..) {
+            vm.invoke(&func, ()).ok();
+        }
+    }
+
+    #[cfg(unix)]
+    #[pyfunction(name = "fork")]
+    fn fork_impl(vm: &VirtualMachine) -> PyResult<i32> {
+        _before_fork(vm);
+        match fork() {
+            Ok(Fork::Parent(child_pid)) => {
+                _after_fork_parent(vm);
+                Ok(child_pid)
+            }
+            Ok(Fork::Child) => {
+                _after_fork_child(vm);
+                Ok(0)
+            }
+            Err(err) => nix::errno::Errno::result(err).map_err(|e| e.into_pyexception(vm)),
+        }
+    }
+
     #[cfg(any(target_os = "macos", target_os = "ios"))]
     fn getgroups_impl() -> nix::Result<Vec<Gid>> {
         use libc::{c_int, gid_t};
@@ -868,13 +939,13 @@ pub mod module {
     fn execve(
         path: PyPathLike,
         argv: Either<PyListRef, PyTupleRef>,
-        env: PyDictRef,
+        env: crate::function::ArgMapping,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
         let path = path.into_cstring(vm)?;
 
         let argv = vm.extract_elements_with(argv.as_ref(), |obj| {
-            PyStrRef::try_from_object(vm, obj)?.to_cstring(vm)
+            ArgStrOrBytesLike::try_from_object(vm, obj)?.to_cstring(vm)
         })?;
         let argv: Vec<&CStr> = argv.iter().map(|entry| entry.as_c_str()).collect();
 
@@ -888,8 +959,22 @@ pub mod module {
             );
         }
 
-        let env = env
+        let env = env.mapping();
+        let keys = env.keys(vm)?;
+        let values = env.values(vm)?;
+
+        let keys = PyListRef::try_from_object(vm, keys)
+            .map_err(|_| vm.new_type_error("env.keys() is not a list".to_owned()))?
+            .borrow_vec()
+            .to_vec();
+        let values = PyListRef::try_from_object(vm, values)
+            .map_err(|_| vm.new_type_error("env.values() is not a list".to_owned()))?
+            .borrow_vec()
+            .to_vec();
+
+        let env = keys
             .into_iter()
+            .zip(values.into_iter())
             .map(|(k, v)| -> PyResult<_> {
                 let (key, value) = (
                     PyPathLike::try_from_object(vm, k)?.into_bytes(),
@@ -1220,8 +1305,9 @@ pub mod module {
         env: crate::function::ArgMapping,
         vm: &VirtualMachine,
     ) -> PyResult<Vec<CString>> {
-        let keys = env.mapping().keys(vm)?;
-        let values = env.mapping().values(vm)?;
+        let env = env.mapping();
+        let keys = env.keys(vm)?;
+        let values = env.values(vm)?;
 
         let keys = PyListRef::try_from_object(vm, keys)
             .map_err(|_| vm.new_type_error("env.keys() is not a list".to_owned()))?

diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -87,6 +87,9 @@ pub struct PyGlobalState {
     pub hash_secret: HashSecret,
     pub atexit_funcs: PyMutex<Vec<(PyObjectRef, FuncArgs)>>,
     pub codec_registry: CodecsRegistry,
+    pub fork_before: PyMutex<Vec<PyObjectRef>>,
+    pub fork_after_parent: PyMutex<Vec<PyObjectRef>>,
+    pub fork_after_child: PyMutex<Vec<PyObjectRef>>,
     pub finalizing: AtomicBool,
 }
 
@@ -160,6 +163,9 @@ impl VirtualMachine {
                 hash_secret,
                 atexit_funcs: PyMutex::default(),
                 codec_registry,
+                fork_before: PyMutex::default(),
+                fork_after_parent: PyMutex::default(),
+                fork_after_child: PyMutex::default(),
                 finalizing: AtomicBool::new(false),
             }),
             initialized: false,