Skip to content

Commit 7c61ef6

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)
1 parent 1813ea2 commit 7c61ef6

File tree

5 files changed

+39
-14
lines changed

5 files changed

+39
-14
lines changed

‎crates/vm/src/signal.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ impl Drop for SignalHandlerGuard {
3737
#[cfg_attr(feature = "flame-it", flame)]
3838
#[inline(always)]
3939
pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
40-
if vm.signal_handlers.is_none() {
40+
if vm.signal_handlers.get().is_none() {
4141
return Ok(());
4242
}
4343

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

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

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -711,6 +711,12 @@ 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.get_or_init(|| {
717+
Box::new(const { core::cell::RefCell::new([const { None }; crate::signal::NSIG]) })
718+
});
719+
714720
let after_forkers_child: Vec<PyObjectRef> = vm.state.after_forkers_child.lock().clone();
715721
run_at_forkers(after_forkers_child, false, vm);
716722
}

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

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ pub(crate) mod _signal {
1212
builtins::PyTypeRef,
1313
function::{ArgIntoFloat, OptionalArg},
1414
};
15+
use core::cell::RefCell;
1516
use core::sync::atomic::{self, Ordering};
1617

1718
#[cfg(any(unix, windows))]
@@ -177,7 +178,7 @@ pub(crate) mod _signal {
177178
} else {
178179
None
179180
};
180-
vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
181+
vm.signal_handlers.get().unwrap().borrow_mut()[signum] = py_handler;
181182
}
182183

183184
let int_handler = module
@@ -220,10 +221,9 @@ pub(crate) mod _signal {
220221
return Err(vm.new_value_error(format!("signal number {} out of range", signalnum)));
221222
}
222223
}
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"))?;
224+
if !vm.is_main_thread() {
225+
return Err(vm.new_value_error("signal only works in main thread"));
226+
}
227227

228228
let sig_handler =
229229
match usize::try_from_borrowed_object(vm, &handler).ok() {
@@ -245,17 +245,22 @@ pub(crate) mod _signal {
245245
siginterrupt(signalnum, 1);
246246
}
247247

248+
let signal_handlers = vm
249+
.signal_handlers
250+
.get_or_init(|| Box::new(const { RefCell::new([const { None }; signal::NSIG]) }));
248251
let old_handler = signal_handlers.borrow_mut()[signalnum as usize].replace(handler);
249252
Ok(old_handler)
250253
}
251254

252255
#[pyfunction]
253256
fn getsignal(signalnum: i32, vm: &VirtualMachine) -> PyResult {
254257
signal::assert_in_range(signalnum, vm)?;
258+
if !vm.is_main_thread() {
259+
return Err(vm.new_value_error("getsignal only works in main thread"));
260+
}
255261
let signal_handlers = vm
256262
.signal_handlers
257-
.as_deref()
258-
.ok_or_else(|| vm.new_value_error("getsignal only works in main thread"))?;
263+
.get_or_init(|| Box::new(const { RefCell::new([const { None }; signal::NSIG]) }));
259264
let handler = signal_handlers.borrow()[signalnum as usize]
260265
.clone()
261266
.unwrap_or_else(|| vm.ctx.none());
@@ -372,7 +377,7 @@ pub(crate) mod _signal {
372377
#[cfg(not(windows))]
373378
let fd = args.fd;
374379

375-
if vm.signal_handlers.is_none() {
380+
if !vm.is_main_thread() {
376381
return Err(vm.new_value_error("signal only works in main thread"));
377382
}
378383

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

Lines changed: 17 additions & 3 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,7 +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(
187+
let signal_handlers = OnceCell::from(Box::new(
174188
// putting it in a const optimizes better, prevents linear initialization of the array
175189
const { RefCell::new([const { None }; signal::NSIG]) },
176190
));

‎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)