use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
#[derive(Debug, Default)]
pub struct ShutdownSignal {
triggered: Arc<AtomicBool>,
}
impl Clone for ShutdownSignal {
fn clone(&self) -> Self {
Self {
triggered: Arc::clone(&self.triggered),
}
}
}
impl ShutdownSignal {
pub fn new() -> Self {
Self::default()
}
pub fn is_triggered(&self) -> bool {
self.triggered.load(Ordering::SeqCst)
}
pub fn trigger(&self) {
self.triggered.store(true, Ordering::SeqCst);
tracing::info!("shutdown signal triggered");
}
#[cfg(unix)]
#[allow(unsafe_code)] pub fn install(&self, _drain_timeout: Duration) {
let signal = self.clone();
unsafe {
libc_signal(libc_signum::SIGINT, move || {
signal.trigger();
});
let signal2 = self.clone();
libc_signal(libc_signum::SIGTERM, move || {
signal2.trigger();
});
}
}
#[cfg(not(unix))]
pub fn install(&self, _drain_timeout: Duration) {}
pub fn wait_for_trigger(&self, timeout: Duration) -> bool {
let deadline = Instant::now() + timeout;
while Instant::now() < deadline {
if self.is_triggered() {
return true;
}
std::thread::sleep(Duration::from_millis(100));
}
self.is_triggered()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DrainResult {
Drained { sessions_completed: usize },
TimedOut { force_expired: usize },
}
pub struct ShutdownCoordinator {
drain_timeout: Duration,
}
impl ShutdownCoordinator {
pub fn new(drain_timeout: Duration) -> Self {
Self { drain_timeout }
}
pub fn drain_timeout(&self) -> Duration {
self.drain_timeout
}
pub fn await_signal(&self, signal: &ShutdownSignal) -> bool {
signal.wait_for_trigger(self.drain_timeout)
}
}
#[cfg(unix)]
mod libc_signum {
pub const SIGINT: i32 = 2;
pub const SIGTERM: i32 = 15;
}
#[cfg(unix)]
#[allow(unsafe_code)] unsafe fn libc_signal(signum: i32, handler: impl Fn() + Send + 'static) {
use std::sync::OnceLock;
static FLAG: OnceLock<Arc<AtomicBool>> = OnceLock::new();
let _ = signum;
let _ = handler;
}
#[cfg(not(unix))]
mod libc_signum {
pub const SIGINT: i32 = 0;
pub const SIGTERM: i32 = 0;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_signal_starts_untriggered() {
let signal = ShutdownSignal::new();
assert!(!signal.is_triggered());
}
#[test]
fn trigger_sets_flag() {
let signal = ShutdownSignal::new();
signal.trigger();
assert!(signal.is_triggered());
}
#[test]
fn clone_shares_state() {
let signal = ShutdownSignal::new();
let clone = signal.clone();
signal.trigger();
assert!(clone.is_triggered());
}
#[test]
fn trigger_is_idempotent() {
let signal = ShutdownSignal::new();
signal.trigger();
signal.trigger();
signal.trigger();
assert!(signal.is_triggered());
}
#[test]
fn wait_for_trigger_returns_immediately_if_triggered() {
let signal = ShutdownSignal::new();
signal.trigger();
let result = signal.wait_for_trigger(Duration::from_secs(10));
assert!(result);
}
#[test]
fn wait_for_trigger_times_out() {
let signal = ShutdownSignal::new();
let result = signal.wait_for_trigger(Duration::from_millis(50));
assert!(!result);
}
#[test]
fn wait_for_trigger_returns_after_external_trigger() {
let signal = ShutdownSignal::new();
let clone = signal.clone();
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(50));
clone.trigger();
});
let result = signal.wait_for_trigger(Duration::from_secs(2));
assert!(result);
}
#[test]
fn shutdown_coordinator_has_drain_timeout() {
let sc = ShutdownCoordinator::new(Duration::from_secs(30));
assert_eq!(sc.drain_timeout(), Duration::from_secs(30));
}
#[test]
fn drain_result_drained_variant() {
let r = DrainResult::Drained {
sessions_completed: 5,
};
assert_eq!(
r,
DrainResult::Drained {
sessions_completed: 5
}
);
}
#[test]
fn drain_result_timed_out_variant() {
let r = DrainResult::TimedOut { force_expired: 2 };
assert_eq!(r, DrainResult::TimedOut { force_expired: 2 });
}
#[test]
fn install_does_not_panic_on_unix() {
let signal = ShutdownSignal::new();
signal.install(Duration::from_secs(5));
}
}