use anyhow::{Context, Result};
use signal_hook::consts::{SIGINT, SIGTERM};
use signal_hook::iterator::Signals;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
const SIGINT_WINDOW: Duration = Duration::from_millis(100);
#[derive(Clone, Debug, Default)]
pub struct RunState(Arc<AtomicUsize>);
impl RunState {
const RUNNING: usize = 0;
const STOPPING: usize = 1;
const SHUTDOWN: usize = 2;
pub fn new() -> Self {
Self(Arc::new(AtomicUsize::new(Self::RUNNING)))
}
fn load(&self) -> usize {
self.0.load(Ordering::SeqCst)
}
fn advance(&self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
pub fn shutdown(&self) {
self.0.store(Self::SHUTDOWN, Ordering::SeqCst);
}
pub fn interrupted(&self) -> bool {
self.load() != Self::RUNNING
}
pub(crate) fn is_stopping(&self) -> bool {
self.load() == Self::STOPPING
}
pub(crate) fn is_shutdown(&self) -> bool {
self.load() >= Self::SHUTDOWN
}
pub fn register_signals(&self) -> Result<()> {
let mut signals =
Signals::new([SIGINT, SIGTERM]).context("Failed to register signal handlers")?;
let state = self.clone();
crate::spawn_named("signals", move || {
let mut last: Option<Instant> = None;
for sig in signals.forever() {
match sig {
SIGINT => {
let now = Instant::now();
if last.is_none_or(|t| now.duration_since(t) > SIGINT_WINDOW) {
last = Some(now);
state.advance();
}
}
SIGTERM => state.shutdown(),
_ => {}
}
}
});
Ok(())
}
}