Skip to main content

ntex_rt/
signals.rs

1#![allow(static_mut_refs)]
2use std::{cell::RefCell, future::poll_fn, panic, sync::Arc, task::Poll};
3
4use atomic_waker::AtomicWaker;
5use ntex_error::Backtrace;
6
7use crate::System;
8
9thread_local! {
10    static STOP: RefCell<Option<oneshot::Sender<()>>> = const { RefCell::new(None) };
11    static HANDLERS: RefCell<Vec<oneshot::Sender<Arc<[Signal]>>>> = RefCell::default();
12}
13
14static mut CUR_SYS: Option<System> = None;
15static mut SIGS: [Option<Signal>; 10] = [const { None }; 10];
16static HND_WAKER: AtomicWaker = AtomicWaker::new();
17
18/// Different types of process signals
19#[derive(Clone, Debug)]
20pub enum Signal {
21    /// SIGHUP
22    Hup,
23    /// SIGINT
24    Int,
25    /// SIGTERM
26    Term,
27    /// SIGQUIT
28    Quit,
29    /// Application panic
30    Panic(PanicSource),
31}
32
33/// Different types of panics
34#[derive(Clone, Debug)]
35pub enum PanicSource {
36    /// SIGSEGV or SIGABRT is received
37    Sig(&'static str),
38    /// Application panic
39    App(Arc<str>, Backtrace),
40}
41
42/// Register signal handler.
43///
44/// Signals are handled by oneshots, you have to re-register
45/// interest after each signal.
46pub fn signal() -> oneshot::AsyncReceiver<Arc<[Signal]>> {
47    let (tx, rx) = oneshot::async_channel();
48    System::current().handle().spawn(async move {
49        HANDLERS.with(|handlers| {
50            handlers.borrow_mut().push(tx);
51        });
52    });
53
54    rx
55}
56
57/// Check if signal handling is enabled.
58pub fn is_enabled() -> bool {
59    unsafe { CUR_SYS.is_some() }
60}
61
62fn register_system(sys: &System) -> bool {
63    unsafe {
64        if CUR_SYS.is_some() {
65            false
66        } else {
67            CUR_SYS = Some(sys.clone());
68
69            let (tx, rx) = oneshot::async_channel();
70            sys.handle().spawn(signals(rx));
71            STOP.with(|stop| {
72                *stop.borrow_mut() = Some(tx);
73            });
74            true
75        }
76    }
77}
78
79fn unregister_system(sys: &System) -> bool {
80    unsafe {
81        if let Some(cur) = CUR_SYS.take() {
82            if cur.id() == sys.id() {
83                sys.handle().spawn(async move {
84                    STOP.with(|stop| {
85                        if let Some(tx) = stop.borrow_mut().take() {
86                            let _ = tx.send(());
87                        }
88                    });
89                });
90                true
91            } else {
92                CUR_SYS = Some(cur);
93                false
94            }
95        } else {
96            false
97        }
98    }
99}
100
101fn handle_signal(sig: Signal) {
102    unsafe {
103        for s in &mut SIGS {
104            if s.is_none() {
105                *s = Some(sig);
106                break;
107            }
108        }
109        HND_WAKER.wake();
110    }
111}
112
113#[cfg(target_family = "unix")]
114static mut SIG_HANDLERS: [Option<signal_hook::SigId>; 10] = [None; 10];
115
116#[cfg(target_family = "unix")]
117/// Register signal handler.
118pub(crate) fn start(sys: &System) {
119    static ONCE: std::sync::Once = std::sync::Once::new();
120
121    if register_system(sys) {
122        use nix::sys::signal;
123        use signal_hook::consts::signal::{SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGUSR2};
124        use signal_hook::low_level::register;
125
126        ONCE.call_once(|| {
127            // Use u128 for alignment.
128            let buf = Vec::leak(vec![0u128; 4096]);
129            let stack = libc::stack_t {
130                ss_sp: buf.as_ptr() as *mut libc::c_void,
131                ss_flags: 0,
132                ss_size: std::mem::size_of_val(buf),
133            };
134            let mut old = libc::stack_t {
135                ss_sp: std::ptr::null_mut(),
136                ss_flags: 0,
137                ss_size: 0,
138            };
139            let result = unsafe { libc::sigaltstack(&raw const stack, &raw mut old) };
140            if result != 0 {
141                log::error!("Cannot set signal stack");
142            }
143
144            let sig_action = signal::SigAction::new(
145                signal::SigHandler::Handler(sig_segv),
146                signal::SaFlags::SA_NODEFER | signal::SaFlags::SA_ONSTACK,
147                signal::SigSet::empty(),
148            );
149            unsafe {
150                if signal::sigaction(signal::SIGSEGV, &sig_action).is_err() {
151                    log::error!("Cannot install signal handler for SIGSEGV");
152                }
153                if signal::sigaction(signal::SIGABRT, &sig_action).is_err() {
154                    log::error!("Cannot install signal handler for SIGABRT");
155                }
156            }
157        });
158
159        for (idx, s, sig) in [
160            (0, SIGHUP, Signal::Hup),
161            (1, SIGINT, Signal::Int),
162            (2, SIGTERM, Signal::Term),
163            (3, SIGQUIT, Signal::Quit),
164        ] {
165            unsafe {
166                let sig2 = sig.clone();
167                match register(s, move || handle_signal(sig.clone())) {
168                    Ok(s) => SIG_HANDLERS[idx] = Some(s),
169                    Err(e) => {
170                        log::error!(
171                            "Cannot install signal handler for {sig2:?} with {e:?}"
172                        );
173                    }
174                }
175            }
176        }
177
178        unsafe {
179            match register(SIGUSR2, || crate::system::sig_usr2()) {
180                Ok(s) => SIG_HANDLERS[5] = Some(s),
181                Err(_) => log::error!("Cannot install signal handler for SIGUSR2"),
182            }
183        }
184    }
185}
186
187#[cfg(target_family = "unix")]
188/// Unregister signal handler.
189pub(crate) fn stop(sys: &System) {
190    if unregister_system(sys) {
191        use signal_hook::low_level::unregister;
192
193        unsafe {
194            for sig in &mut SIG_HANDLERS {
195                if let Some(s) = sig.take() {
196                    let _ = unregister(s);
197                }
198            }
199        }
200    }
201}
202
203#[cfg(target_family = "windows")]
204/// Register signal handler.
205///
206/// Signals are handled by oneshots, you have to re-register
207/// after each signal.
208pub(crate) fn start(sys: &System) {
209    if register_system(sys) {
210        ctrlc::set_handler(move || handle_signal(Signal::Int))
211            .expect("Error setting Ctrl-C handler");
212    }
213}
214
215#[cfg(target_family = "windows")]
216/// Unregister signal handler.
217pub(crate) fn stop(sys: &System) {
218    if unregister_system(sys) {
219        log::info!("Signals handling is disabled");
220    }
221}
222
223async fn signals(rx: oneshot::AsyncReceiver<()>) {
224    let mut rx = std::pin::pin!(rx);
225
226    poll_fn(|cx| {
227        if rx.as_mut().poll(cx).is_ready() {
228            Poll::Ready(())
229        } else {
230            HND_WAKER.register(cx.waker());
231
232            let mut sigs = Vec::new();
233            unsafe {
234                for sig in &mut SIGS {
235                    if let Some(sig) = sig.take() {
236                        sigs.push(sig);
237                    }
238                }
239            }
240            if !sigs.is_empty() {
241                let sigs: Arc<[Signal]> = Arc::from(sigs);
242
243                HANDLERS.with(|handlers| {
244                    for tx in handlers.borrow_mut().drain(..) {
245                        let _ = tx.send(sigs.clone());
246                    }
247                });
248            }
249
250            Poll::Pending
251        }
252    })
253    .await;
254}
255
256#[cfg(target_family = "unix")]
257extern "C" fn sig_segv(v: i32) {
258    if v == 6 {
259        eprintln!("SIGABRT Received:\n{:?}", backtrace::Backtrace::new());
260        handle_signal(Signal::Panic(PanicSource::Sig("SIGABRT")));
261    } else {
262        eprintln!("SIGSEGV Received:\n{:?}", backtrace::Backtrace::new());
263        handle_signal(Signal::Panic(PanicSource::Sig("SIGSEGV")));
264    }
265}
266
267pub(crate) fn enable_panic_handling() {
268    panic::set_hook(Box::new(|panic_info| {
269        let info: Arc<str> = if let Some(s) = panic_info.payload().downcast_ref::<&str>() {
270            Arc::from(s.to_string())
271        } else if let Some(s) = panic_info.payload().downcast_ref::<String>() {
272            Arc::from(s.clone())
273        } else {
274            Arc::from("panic")
275        };
276        let bt = if let Some(loc) = panic_info.location() {
277            let s = Box::new(loc.file().to_string());
278            let filename = Box::leak(s);
279            Backtrace::with_filename(filename)
280        } else {
281            Backtrace::new(panic::Location::caller())
282        };
283
284        handle_signal(Signal::Panic(PanicSource::App(info, bt)));
285    }));
286}