RustPython · youknowone · Apr 24, 2023 · Apr 15, 2023 · Apr 16, 2023 · Apr 16, 2023
diff --git a/Lib/concurrent/futures/thread.py b/Lib/concurrent/futures/thread.py
@@ -37,7 +37,8 @@ def _python_exit():
 threading._register_atexit(_python_exit)
 
 # At fork, reinitialize the `_global_shutdown_lock` lock in the child process
-if hasattr(os, 'register_at_fork'):
+# TODO RUSTPYTHON - _at_fork_reinit is not implemented yet
+if hasattr(os, 'register_at_fork') and hasattr(_global_shutdown_lock, '_at_fork_reinit'):
     os.register_at_fork(before=_global_shutdown_lock.acquire,
                         after_in_child=_global_shutdown_lock._at_fork_reinit,
                         after_in_parent=_global_shutdown_lock.release)

diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -4611,6 +4611,34 @@ def test_fork(self):
         assert_python_ok("-c", code)
         assert_python_ok("-c", code, PYTHONMALLOC="malloc_debug")
 
+    @unittest.skipUnless(sys.platform in ("linux", "darwin"),
+                         "Only Linux and macOS detect this today.")
+    def test_fork_warns_when_non_python_thread_exists(self):
+        code = """if 1:
+            import os, threading, warnings
+            from _testcapi import _spawn_pthread_waiter, _end_spawned_pthread
+            _spawn_pthread_waiter()
+            try:
+                with warnings.catch_warnings(record=True) as ws:
+                    warnings.filterwarnings(
+                            "always", category=DeprecationWarning)
+                    if os.fork() == 0:
+                        assert not ws, f"unexpected warnings in child: {ws}"
+                        os._exit(0)  # child
+                    else:
+                        assert ws[0].category == DeprecationWarning, ws[0]
+                        assert 'fork' in str(ws[0].message), ws[0]
+                        # Waiting allows an error in the child to hit stderr.
+                        exitcode = os.wait()[1]
+                        assert exitcode == 0, f"child exited {exitcode}"
+                assert threading.active_count() == 1, threading.enumerate()
+            finally:
+                _end_spawned_pthread()
+        """
+        _, out, err = assert_python_ok("-c", code, PYTHONOPTIMIZE='0')
+        self.assertEqual(err.decode("utf-8"), "")
+        self.assertEqual(out.decode("utf-8"), "")
+
 
 # Only test if the C version is provided, otherwise TestPEP519 already tested
 # the pure Python implementation.

diff --git a/Lib/test/test_socketserver.py b/Lib/test/test_socketserver.py
@@ -55,7 +55,7 @@ class ForkingUnixDatagramServer(socketserver.ForkingMixIn,
                                     socketserver.UnixDatagramServer):
         pass
 
-
+@test.support.requires_fork() # TODO: RUSTPYTHON, os.fork is currently only supported on Unix-based systems
 @contextlib.contextmanager
 def simple_subprocess(testcase):
     """Tests that a custom child process is not waited on (Issue 1540386)"""

diff --git a/Lib/test/test_uuid.py b/Lib/test/test_uuid.py
@@ -641,7 +641,7 @@ def test_uuid5(self):
 
     # TODO: RUSTPYTHON
     @unittest.expectedFailure
-    @unittest.skipUnless(os.name == 'posix', 'requires Posix')
+    @support.requires_fork() # TODO: RUSTPYTHON, os.fork is currently only supported on Unix-based systems
     def testIssue8621(self):
         # On at least some versions of OSX self.uuid.uuid4 generates
         # the same sequence of UUIDs in the parent and any

diff --git a/vm/src/stdlib/os.rs b/vm/src/stdlib/os.rs
@@ -1127,6 +1127,87 @@ pub(super) mod _os {
         OutputMode::String.process_path(curdir_inner(vm)?, vm)
     }
 
