Skip to content

Commit c50f6a9

Browse files
youknowoneyouknowone
committed
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
1 parent bcb5a3d commit c50f6a9

File tree

5 files changed

+41
-22
lines changed

5 files changed

+41
-22
lines changed

‎crates/vm/src/signal.rs‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,17 @@
11
#![cfg_attr(target_os = "wasi", allow(dead_code))]
2-
use crate::{PyResult, VirtualMachine};
2+
use crate::{PyObjectRef, PyResult, VirtualMachine};
33
use alloc::fmt;
4-
use core::cell::Cell;
4+
use core::cell::{Cell, RefCell};
55
#[cfg(windows)]
66
use core::sync::atomic::AtomicIsize;
77
use core::sync::atomic::{AtomicBool, Ordering};
88
use std::sync::mpsc;
99

1010
pub(crate) const NSIG: usize = 64;
11+
12+
pub(crate) fn new_signal_handlers() -> Box<RefCell<[Option<PyObjectRef>; NSIG]>> {
13+
Box::new(const { RefCell::new([const { None }; NSIG]) })
14+
}
1115
static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
1216
// hack to get around const array repeat expressions, rust issue #79270
1317
#[allow(
@@ -37,7 +41,7 @@ impl Drop for SignalHandlerGuard {
3741
#[cfg_attr(feature = "flame-it", flame)]
3842
#[inline(always)]
3943
pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
40-
if vm.signal_handlers.is_none() {
44+
if vm.signal_handlers.get().is_none() {
4145
return Ok(());
4246
}
4347

@@ -58,7 +62,7 @@ fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
5862
let _guard = SignalHandlerGuard;
5963

6064
// unwrap should never fail since we check above
61-
let signal_handlers = vm.signal_handlers.as_ref().unwrap().borrow();
65+
let signal_handlers = vm.signal_handlers.get().unwrap().borrow();
6266
for (signum, trigger) in TRIGGERS.iter().enumerate().skip(1) {
6367
let triggered = trigger.swap(false, Ordering::Relaxed);
6468
if triggered

‎crates/vm/src/stdlib/posix.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -711,6 +711,11 @@ pub mod module {
711711
#[cfg(feature = "threading")]
712712
crate::stdlib::thread::after_fork_child(vm);
713713

714+
// Initialize signal handlers for the child's main thread.
715+
// When forked from a worker thread, the OnceCell is empty.
716+
vm.signal_handlers
717+
.get_or_init(crate::signal::new_signal_handlers);
718+
714719
let after_forkers_child: Vec<PyObjectRef> = vm.state.after_forkers_child.lock().clone();
715720
run_at_forkers(after_forkers_child, false, vm);
716721
}

‎crates/vm/src/stdlib/signal.rs‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,9 @@ pub(crate) mod _signal {
177177
} else {
178178
None
179179
};
180-
vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
180+
vm.signal_handlers
181+
.get_or_init(signal::new_signal_handlers)
182+
.borrow_mut()[signum] = py_handler;
181183
}
182184

183185
let int_handler = module
@@ -220,10 +222,9 @@ pub(crate) mod _signal {
220222
return Err(vm.new_value_error(format!("signal number {} out of range", signalnum)));
221223
}
222224
}
223-
let signal_handlers = vm
224-
.signal_handlers
225-
.as_deref()
226-
.ok_or_else(|| vm.new_value_error("signal only works in main thread"))?;
225+
if !vm.is_main_thread() {
226+
return Err(vm.new_value_error("signal only works in main thread"));
227+
}
227228

228229
let sig_handler =
229230
match usize::try_from_borrowed_object(vm, &handler).ok() {
@@ -245,17 +246,15 @@ pub(crate) mod _signal {
245246
siginterrupt(signalnum, 1);
246247
}
247248

249+
let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
248250
let old_handler = signal_handlers.borrow_mut()[signalnum as usize].replace(handler);
249251
Ok(old_handler)
250252
}
251253

252254
#[pyfunction]
253255
fn getsignal(signalnum: i32, vm: &VirtualMachine) -> PyResult {
254256
signal::assert_in_range(signalnum, vm)?;
255-
let signal_handlers = vm
256-
.signal_handlers
257-
.as_deref()
258-
.ok_or_else(|| vm.new_value_error("getsignal only works in main thread"))?;
257+
let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
259258
let handler = signal_handlers.borrow()[signalnum as usize]
260259
.clone()
261260
.unwrap_or_else(|| vm.ctx.none());
@@ -372,8 +371,8 @@ pub(crate) mod _signal {
372371
#[cfg(not(windows))]
373372
let fd = args.fd;
374373

375-
if vm.signal_handlers.is_none() {
376-
return Err(vm.new_value_error("signal only works in main thread"));
374+
if !vm.is_main_thread() {
375+
return Err(vm.new_value_error("set_wakeup_fd only works in main thread"));
377376
}
378377

379378
#[cfg(windows)]

‎crates/vm/src/vm/mod.rs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ use crate::{
4141
};
4242
use alloc::{borrow::Cow, collections::BTreeMap};
4343
use core::{
44-
cell::{Cell, Ref, RefCell},
44+
cell::{Cell, OnceCell, Ref, RefCell},
4545
sync::atomic::{AtomicBool, Ordering},
4646
};
4747
use crossbeam_utils::atomic::AtomicCell;
@@ -81,7 +81,7 @@ pub struct VirtualMachine {
8181
pub trace_func: RefCell<PyObjectRef>,
8282
pub use_tracing: Cell<bool>,
8383
pub recursion_limit: Cell<usize>,
84-
pub(crate) signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
84+
pub(crate) signal_handlers: OnceCell<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
8585
pub(crate) signal_rx: Option<signal::UserSignalReceiver>,
8686
pub repr_guards: RefCell<HashSet<usize>>,
8787
pub state: PyRc<PyGlobalState>,
@@ -148,6 +148,20 @@ pub fn process_hash_secret_seed() -> u32 {
148148
}
149149

150150
impl VirtualMachine {
151+
/// Check whether the current thread is the main thread.
152+
/// Mirrors `_Py_ThreadCanHandleSignals`.
153+
#[allow(dead_code)]
154+
pub(crate) fn is_main_thread(&self) -> bool {
155+
#[cfg(feature = "threading")]
156+
{
157+
crate::stdlib::thread::get_ident() == self.state.main_thread_ident.load()
158+
}
159+
#[cfg(not(feature = "threading"))]
160+
{
161+
true
162+
}
163+
}
164+
151165
/// Create a new `VirtualMachine` structure.
152166
pub(crate) fn new(ctx: PyRc<Context>, state: PyRc<PyGlobalState>) -> Self {
153167
flame_guard!("new VirtualMachine");
@@ -170,10 +184,7 @@ impl VirtualMachine {
170184
let importlib = ctx.none();
171185
let profile_func = RefCell::new(ctx.none());
172186
let trace_func = RefCell::new(ctx.none());
173-
let signal_handlers = Some(Box::new(
174-
// putting it in a const optimizes better, prevents linear initialization of the array
175-
const { RefCell::new([const { None }; signal::NSIG]) },
176-
));
187+
let signal_handlers = OnceCell::from(signal::new_signal_handlers());
177188

178189
let vm = Self {
179190
builtins,

‎crates/vm/src/vm/thread.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,7 +275,7 @@ impl VirtualMachine {
275275
trace_func: RefCell::new(global_trace.unwrap_or_else(|| self.ctx.none())),
276276
use_tracing: Cell::new(use_tracing),
277277
recursion_limit: self.recursion_limit.clone(),
278-
signal_handlers: None,
278+
signal_handlers: core::cell::OnceCell::new(),
279279
signal_rx: None,
280280
repr_guards: RefCell::default(),
281281
state: self.state.clone(),

0 commit comments

Comments
 (0)