ntex_rt/
signals.rs

1#![allow(static_mut_refs)]
2use std::{cell::RefCell, future::poll_fn, sync::Arc, task::Poll};
3
4use atomic_waker::AtomicWaker;
5
6use crate::System;
7
8thread_local! {
9    static STOP: RefCell<Option<oneshot::Sender<()>>> = const { RefCell::new(None) };
10    static HANDLERS: RefCell<Vec<oneshot::Sender<Arc<[Signal]>>>> = RefCell::default();
11}
12
13static mut CUR_SYS: Option<System> = None;
14static mut SIGS: [Option<Signal>; 10] = [None; 10];
15static HND_WAKER: AtomicWaker = AtomicWaker::new();
16
17/// Different types of process signals
18#[derive(PartialEq, Eq, Clone, Copy, Debug)]
19pub enum Signal {
20    /// SIGHUP
21    Hup,
22    /// SIGINT
23    Int,
24    /// SIGTERM
25    Term,
26    /// SIGQUIT
27    Quit,
28}
29
30/// Register signal handler.
31///
32/// Signals are handled by oneshots, you have to re-register
33/// interest after each signal.
34pub fn signal() -> oneshot::AsyncReceiver<Arc<[Signal]>> {
35    let (tx, rx) = oneshot::async_channel();
36    System::current().handle().spawn(async move {
37        HANDLERS.with(|handlers| {
38            handlers.borrow_mut().push(tx);
39        });
40    });
41
42    rx
43}
44
45/// Check if signal handling is enabled.
46pub fn is_enabled() -> bool {
47    unsafe { CUR_SYS.is_some() }
48}
49
50fn register_system(sys: &System) -> bool {
51    unsafe {
52        if CUR_SYS.is_some() {
53            false
54        } else {
55            CUR_SYS = Some(sys.clone());
56
57            let (tx, rx) = oneshot::async_channel();
58            sys.handle().spawn(signals(rx));
59            STOP.with(|stop| {
60                *stop.borrow_mut() = Some(tx);
61            });
62            true
63        }
64    }
65}
66
67fn unregister_system(sys: &System) -> bool {
68    unsafe {
69        if let Some(cur) = CUR_SYS.take() {
70            if cur.id() == sys.id() {
71                sys.handle().spawn(async move {
72                    STOP.with(|stop| {
73                        if let Some(tx) = stop.borrow_mut().take() {
74                            let _ = tx.send(());
75                        }
76                    });
77                });
78                true
79            } else {
80                CUR_SYS = Some(cur);
81                false
82            }
83        } else {
84            false
85        }
86    }
87}
88
89fn handle_signal(sig: Signal) {
90    unsafe {
91        for s in &mut SIGS {
92            if s.is_none() {
93                *s = Some(sig);
94                break;
95            }
96        }
97        HND_WAKER.wake();
98    }
99}
100
101#[cfg(target_family = "unix")]
102static mut SIG_HANDLERS: [Option<signal_hook::SigId>; 10] = [None; 10];
103
104#[cfg(target_family = "unix")]
105/// Register signal handler.
106pub(crate) fn start(sys: &System) {
107    if register_system(sys) {
108        use signal_hook::consts::signal::{SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGUSR2};
109        use signal_hook::low_level::register;
110
111        for (idx, s, sig) in [
112            (0, SIGHUP, Signal::Hup),
113            (1, SIGINT, Signal::Int),
114            (2, SIGTERM, Signal::Term),
115            (3, SIGQUIT, Signal::Quit),
116        ] {
117            unsafe {
118                match register(s, move || handle_signal(sig)) {
119                    Ok(s) => SIG_HANDLERS[idx] = Some(s),
120                    Err(e) => {
121                        log::error!("Cannot install signal handler for {sig:?} with {e:?}");
122                    }
123                }
124            }
125        }
126
127        unsafe {
128            match register(SIGUSR2, || crate::system::sig_usr2()) {
129                Ok(s) => SIG_HANDLERS[5] = Some(s),
130                Err(_) => log::error!("Cannot install signal handler for SIGUSR2"),
131            }
132        }
133    }
134}
135
136#[cfg(target_family = "unix")]
137/// Unregister signal handler.
138pub(crate) fn stop(sys: &System) {
139    if unregister_system(sys) {
140        use signal_hook::low_level::unregister;
141
142        unsafe {
143            for sig in &mut SIG_HANDLERS {
144                if let Some(s) = sig.take() {
145                    let _ = unregister(s);
146                }
147            }
148        }
149    }
150}
151
152#[cfg(target_family = "windows")]
153/// Register signal handler.
154///
155/// Signals are handled by oneshots, you have to re-register
156/// after each signal.
157pub(crate) fn start(sys: &System) {
158    if register_system(sys) {
159        ctrlc::set_handler(move || handle_signal(Signal::Int))
160            .expect("Error setting Ctrl-C handler");
161    }
162}
163
164#[cfg(target_family = "windows")]
165/// Unregister signal handler.
166pub(crate) fn stop(sys: &System) {
167    if unregister_system(sys) {
168        log::info!("Signals handling is disabled");
169    }
170}
171
172async fn signals(rx: oneshot::AsyncReceiver<()>) {
173    let mut rx = std::pin::pin!(rx);
174
175    poll_fn(|cx| {
176        if rx.as_mut().poll(cx).is_ready() {
177            Poll::Ready(())
178        } else {
179            HND_WAKER.register(cx.waker());
180
181            let mut sigs = Vec::new();
182            unsafe {
183                for sig in &mut SIGS {
184                    if let Some(sig) = sig.take() {
185                        sigs.push(sig);
186                    }
187                }
188            }
189            if !sigs.is_empty() {
190                let sigs: Arc<[Signal]> = Arc::from(sigs);
191
192                HANDLERS.with(|handlers| {
193                    for tx in handlers.borrow_mut().drain(..) {
194                        let _ = tx.send(sigs.clone());
195                    }
196                });
197            }
198
199            Poll::Pending
200        }
201    })
202    .await;
203}
ntex_rt/signals.rs

ntex_rt/
signals.rs
