use std::io;
use std::sync::{Arc, Condvar, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StopCause {
Interrupt,
Terminate,
StdinClosed,
ParentGone,
Requested,
}
struct Inner {
state: Mutex<Option<StopCause>>,
changed: Condvar,
}
#[derive(Clone)]
pub struct StopSource {
inner: Arc<Inner>,
}
#[derive(Clone)]
pub struct StopToken {
inner: Arc<Inner>,
}
impl Default for StopSource {
fn default() -> Self {
Self::new()
}
}
impl StopSource {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
state: Mutex::new(None),
changed: Condvar::new(),
}),
}
}
pub fn token(&self) -> StopToken {
StopToken {
inner: self.inner.clone(),
}
}
pub fn trigger(&self, cause: StopCause) -> bool {
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
if state.is_some() {
return false;
}
*state = Some(cause);
self.inner.changed.notify_all();
true
}
}
impl StopToken {
pub fn cause(&self) -> Option<StopCause> {
*self.inner.state.lock().unwrap_or_else(|e| e.into_inner())
}
pub fn is_stopped(&self) -> bool {
self.cause().is_some()
}
pub fn wait(&self, timeout: Duration) -> Option<StopCause> {
let deadline = Instant::now() + timeout;
let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
loop {
if let Some(cause) = *state {
return Some(cause);
}
let now = Instant::now();
if now >= deadline {
return None;
}
let (next, _) = self
.inner
.changed
.wait_timeout(state, deadline - now)
.unwrap_or_else(|e| e.into_inner());
state = next;
}
}
}
#[cfg(unix)]
static PENDING_SIGNAL: std::sync::atomic::AtomicI32 = std::sync::atomic::AtomicI32::new(0);
#[cfg(unix)]
const SIGINT: i32 = 2;
#[cfg(unix)]
const SIGTERM: i32 = 15;
#[cfg(unix)]
extern "C" fn on_unix_signal(signal: i32) {
PENDING_SIGNAL.store(signal, std::sync::atomic::Ordering::SeqCst);
}
#[cfg(unix)]
extern "C" {
fn signal(signum: i32, handler: usize) -> usize;
}
#[cfg(windows)]
static PENDING_CONSOLE_EVENT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
#[cfg(windows)]
unsafe extern "system" fn on_console_event(event: u32) -> i32 {
PENDING_CONSOLE_EVENT.store(event + 1, std::sync::atomic::Ordering::SeqCst);
1 }
#[cfg(windows)]
extern "system" {
fn SetConsoleCtrlHandler(
handler: Option<unsafe extern "system" fn(u32) -> i32>,
add: i32,
) -> i32;
}
pub fn stop_on_signals(source: &StopSource) -> io::Result<JoinHandle<()>> {
install_os_handlers()?;
let source = source.clone();
std::thread::Builder::new()
.name("rk-stop-signals".into())
.spawn(move || loop {
if source.token().is_stopped() {
break;
}
if let Some(cause) = take_pending_signal() {
source.trigger(cause);
break;
}
std::thread::sleep(Duration::from_millis(50));
})
}
pub fn stop_on_stdin_close(source: &StopSource) -> JoinHandle<()> {
let source = source.clone();
crate::private::watch_stdin_eof(move || {
source.trigger(StopCause::StdinClosed);
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ParentFence {
pid: u32,
started: Option<(u64, u64)>,
}
impl ParentFence {
pub fn capture(pid: u32) -> Self {
Self {
pid,
started: crate::host_identity::process_start_marker(pid),
}
}
pub fn from_parts(pid: u32, started: Option<(u64, u64)>) -> Self {
Self { pid, started }
}
pub fn pid(&self) -> u32 {
self.pid
}
pub fn started(&self) -> Option<(u64, u64)> {
self.started
}
pub fn parent_gone(&self) -> bool {
parent_gone(
self.started,
crate::host_identity::process_start_marker(self.pid),
crate::transport::pid_alive(self.pid),
)
}
}
fn parent_gone(captured: Option<(u64, u64)>, current: Option<(u64, u64)>, pid_alive: bool) -> bool {
match captured {
Some(marker) => current != Some(marker),
None => !pid_alive,
}
}
pub fn stop_when_parent_gone(
source: &StopSource,
fence: ParentFence,
poll: Duration,
) -> JoinHandle<()> {
let source = source.clone();
let poll = poll.max(Duration::from_millis(10));
std::thread::spawn(move || loop {
if source.token().is_stopped() {
break;
}
if fence.parent_gone() {
source.trigger(StopCause::ParentGone);
break;
}
std::thread::sleep(poll);
})
}
#[cfg(unix)]
fn install_os_handlers() -> io::Result<()> {
for sig in [SIGINT, SIGTERM] {
let handler = on_unix_signal as extern "C" fn(i32) as usize;
let previous = unsafe { signal(sig, handler) };
if previous == usize::MAX {
return Err(io::Error::last_os_error());
}
}
Ok(())
}
#[cfg(unix)]
fn take_pending_signal() -> Option<StopCause> {
match PENDING_SIGNAL.swap(0, std::sync::atomic::Ordering::SeqCst) {
0 => None,
SIGINT => Some(StopCause::Interrupt),
_ => Some(StopCause::Terminate),
}
}
#[cfg(windows)]
fn install_os_handlers() -> io::Result<()> {
let ok = unsafe { SetConsoleCtrlHandler(Some(on_console_event), 1) };
if ok == 0 {
return Err(io::Error::last_os_error());
}
Ok(())
}
#[cfg(windows)]
fn take_pending_signal() -> Option<StopCause> {
match PENDING_CONSOLE_EVENT.swap(0, std::sync::atomic::Ordering::SeqCst) {
0 => None,
1 | 2 => Some(StopCause::Interrupt),
_ => Some(StopCause::Terminate),
}
}
#[cfg(not(any(unix, windows)))]
fn install_os_handlers() -> io::Result<()> {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"stop signals are not implemented on this platform",
))
}
#[cfg(not(any(unix, windows)))]
fn take_pending_signal() -> Option<StopCause> {
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_first_trigger_wins() {
let source = StopSource::new();
assert!(source.trigger(StopCause::Terminate));
assert!(!source.trigger(StopCause::StdinClosed));
assert_eq!(source.token().cause(), Some(StopCause::Terminate));
}
#[test]
fn a_reused_parent_pid_is_gone_but_the_same_process_is_not() {
let mine = Some((10, 20));
assert!(!parent_gone(mine, mine, true));
assert!(parent_gone(mine, Some((99, 1)), true));
assert!(parent_gone(mine, None, false));
assert!(!parent_gone(None, None, true));
assert!(parent_gone(None, None, false));
}
#[test]
fn a_fence_on_this_process_is_not_gone_and_a_watcher_fires_for_a_dead_parent() {
let fence = ParentFence::capture(std::process::id());
assert_eq!(fence.pid(), std::process::id());
assert!(!fence.parent_gone());
let reused = ParentFence::from_parts(std::process::id(), Some((u64::MAX, u64::MAX)));
if fence.started().is_some() {
assert!(reused.parent_gone());
let source = StopSource::new();
let token = source.token();
let watcher = stop_when_parent_gone(&source, reused, Duration::from_millis(10));
assert_eq!(
token.wait(Duration::from_secs(5)),
Some(StopCause::ParentGone)
);
watcher.join().unwrap();
}
}
#[test]
fn a_wait_times_out_until_triggered() {
let source = StopSource::new();
let token = source.token();
assert_eq!(token.wait(Duration::from_millis(5)), None);
assert!(!token.is_stopped());
let clone = source.clone();
let t = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(10));
clone.trigger(StopCause::Requested);
});
assert_eq!(
token.wait(Duration::from_secs(5)),
Some(StopCause::Requested)
);
t.join().unwrap();
}
}