1use std::io;
16use std::sync::{Arc, Condvar, Mutex};
17use std::thread::JoinHandle;
18use std::time::{Duration, Instant};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum StopCause {
23 Interrupt,
25 Terminate,
27 StdinClosed,
29 ParentGone,
31 Requested,
33}
34
35struct Inner {
36 state: Mutex<Option<StopCause>>,
37 changed: Condvar,
38}
39
40#[derive(Clone)]
42pub struct StopSource {
43 inner: Arc<Inner>,
44}
45
46#[derive(Clone)]
48pub struct StopToken {
49 inner: Arc<Inner>,
50}
51
52impl Default for StopSource {
53 fn default() -> Self {
54 Self::new()
55 }
56}
57
58impl StopSource {
59 pub fn new() -> Self {
60 Self {
61 inner: Arc::new(Inner {
62 state: Mutex::new(None),
63 changed: Condvar::new(),
64 }),
65 }
66 }
67
68 pub fn token(&self) -> StopToken {
70 StopToken {
71 inner: self.inner.clone(),
72 }
73 }
74
75 pub fn trigger(&self, cause: StopCause) -> bool {
77 let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
78 if state.is_some() {
79 return false;
80 }
81 *state = Some(cause);
82 self.inner.changed.notify_all();
83 true
84 }
85}
86
87impl StopToken {
88 pub fn cause(&self) -> Option<StopCause> {
90 *self.inner.state.lock().unwrap_or_else(|e| e.into_inner())
91 }
92
93 pub fn is_stopped(&self) -> bool {
94 self.cause().is_some()
95 }
96
97 pub fn wait(&self, timeout: Duration) -> Option<StopCause> {
100 let deadline = Instant::now() + timeout;
101 let mut state = self.inner.state.lock().unwrap_or_else(|e| e.into_inner());
102 loop {
103 if let Some(cause) = *state {
104 return Some(cause);
105 }
106 let now = Instant::now();
107 if now >= deadline {
108 return None;
109 }
110 let (next, _) = self
111 .inner
112 .changed
113 .wait_timeout(state, deadline - now)
114 .unwrap_or_else(|e| e.into_inner());
115 state = next;
116 }
117 }
118}
119
120#[cfg(unix)]
122static PENDING_SIGNAL: std::sync::atomic::AtomicI32 = std::sync::atomic::AtomicI32::new(0);
123
124#[cfg(unix)]
125const SIGINT: i32 = 2;
126#[cfg(unix)]
127const SIGTERM: i32 = 15;
128
129#[cfg(unix)]
130extern "C" fn on_unix_signal(signal: i32) {
131 PENDING_SIGNAL.store(signal, std::sync::atomic::Ordering::SeqCst);
133}
134
135#[cfg(unix)]
136extern "C" {
137 fn signal(signum: i32, handler: usize) -> usize;
138}
139
140#[cfg(windows)]
141static PENDING_CONSOLE_EVENT: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
142
143#[cfg(windows)]
144unsafe extern "system" fn on_console_event(event: u32) -> i32 {
145 PENDING_CONSOLE_EVENT.store(event + 1, std::sync::atomic::Ordering::SeqCst);
146 1 }
148
149#[cfg(windows)]
150extern "system" {
151 fn SetConsoleCtrlHandler(
152 handler: Option<unsafe extern "system" fn(u32) -> i32>,
153 add: i32,
154 ) -> i32;
155}
156
157pub fn stop_on_signals(source: &StopSource) -> io::Result<JoinHandle<()>> {
160 install_os_handlers()?;
161 let source = source.clone();
162 std::thread::Builder::new()
163 .name("rk-stop-signals".into())
164 .spawn(move || loop {
165 if source.token().is_stopped() {
166 break;
167 }
168 if let Some(cause) = take_pending_signal() {
169 source.trigger(cause);
170 break;
171 }
172 std::thread::sleep(Duration::from_millis(50));
173 })
174}
175
176pub fn stop_on_stdin_close(source: &StopSource) -> JoinHandle<()> {
178 let source = source.clone();
179 crate::private::watch_stdin_eof(move || {
180 source.trigger(StopCause::StdinClosed);
181 })
182}
183
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
190pub struct ParentFence {
191 pid: u32,
192 started: Option<(u64, u64)>,
193}
194
195impl ParentFence {
196 pub fn capture(pid: u32) -> Self {
200 Self {
201 pid,
202 started: crate::host_identity::process_start_marker(pid),
203 }
204 }
205
206 pub fn from_parts(pid: u32, started: Option<(u64, u64)>) -> Self {
209 Self { pid, started }
210 }
211
212 pub fn pid(&self) -> u32 {
213 self.pid
214 }
215
216 pub fn started(&self) -> Option<(u64, u64)> {
218 self.started
219 }
220
221 pub fn parent_gone(&self) -> bool {
225 parent_gone(
226 self.started,
227 crate::host_identity::process_start_marker(self.pid),
228 crate::transport::pid_alive(self.pid),
229 )
230 }
231}
232
233fn parent_gone(captured: Option<(u64, u64)>, current: Option<(u64, u64)>, pid_alive: bool) -> bool {
235 match captured {
236 Some(marker) => current != Some(marker),
239 None => !pid_alive,
240 }
241}
242
243pub fn stop_when_parent_gone(
247 source: &StopSource,
248 fence: ParentFence,
249 poll: Duration,
250) -> JoinHandle<()> {
251 let source = source.clone();
252 let poll = poll.max(Duration::from_millis(10));
253 std::thread::spawn(move || loop {
254 if source.token().is_stopped() {
255 break;
256 }
257 if fence.parent_gone() {
258 source.trigger(StopCause::ParentGone);
259 break;
260 }
261 std::thread::sleep(poll);
262 })
263}
264
265#[cfg(unix)]
266fn install_os_handlers() -> io::Result<()> {
267 for sig in [SIGINT, SIGTERM] {
268 let handler = on_unix_signal as extern "C" fn(i32) as usize;
269 let previous = unsafe { signal(sig, handler) };
272 if previous == usize::MAX {
273 return Err(io::Error::last_os_error());
274 }
275 }
276 Ok(())
277}
278
279#[cfg(unix)]
280fn take_pending_signal() -> Option<StopCause> {
281 match PENDING_SIGNAL.swap(0, std::sync::atomic::Ordering::SeqCst) {
282 0 => None,
283 SIGINT => Some(StopCause::Interrupt),
284 _ => Some(StopCause::Terminate),
285 }
286}
287
288#[cfg(windows)]
289fn install_os_handlers() -> io::Result<()> {
290 let ok = unsafe { SetConsoleCtrlHandler(Some(on_console_event), 1) };
292 if ok == 0 {
293 return Err(io::Error::last_os_error());
294 }
295 Ok(())
296}
297
298#[cfg(windows)]
299fn take_pending_signal() -> Option<StopCause> {
300 match PENDING_CONSOLE_EVENT.swap(0, std::sync::atomic::Ordering::SeqCst) {
303 0 => None,
304 1 | 2 => Some(StopCause::Interrupt),
305 _ => Some(StopCause::Terminate),
306 }
307}
308
309#[cfg(not(any(unix, windows)))]
310fn install_os_handlers() -> io::Result<()> {
311 Err(io::Error::new(
312 io::ErrorKind::Unsupported,
313 "stop signals are not implemented on this platform",
314 ))
315}
316
317#[cfg(not(any(unix, windows)))]
318fn take_pending_signal() -> Option<StopCause> {
319 None
320}