use std::sync::Arc;
#[cfg(unix)]
use std::sync::atomic::AtomicBool;
use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
#[cfg(unix)]
use std::time::Duration;
use tokio_util::sync::CancellationToken;
#[cfg(windows)]
use tokio::signal::windows::{CtrlBreak, CtrlC};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SignalAction {
Graceful,
Force,
}
static SUPPRESSED_INTERRUPTS: AtomicUsize = AtomicUsize::new(0);
#[cfg(unix)]
const SIGNAL_POLL_INTERVAL: Duration = Duration::from_millis(10);
pub struct SuppressInterruptGuard;
pub fn suppress_interrupts() -> SuppressInterruptGuard {
SUPPRESSED_INTERRUPTS.fetch_add(1, Ordering::AcqRel);
SuppressInterruptGuard
}
impl Drop for SuppressInterruptGuard {
fn drop(&mut self) {
SUPPRESSED_INTERRUPTS.fetch_sub(1, Ordering::AcqRel);
}
}
fn next_signal_action(signals_seen: &AtomicU8) -> SignalAction {
match signals_seen.fetch_add(1, Ordering::AcqRel) {
0 => SignalAction::Graceful,
_ => SignalAction::Force,
}
}
#[derive(Clone, Debug)]
pub struct Coordinator {
token: CancellationToken,
signals_seen: Arc<AtomicU8>,
}
impl Coordinator {
pub fn install() -> Self {
let coordinator = Self {
token: CancellationToken::new(),
signals_seen: Arc::new(AtomicU8::new(0)),
};
#[cfg(unix)]
install_unix_signals(&coordinator);
#[cfg(windows)]
{
let (ctrl_c, ctrl_break) = install_windows_signals();
let listener = coordinator.clone();
tokio::spawn(async move { listener.listen(ctrl_c, ctrl_break).await });
}
coordinator
}
pub fn token(&self) -> CancellationToken {
self.token.clone()
}
fn received_signal(&self, signal: i32) {
#[cfg(unix)]
if signal == libc::SIGINT && SUPPRESSED_INTERRUPTS.load(Ordering::Acquire) > 0 {
return;
}
#[cfg(windows)]
if signal == 0 && SUPPRESSED_INTERRUPTS.load(Ordering::Acquire) > 0 {
return;
}
#[cfg(unix)]
if signal == libc::SIGHUP {
std::process::exit(exit_code(signal));
}
match next_signal_action(&self.signals_seen) {
SignalAction::Graceful => self.token.cancel(),
SignalAction::Force => std::process::exit(exit_code(signal)),
}
}
#[cfg(windows)]
async fn listen(self, mut ctrl_c: CtrlC, mut ctrl_break: CtrlBreak) {
loop {
tokio::select! {
_ = ctrl_c.recv() => self.received_signal(0),
_ = ctrl_break.recv() => self.received_signal(1),
}
}
}
}
#[cfg(unix)]
fn install_unix_signals(coordinator: &Coordinator) {
let requested = Arc::new(AtomicBool::new(false));
for signal in [libc::SIGINT, libc::SIGTERM] {
signal_hook::flag::register_conditional_shutdown(
signal,
exit_code(signal),
requested.clone(),
)
.expect("install forced termination signal handler");
let requested = requested.clone();
unsafe {
signal_hook::low_level::register(signal, move || {
if signal != libc::SIGINT || SUPPRESSED_INTERRUPTS.load(Ordering::Acquire) == 0 {
requested.store(true, Ordering::SeqCst);
}
})
}
.expect("install graceful termination signal handler");
}
signal_hook::flag::register_conditional_shutdown(
libc::SIGHUP,
exit_code(libc::SIGHUP),
Arc::new(AtomicBool::new(true)),
)
.expect("install hangup signal handler");
let listener = coordinator.clone();
std::thread::Builder::new()
.name("hel-termination".to_string())
.spawn(move || {
while !requested.load(Ordering::Acquire) {
std::thread::sleep(SIGNAL_POLL_INTERVAL);
}
listener.received_signal(0);
})
.expect("spawn termination signal listener");
}
#[cfg(windows)]
fn install_windows_signals() -> (CtrlC, CtrlBreak) {
use tokio::signal::windows::{ctrl_break, ctrl_c};
(
ctrl_c().expect("install Ctrl-C listener"),
ctrl_break().expect("install Ctrl-Break listener"),
)
}
#[cfg(unix)]
const fn exit_code(signal: i32) -> i32 {
128 + signal
}
#[cfg(not(unix))]
const fn exit_code(_signal: i32) -> i32 {
1
}
#[cfg(test)]
mod tests {
use super::*;
static INTERRUPT_SUPPRESSION_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
#[test]
fn first_then_repeated_signal_transitions_to_force() {
let signals_seen = AtomicU8::new(0);
assert_eq!(next_signal_action(&signals_seen), SignalAction::Graceful);
assert_eq!(next_signal_action(&signals_seen), SignalAction::Force);
}
#[test]
fn interrupt_suppression_is_scoped() {
let _lock = INTERRUPT_SUPPRESSION_TEST_LOCK.lock().unwrap();
assert_eq!(SUPPRESSED_INTERRUPTS.load(Ordering::Acquire), 0);
{
let _guard = suppress_interrupts();
assert_eq!(SUPPRESSED_INTERRUPTS.load(Ordering::Acquire), 1);
}
assert_eq!(SUPPRESSED_INTERRUPTS.load(Ordering::Acquire), 0);
}
#[cfg(unix)]
#[test]
fn suppressed_interrupt_does_not_advance_shutdown() {
let _lock = INTERRUPT_SUPPRESSION_TEST_LOCK.lock().unwrap();
let coordinator = Coordinator {
token: CancellationToken::new(),
signals_seen: Arc::new(AtomicU8::new(0)),
};
let guard = suppress_interrupts();
coordinator.received_signal(libc::SIGINT);
assert!(!coordinator.token().is_cancelled());
assert_eq!(coordinator.signals_seen.load(Ordering::Acquire), 0);
drop(guard);
coordinator.received_signal(libc::SIGINT);
assert!(coordinator.token().is_cancelled());
assert_eq!(coordinator.signals_seen.load(Ordering::Acquire), 1);
}
#[tokio::test]
async fn coordinator_cancellation_fans_out_to_late_subscribers() {
let lock = INTERRUPT_SUPPRESSION_TEST_LOCK.lock().unwrap();
let coordinator = Coordinator {
token: CancellationToken::new(),
signals_seen: Arc::new(AtomicU8::new(0)),
};
let early = coordinator.token().child_token();
coordinator.received_signal(0);
drop(lock);
let late = coordinator.token().child_token();
early.cancelled().await;
late.cancelled().await;
}
#[cfg(windows)]
#[tokio::test]
async fn windows_signal_streams_register_synchronously() {
let (_ctrl_c, _ctrl_break) = install_windows_signals();
}
}