diff --git a/Lib/test/libregrtest/save_env.py b/Lib/test/libregrtest/save_env.py index 0ec12d7c475..eeb1f17d702 100644 --- a/Lib/test/libregrtest/save_env.py +++ b/Lib/test/libregrtest/save_env.py @@ -240,7 +240,9 @@ def get_multiprocessing_process__dangling(self): # 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() diff --git a/crates/vm/src/signal.rs b/crates/vm/src/signal.rs index 57f690bd9df..f146a7321a0 100644 --- a/crates/vm/src/signal.rs +++ b/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; 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 diff --git a/crates/vm/src/stdlib/posix.rs b/crates/vm/src/stdlib/posix.rs index d187d868e5d..6f8f0ef526d 100644 --- a/crates/vm/src/stdlib/posix.rs +++ b/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 = vm.state.after_forkers_child.lock().clone(); run_at_forkers(after_forkers_child, false, vm); } diff --git a/crates/vm/src/stdlib/signal.rs b/crates/vm/src/stdlib/signal.rs index bd93aa194ec..9e51cd3b425 100644 --- a/crates/vm/src/stdlib/signal.rs +++ b/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)] diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index fb190cbd5c9..aa689e85f8c 100644 --- a/crates/vm/src/vm/mod.rs +++ b/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, pub use_tracing: Cell, pub recursion_limit: Cell, - pub(crate) signal_handlers: Option; signal::NSIG]>>>, + pub(crate) signal_handlers: OnceCell; signal::NSIG]>>>, pub(crate) signal_rx: Option, pub repr_guards: RefCell>, pub state: PyRc, @@ -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, state: PyRc) -> 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, diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 334a03701b1..af69fa8d8e5 100644 --- a/crates/vm/src/vm/thread.rs +++ b/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(),