Set ThreadHandle state to Running in parent thread after spawn

RustPython · youknowone · Mar 8, 2026 · Mar 7, 2026 · Mar 7, 2026 · Mar 7, 2026
commit e37b15fedd1df5d6027ca2ffc2d181715e691072
diff --git a/crates/vm/src/stdlib/posix.rs b/crates/vm/src/stdlib/posix.rs
@@ -872,57 +872,129 @@ pub mod module {
         run_at_forkers(after_forkers_parent, false, vm);
     }
 
-    /// Warn if forking from a multi-threaded process
-    fn warn_if_multi_threaded(name: &str, vm: &VirtualMachine) {
-        // Only check threading if it was already imported
-        // Avoid vm.import() which can execute arbitrary Python code in the fork path
-        let threading = match vm
-            .sys_module
-            .get_attr("modules", vm)
-            .and_then(|m| m.get_item("threading", vm))
+    /// Best-effort number of OS threads in this process.
+    /// Returns <= 0 when unavailable, mirroring CPython fallback behavior.
+    fn get_number_of_os_threads() -> isize {
+        #[cfg(target_os = "macos")]
         {
-            Ok(m) => m,
-            Err(_) => return,
-        };
-        let active = threading.get_attr("_active", vm).ok();
-        let limbo = threading.get_attr("_limbo", vm).ok();
+            type MachPortT = libc::c_uint;
+            type KernReturnT = libc::c_int;
+            type MachMsgTypeNumberT = libc::c_uint;
+            type ThreadActArrayT = *mut MachPortT;
+            const KERN_SUCCESS: KernReturnT = 0;
+            unsafe extern "C" {
+                fn mach_task_self() -> MachPortT;
+                fn task_for_pid(
+                    task: MachPortT,
+                    pid: libc::c_int,
+                    target_task: *mut MachPortT,
+                ) -> KernReturnT;
+                fn task_threads(
+                    target_task: MachPortT,
+                    act_list: *mut ThreadActArrayT,
+                    act_list_cnt: *mut MachMsgTypeNumberT,
+                ) -> KernReturnT;
+                fn vm_deallocate(
+                    target_task: MachPortT,
+                    address: libc::uintptr_t,
+                    size: libc::uintptr_t,
+                ) -> KernReturnT;
+            }
 
-        let count_dict = |obj: Option<crate::PyObjectRef>| -> usize {
-            obj.and_then(|o| o.length_opt(vm))
-                .and_then(|r| r.ok())
-                .unwrap_or(0)
-        };
+            let self_task = unsafe { mach_task_self() };
+            let mut proc_task: MachPortT = 0;
+            if unsafe { task_for_pid(self_task, libc::getpid(), &mut proc_task) } == KERN_SUCCESS {
+                let mut threads: ThreadActArrayT = core::ptr::null_mut();
+                let mut n_threads: MachMsgTypeNumberT = 0;
+                if unsafe { task_threads(proc_task, &mut threads, &mut n_threads) } == KERN_SUCCESS
+                {
+                    if !threads.is_null() {
+                        let _ = unsafe {
+                            vm_deallocate(
+                                self_task,
+                                threads as libc::uintptr_t,
+                                (n_threads as usize * core::mem::size_of::<MachPortT>())
+                                    as libc::uintptr_t,
+                            )
+                        };
+                    }
+                    return n_threads as isize;
+                }
+            }
+            0
+        }
+        #[cfg(target_os = "linux")]
+        {
+            use std::io::Read as _;
+            let mut file = match std::fs::File::open("/proc/self/stat") {
+                Ok(f) => f,
+                Err(_) => return 0,
+            };
+            let mut buf = [0u8; 160];
+            let n = match file.read(&mut buf) {
+                Ok(n) => n,
+                Err(_) => return 0,
+            };
+            let line = match std::str::from_utf8(&buf[..n]) {
+                Ok(s) => s,
+                Err(_) => return 0,
+            };
+            if let Some(field) = line.split_whitespace().nth(19) {
+                return field.parse::<isize>().unwrap_or(0);
+            }
+            return 0;
+        }
+        #[cfg(not(any(target_os = "macos", target_os = "linux")))]
+        {
+            0
+        }
+    }
 
-        let num_threads = count_dict(active) + count_dict(limbo);
-        if num_threads > 1 {
-            // Use Python warnings module to ensure filters are applied correctly
-            let Ok(warnings) = vm.import("warnings", 0) else {
-                return;
+    /// Warn if forking from a multi-threaded process.
+    /// `num_os_threads` should be captured before parent after-fork hooks run.
+    fn warn_if_multi_threaded(name: &str, num_os_threads: isize, vm: &VirtualMachine) {
+        let num_threads = if num_os_threads > 0 {
+            num_os_threads as usize
+        } else {
+            // CPython fallback: if OS-level count isn't available, use the
+            // threading module's active+limbo view.
+            // Only check threading if it was already imported. Avoid vm.import()
+            // which can execute arbitrary Python code in the fork path.
+            let threading = match vm
+                .sys_module
+                .get_attr("modules", vm)
+                .and_then(|m| m.get_item("threading", vm))
+            {
+                Ok(m) => m,
+                Err(_) => return,
             };
-            let Ok(warn_fn) = warnings.get_attr("warn", vm) else {
-                return;
+            let active = threading.get_attr("_active", vm).ok();
+            let limbo = threading.get_attr("_limbo", vm).ok();
+
+            // Match threading module internals and avoid sequence overcounting:
+            // count only dict-backed _active/_limbo containers.
+            let count_dict = |obj: Option<crate::PyObjectRef>| -> usize {
+                obj.and_then(|o| {
+                    o.downcast_ref::<crate::builtins::PyDict>()
+                        .map(|d| d.__len__())
+                })
+                .unwrap_or(0)
             };
 
+            count_dict(active) + count_dict(limbo)
+        };
+
+        if num_threads > 1 {
             let pid = unsafe { libc::getpid() };
             let msg = format!(
                 "This process (pid={}) is multi-threaded, use of {}() may lead to deadlocks in the child.",
                 pid, name
             );
 
-            // Call warnings.warn(message, DeprecationWarning, stacklevel=2)
-            // stacklevel=2 to point to the caller of fork()
-            let args = crate::function::FuncArgs::new(
-                vec![
-                    vm.ctx.new_str(msg).into(),
-                    vm.ctx.exceptions.deprecation_warning.as_object().to_owned(),
-                ],
-                crate::function::KwArgs::new(
-                    [("stacklevel".to_owned(), vm.ctx.new_int(2).into())]
-                        .into_iter()
-                        .collect(),
-                ),
-            );
-            let _ = warn_fn.call(args, vm);
+            // Match PyErr_WarnFormat(..., stacklevel=1) in CPython.
+            // Best effort: ignore failures like CPython does in this path.
+            let _ =
+                crate::stdlib::warnings::warn(vm.ctx.exceptions.deprecation_warning, msg, 1, vm);
         }
     }
 
@@ -953,9 +1025,12 @@ pub mod module {
         if pid == 0 {
             py_os_after_fork_child(vm);
         } else {
+            // Match CPython timing: capture this before parent after-fork hooks
+            // in case those hooks start threads.
+            let num_os_threads = get_number_of_os_threads();
             py_os_after_fork_parent(vm);
             // Match CPython timing: warn only after parent callback path resumes world.
-            warn_if_multi_threaded("fork", vm);
+            warn_if_multi_threaded("fork", num_os_threads, vm);
         }
         if pid == -1 {
             Err(nix::Error::from_raw(saved_errno).into_pyexception(vm))