Skip to main content

rustpython_vm/
signal.rs

1#![cfg_attr(target_os = "wasi", allow(dead_code))]
2use crate::{PyObjectRef, PyResult, VirtualMachine};
3use alloc::fmt;
4use core::cell::{Cell, RefCell};
5#[cfg(windows)]
6use core::sync::atomic::AtomicIsize;
7use core::sync::atomic::{AtomicBool, Ordering};
8use std::sync::mpsc;
9
10pub(crate) const NSIG: usize = 64;
11
12pub(crate) fn new_signal_handlers() -> Box<RefCell<[Option<PyObjectRef>; NSIG]>> {
13    Box::new(const { RefCell::new([const { None }; NSIG]) })
14}
15static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
16// hack to get around const array repeat expressions, rust issue #79270
17#[allow(
18    clippy::declare_interior_mutable_const,
19    reason = "workaround for const array repeat limitation (rust issue #79270)"
20)]
21const ATOMIC_FALSE: AtomicBool = AtomicBool::new(false);
22pub(crate) static TRIGGERS: [AtomicBool; NSIG] = [ATOMIC_FALSE; NSIG];
23
24#[cfg(windows)]
25static SIGINT_EVENT: AtomicIsize = AtomicIsize::new(0);
26
27thread_local! {
28    /// Prevent recursive signal handler invocation. When a Python signal
29    /// handler is running, new signals are deferred until it completes.
30    static IN_SIGNAL_HANDLER: Cell<bool> = const { Cell::new(false) };
31}
32
33struct SignalHandlerGuard;
34
35impl Drop for SignalHandlerGuard {
36    fn drop(&mut self) {
37        IN_SIGNAL_HANDLER.with(|h| h.set(false));
38    }
39}
40
41#[cfg_attr(feature = "flame-it", flame)]
42#[inline(always)]
43pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
44    if vm.signal_handlers.get().is_none() {
45        return Ok(());
46    }
47
48    // Read-only check first: avoids cache-line invalidation on every
49    // instruction when no signal is pending (the common case).
50    if !ANY_TRIGGERED.load(Ordering::Relaxed) {
51        return Ok(());
52    }
53    // Atomic RMW only when a signal is actually pending.
54    if !ANY_TRIGGERED.swap(false, Ordering::Acquire) {
55        return Ok(());
56    }
57
58    trigger_signals(vm)
59}
60
61#[inline(never)]
62#[cold]
63fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
64    if IN_SIGNAL_HANDLER.with(|h| h.replace(true)) {
65        // Already inside a signal handler — defer pending signals
66        set_triggered();
67        return Ok(());
68    }
69    let _guard = SignalHandlerGuard;
70
71    // unwrap should never fail since we check above
72    let signal_handlers = vm.signal_handlers.get().unwrap().borrow();
73    for (signum, trigger) in TRIGGERS.iter().enumerate().skip(1) {
74        let triggered = trigger.swap(false, Ordering::Relaxed);
75        if triggered
76            && let Some(handler) = &signal_handlers[signum]
77            && let Some(callable) = handler.to_callable()
78        {
79            callable.invoke((signum, vm.ctx.none()), vm)?;
80        }
81    }
82    if let Some(signal_rx) = &vm.signal_rx {
83        for f in signal_rx.rx.try_iter() {
84            f(vm)?;
85        }
86    }
87    Ok(())
88}
89
90pub(crate) fn set_triggered() {
91    ANY_TRIGGERED.store(true, Ordering::Release);
92}
93
94#[inline(always)]
95pub(crate) fn is_triggered() -> bool {
96    ANY_TRIGGERED.load(Ordering::Relaxed)
97}
98
99/// Reset all signal trigger state after fork in child process.
100/// Stale triggers from the parent must not fire in the child.
101#[cfg(unix)]
102#[cfg(feature = "host_env")]
103pub(crate) fn clear_after_fork() {
104    ANY_TRIGGERED.store(false, Ordering::Release);
105    for trigger in &TRIGGERS {
106        trigger.store(false, Ordering::Relaxed);
107    }
108}
109
110pub fn assert_in_range(signum: i32, vm: &VirtualMachine) -> PyResult<()> {
111    if (1..NSIG as i32).contains(&signum) {
112        Ok(())
113    } else {
114        Err(vm.new_value_error("signal number out of range"))
115    }
116}
117
118/// Similar to `PyErr_SetInterruptEx` in CPython
119///
120/// Missing signal handler for the given signal number is silently ignored.
121#[allow(dead_code)]
122#[cfg(all(not(target_arch = "wasm32"), feature = "host_env"))]
123pub fn set_interrupt_ex(signum: i32, vm: &VirtualMachine) -> PyResult<()> {
124    use crate::stdlib::_signal::_signal::{SIG_DFL, SIG_IGN, run_signal};
125    assert_in_range(signum, vm)?;
126
127    match signum as usize {
128        SIG_DFL | SIG_IGN => Ok(()),
129        _ => {
130            // interrupt the main thread with given signal number
131            run_signal(signum);
132            Ok(())
133        }
134    }
135}
136
137pub type UserSignal = Box<dyn FnOnce(&VirtualMachine) -> PyResult<()> + Send>;
138
139#[derive(Clone, Debug)]
140pub struct UserSignalSender {
141    tx: mpsc::Sender<UserSignal>,
142}
143
144#[derive(Debug)]
145pub struct UserSignalReceiver {
146    rx: mpsc::Receiver<UserSignal>,
147}
148
149impl UserSignalSender {
150    pub fn send(&self, sig: UserSignal) -> Result<(), UserSignalSendError> {
151        self.tx
152            .send(sig)
153            .map_err(|mpsc::SendError(sig)| UserSignalSendError(sig))?;
154        set_triggered();
155        Ok(())
156    }
157}
158
159pub struct UserSignalSendError(pub UserSignal);
160
161impl fmt::Debug for UserSignalSendError {
162    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
163        f.debug_struct("UserSignalSendError")
164            .finish_non_exhaustive()
165    }
166}
167
168impl fmt::Display for UserSignalSendError {
169    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
170        f.write_str("sending a signal to a exited vm")
171    }
172}
173
174pub fn user_signal_channel() -> (UserSignalSender, UserSignalReceiver) {
175    let (tx, rx) = mpsc::channel();
176    (UserSignalSender { tx }, UserSignalReceiver { rx })
177}
178
179#[cfg(windows)]
180pub fn set_sigint_event(handle: isize) {
181    SIGINT_EVENT.store(handle, Ordering::Release);
182}
183
184#[cfg(windows)]
185pub fn get_sigint_event() -> Option<isize> {
186    let handle = SIGINT_EVENT.load(Ordering::Acquire);
187    if handle == 0 { None } else { Some(handle) }
188}