Skip to content

Commit dc6c6ef

Browse files
youknowoneyouknowone
committed
fix signal
1 parent e8d2a33 commit dc6c6ef

2 files changed

Lines changed: 42 additions & 9 deletions

File tree

‎crates/vm/src/signal.rs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
use crate::{PyResult, VirtualMachine};
33
use alloc::fmt;
44
use core::sync::atomic::{AtomicBool, Ordering};
5+
use std::cell::Cell;
56
use std::sync::mpsc;
67

78
pub(crate) const NSIG: usize = 64;
@@ -11,6 +12,20 @@ static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
1112
const ATOMIC_FALSE: AtomicBool = AtomicBool::new(false);
1213
pub(crate) static TRIGGERS: [AtomicBool; NSIG] = [ATOMIC_FALSE; NSIG];
1314

15+
thread_local! {
16+
/// Prevent recursive signal handler invocation. When a Python signal
17+
/// handler is running, new signals are deferred until it completes.
18+
static IN_SIGNAL_HANDLER: Cell<bool> = const { Cell::new(false) };
19+
}
20+
21+
struct SignalHandlerGuard;
22+
23+
impl Drop for SignalHandlerGuard {
24+
fn drop(&mut self) {
25+
IN_SIGNAL_HANDLER.with(|h| h.set(false));
26+
}
27+
}
28+
1429
#[cfg_attr(feature = "flame-it", flame)]
1530
#[inline(always)]
1631
pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
@@ -27,6 +42,13 @@ pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
2742
#[inline(never)]
2843
#[cold]
2944
fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
45+
if IN_SIGNAL_HANDLER.with(|h| h.replace(true)) {
46+
// Already inside a signal handler — defer pending signals
47+
set_triggered();
48+
return Ok(());
49+
}
50+
let _guard = SignalHandlerGuard;
51+
3052
// unwrap should never fail since we check above
3153
let signal_handlers = vm.signal_handlers.as_ref().unwrap().borrow();
3254
for (signum, trigger) in TRIGGERS.iter().enumerate().skip(1) {

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

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -98,20 +98,31 @@ pub fn cleanup_current_thread_frames(vm: &VirtualMachine) {
9898
}
9999

100100
/// Reinitialize frame slot after fork. Called in child process.
101-
/// Creates a fresh slot and registers it for the current thread.
101+
/// Creates a fresh slot and registers it for the current thread,
102+
/// preserving the current thread's frames from `vm.frames`.
102103
#[cfg(feature = "threading")]
103104
pub fn reinit_frame_slot_after_fork(vm: &VirtualMachine) {
104105
let current_ident = crate::stdlib::thread::get_ident();
105-
let new_slot = Arc::new(parking_lot::Mutex::new(Vec::new()));
106+
// Preserve the current thread's frames across fork
107+
let current_frames: Vec<FrameRef> = vm.frames.borrow().clone();
108+
let new_slot = Arc::new(parking_lot::Mutex::new(current_frames));
106109

107-
// Try to update the global registry. If we can't get the lock
108-
// (parent thread might have been holding it during fork), skip.
109-
if let Some(mut registry) = vm.state.thread_frames.try_lock() {
110-
registry.clear();
111-
registry.insert(current_ident, new_slot.clone());
112-
}
110+
// After fork, only the current thread exists. If the lock was held by
111+
// another thread during fork, force unlock it.
112+
let mut registry = match vm.state.thread_frames.try_lock() {
113+
Some(guard) => guard,
114+
None => {
115+
// SAFETY: After fork in child process, only the current thread
116+
// exists. The lock holder no longer exists.
117+
unsafe { vm.state.thread_frames.force_unlock() };
118+
vm.state.thread_frames.lock()
119+
}
120+
};
121+
registry.clear();
122+
registry.insert(current_ident, new_slot.clone());
123+
drop(registry);
113124

114-
// Always update thread-local to point to the new slot
125+
// Update thread-local to point to the new slot
115126
CURRENT_FRAME_SLOT.with(|s| {
116127
*s.borrow_mut() = Some(new_slot);
117128
});

0 commit comments

Comments
 (0)