Fix signal handler after fork (#7012)

* temp regrtest save_env patch

* Refactor signal_handlers from Option to OnceCell

- Change `signal_handlers` type from `Option<Box<...>>` to
  `OnceCell<Box<...>>` so fork children can lazily initialize it
- Add `VirtualMachine::is_main_thread()` using runtime thread ID
  comparison instead of structural `signal_handlers.is_none()` check
- Initialize signal handlers in `py_os_after_fork_child()` via
  `get_or_init()` for workers that fork
- Use `get_or_init()` in `signal()` and `getsignal()` functions
- Remove `set_wakeup_fd(-1)` special-case workaround (d6d0303)
- Extract `signal::new_signal_handlers()` to deduplicate init expr
- Allow `getsignal()` from any thread (matches CPython behavior)
- Fix `set_wakeup_fd` error message to name the correct function

Lib/test/libregrtest/save_env.py

@@ -240,7 +240,9 @@ class saved_test_environment:
         # Unjoined process objects can survive after process exits
         multiprocessing_process._cleanup()
         # This copies the weakrefs without making any strong reference
-        return multiprocessing_process._dangling.copy()
+        # TODO: RUSTPYTHON - filter out dead processes since gc doesn't clean WeakSet. Revert this line when we have a GC
+        # return multiprocessing_process._dangling.copy()
+        return {p for p in multiprocessing_process._dangling if p.is_alive()}
     def restore_multiprocessing_process__dangling(self, saved):
         multiprocessing_process = self.get_module('multiprocessing.process')
         multiprocessing_process._dangling.clear()

crates/vm/src/signal.rs

@@ -1,13 +1,17 @@
 #![cfg_attr(target_os = "wasi", allow(dead_code))]
-use crate::{PyResult, VirtualMachine};
+use crate::{PyObjectRef, PyResult, VirtualMachine};
 use alloc::fmt;
-use core::cell::Cell;
+use core::cell::{Cell, RefCell};
 #[cfg(windows)]
 use core::sync::atomic::AtomicIsize;
 use core::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 
 pub(crate) const NSIG: usize = 64;
+
+pub(crate) fn new_signal_handlers() -> Box<RefCell<[Option<PyObjectRef>; NSIG]>> {
+    Box::new(const { RefCell::new([const { None }; NSIG]) })
+}
 static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
 // hack to get around const array repeat expressions, rust issue #79270
 #[allow(
@@ -37,7 +41,7 @@ impl Drop for SignalHandlerGuard {
 #[cfg_attr(feature = "flame-it", flame)]
 #[inline(always)]
 pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
-    if vm.signal_handlers.is_none() {
+    if vm.signal_handlers.get().is_none() {
         return Ok(());
     }
 
@@ -58,7 +62,7 @@ fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
     let _guard = SignalHandlerGuard;
 
     // unwrap should never fail since we check above
-    let signal_handlers = vm.signal_handlers.as_ref().unwrap().borrow();
+    let signal_handlers = vm.signal_handlers.get().unwrap().borrow();
     for (signum, trigger) in TRIGGERS.iter().enumerate().skip(1) {
         let triggered = trigger.swap(false, Ordering::Relaxed);
         if triggered

crates/vm/src/stdlib/posix.rs

@@ -711,6 +711,11 @@ pub mod module {
         #[cfg(feature = "threading")]
         crate::stdlib::thread::after_fork_child(vm);
 
+        // Initialize signal handlers for the child's main thread.
+        // When forked from a worker thread, the OnceCell is empty.
+        vm.signal_handlers
+            .get_or_init(crate::signal::new_signal_handlers);
+
         let after_forkers_child: Vec<PyObjectRef> = vm.state.after_forkers_child.lock().clone();
         run_at_forkers(after_forkers_child, false, vm);
     }

crates/vm/src/stdlib/signal.rs

@@ -177,7 +177,9 @@ pub(crate) mod _signal {
                 } else {
                     None
                 };
-                vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
+                vm.signal_handlers
+                    .get_or_init(signal::new_signal_handlers)
+                    .borrow_mut()[signum] = py_handler;
             }
 
             let int_handler = module
@@ -220,10 +222,9 @@ pub(crate) mod _signal {
                 return Err(vm.new_value_error(format!("signal number {} out of range", signalnum)));
             }
         }
-        let signal_handlers = vm
-            .signal_handlers
-            .as_deref()
-            .ok_or_else(|| vm.new_value_error("signal only works in main thread"))?;
+        if !vm.is_main_thread() {
+            return Err(vm.new_value_error("signal only works in main thread"));
+        }
 
         let sig_handler =
             match usize::try_from_borrowed_object(vm, &handler).ok() {
@@ -245,6 +246,7 @@ pub(crate) mod _signal {
             siginterrupt(signalnum, 1);
         }
 
+        let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
         let old_handler = signal_handlers.borrow_mut()[signalnum as usize].replace(handler);
         Ok(old_handler)
     }
@@ -252,10 +254,7 @@ pub(crate) mod _signal {
     #[pyfunction]
     fn getsignal(signalnum: i32, vm: &VirtualMachine) -> PyResult {
         signal::assert_in_range(signalnum, vm)?;
-        let signal_handlers = vm
-            .signal_handlers
-            .as_deref()
-            .ok_or_else(|| vm.new_value_error("getsignal only works in main thread"))?;
+        let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
         let handler = signal_handlers.borrow()[signalnum as usize]
             .clone()
             .unwrap_or_else(|| vm.ctx.none());
@@ -372,8 +371,8 @@ pub(crate) mod _signal {
         #[cfg(not(windows))]
         let fd = args.fd;
 
-        if vm.signal_handlers.is_none() {
-            return Err(vm.new_value_error("signal only works in main thread"));
+        if !vm.is_main_thread() {
+            return Err(vm.new_value_error("set_wakeup_fd only works in main thread"));
         }
 
         #[cfg(windows)]

crates/vm/src/vm/mod.rs

@@ -41,7 +41,7 @@ use crate::{
 };
 use alloc::{borrow::Cow, collections::BTreeMap};
 use core::{
-    cell::{Cell, Ref, RefCell},
+    cell::{Cell, OnceCell, Ref, RefCell},
     sync::atomic::{AtomicBool, Ordering},
 };
 use crossbeam_utils::atomic::AtomicCell;
@@ -81,7 +81,7 @@ pub struct VirtualMachine {
     pub trace_func: RefCell<PyObjectRef>,
     pub use_tracing: Cell<bool>,
     pub recursion_limit: Cell<usize>,
-    pub(crate) signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
+    pub(crate) signal_handlers: OnceCell<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
     pub(crate) signal_rx: Option<signal::UserSignalReceiver>,
     pub repr_guards: RefCell<HashSet<usize>>,
     pub state: PyRc<PyGlobalState>,
@@ -148,6 +148,20 @@ pub fn process_hash_secret_seed() -> u32 {
 }
 
 impl VirtualMachine {
+    /// Check whether the current thread is the main thread.
+    /// Mirrors `_Py_ThreadCanHandleSignals`.
+    #[allow(dead_code)]
+    pub(crate) fn is_main_thread(&self) -> bool {
+        #[cfg(feature = "threading")]
+        {
+            crate::stdlib::thread::get_ident() == self.state.main_thread_ident.load()
+        }
+        #[cfg(not(feature = "threading"))]
+        {
+            true
+        }
+    }
+
     /// Create a new `VirtualMachine` structure.
     pub(crate) fn new(ctx: PyRc<Context>, state: PyRc<PyGlobalState>) -> Self {
         flame_guard!("new VirtualMachine");
@@ -170,10 +184,7 @@ impl VirtualMachine {
         let importlib = ctx.none();
         let profile_func = RefCell::new(ctx.none());
         let trace_func = RefCell::new(ctx.none());
-        let signal_handlers = Some(Box::new(
-            // putting it in a const optimizes better, prevents linear initialization of the array
-            const { RefCell::new([const { None }; signal::NSIG]) },
-        ));
+        let signal_handlers = OnceCell::from(signal::new_signal_handlers());
 
         let vm = Self {
             builtins,

crates/vm/src/vm/thread.rs

@@ -275,7 +275,7 @@ impl VirtualMachine {
             trace_func: RefCell::new(global_trace.unwrap_or_else(|| self.ctx.none())),
             use_tracing: Cell::new(use_tracing),
             recursion_limit: self.recursion_limit.clone(),
-            signal_handlers: None,
+            signal_handlers: core::cell::OnceCell::new(),
             signal_rx: None,
             repr_guards: RefCell::default(),
             state: self.state.clone(),