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
}