+    #[cfg(unix)]
+    #[pyfunction]
+    fn register_at_fork(kwargs: crate::function::KwArgs, vm: &VirtualMachine) -> PyResult<()> {
+        let mut match_found = false; // better way to handle this?
+        for (key, value) in kwargs.into_iter() {
+            if key == "before" {
+                match_found = true;
+                vm.state.before_forkers.lock().push(value.clone());
+            }
+            if key == "after_in_parent" {
+                match_found = true;
+                vm.state.after_forkers_parent.lock().push(value.clone());
+            }
+            if key == "after_in_child" {
+                match_found = true;
+                vm.state.after_forkers_child.lock().push(value.clone());
+            }
+        }
+
+        if !match_found {
+            return Err(vm.new_value_error("At least one argument is required.".to_owned()));
+        }
+
+        Ok(())
+    }
+
+    #[cfg(unix)]
+    fn run_at_forkers(funcs: Vec<PyObjectRef>, vm: &VirtualMachine) {
+        if !funcs.is_empty() {
+            for func in funcs.into_iter().rev() {
+                if let Err(e) = func.call((), vm) {
+                    let exit = e.fast_isinstance(vm.ctx.exceptions.system_exit);
+                    vm.run_unraisable(e, Some("Exception ignored in".to_owned()), func);
+                    if exit {
+                        // Do nothing!
+                    }
+                }
+            }
+        }
+    }
+
+    #[cfg(unix)]
+    fn py_os_before_fork(vm: &VirtualMachine) -> PyResult<()> {
+        let before_forkers: Vec<PyObjectRef> = std::mem::take(&mut *vm.state.before_forkers.lock());
+        run_at_forkers(before_forkers, vm);
+
+        Ok(())
+    }
+
+    #[cfg(unix)]
+    fn py_os_after_fork_child(vm: &VirtualMachine) -> PyResult<()> {
+        let after_forkers_child: Vec<PyObjectRef> =
+            std::mem::take(&mut *vm.state.after_forkers_child.lock());
+        run_at_forkers(after_forkers_child, vm);
+        Ok(())
+    }
+
+    #[cfg(unix)]
+    fn py_os_after_fork_parent(vm: &VirtualMachine) -> PyResult<()> {
+        let after_forkers_parent: Vec<PyObjectRef> =
+            std::mem::take(&mut *vm.state.after_forkers_parent.lock());
+        run_at_forkers(after_forkers_parent, vm);
+        Ok(())
+    }
+
+    #[cfg(unix)]
+    #[pyfunction]
+    fn fork(vm: &VirtualMachine) -> PyResult {
+        let pid: i32;
+        py_os_before_fork(vm)?;
+        unsafe {
+            pid = libc::fork();
+        }
+        if pid == 0 {
+            py_os_after_fork_child(vm)?;
+        } else {
+            py_os_after_fork_parent(vm)?;
+        }
+        Ok(vm.ctx.new_int(pid).into())
+    }
+
     #[pyfunction]
     fn getcwdb(vm: &VirtualMachine) -> PyResult {
         OutputMode::Bytes.process_path(curdir_inner(vm)?, vm)

diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -98,6 +98,9 @@ pub struct PyGlobalState {
     pub finalizing: AtomicBool,
     pub warnings: WarningsState,
     pub override_frozen_modules: AtomicCell<isize>,
+    pub before_forkers: PyMutex<Vec<PyObjectRef>>,
+    pub after_forkers_child: PyMutex<Vec<PyObjectRef>>,
+    pub after_forkers_parent: PyMutex<Vec<PyObjectRef>>,
 }
 
 pub fn process_hash_secret_seed() -> u32 {
@@ -175,6 +178,9 @@ impl VirtualMachine {
                 finalizing: AtomicBool::new(false),
                 warnings,
                 override_frozen_modules: AtomicCell::new(0),
+                before_forkers: PyMutex::default(),
+                after_forkers_child: PyMutex::default(),
+                after_forkers_parent: PyMutex::default(),
             }),
             initialized: false,
             recursion_depth: Cell::new(0),