|
| 1 | +//! Capture who sent SIGTERM/SIGINT (`si_pid`) without fighting tokio's waiter. |
| 2 | +//! |
| 3 | +//! A `SA_SIGINFO` handler records the sender, then writes one byte to a socket |
| 4 | +//! pair so the async setup task can log and exit. We do not use `tokio::signal` |
| 5 | +//! on Unix — the last `sigaction` wins, and we need `si_pid`. |
| 6 | +
|
| 7 | +use std::os::fd::AsRawFd; |
| 8 | +use std::os::unix::net::UnixStream; |
| 9 | +use std::sync::atomic::{AtomicI32, Ordering}; |
| 10 | + |
| 11 | +use tracing::info; |
| 12 | + |
| 13 | +static SENDER_PID: AtomicI32 = AtomicI32::new(0); |
| 14 | +static SIGNAL_NO: AtomicI32 = AtomicI32::new(0); |
| 15 | +static SIGNAL_CODE: AtomicI32 = AtomicI32::new(0); |
| 16 | +static PIPE_WRITE_FD: AtomicI32 = AtomicI32::new(-1); |
| 17 | + |
| 18 | +/// Install the recorder and wait until SIGTERM or SIGINT arrives. |
| 19 | +/// |
| 20 | +/// Returns `true` after the first termination signal. Returns `false` if the |
| 21 | +/// handler could not be installed — the caller must not treat that as exit. |
| 22 | +pub async fn wait_for_term() -> bool { |
| 23 | + let read_end = match install() { |
| 24 | + Ok(fd) => fd, |
| 25 | + Err(e) => { |
| 26 | + tracing::warn!("[Signal] Failed to install SA_SIGINFO handler: {e}"); |
| 27 | + return false; |
| 28 | + } |
| 29 | + }; |
| 30 | + |
| 31 | + let mut reader = match tokio::net::UnixStream::from_std(read_end) { |
| 32 | + Ok(s) => s, |
| 33 | + Err(e) => { |
| 34 | + tracing::warn!("[Signal] Failed to wrap self-pipe: {e}"); |
| 35 | + return false; |
| 36 | + } |
| 37 | + }; |
| 38 | + let mut buf = [0u8; 1]; |
| 39 | + use tokio::io::AsyncReadExt; |
| 40 | + let _ = reader.read_exact(&mut buf).await; |
| 41 | + |
| 42 | + let pid = std::process::id(); |
| 43 | + let ppid = unsafe { libc::getppid() }; |
| 44 | + let sender = SENDER_PID.load(Ordering::SeqCst); |
| 45 | + let sig = SIGNAL_NO.load(Ordering::SeqCst); |
| 46 | + let code = SIGNAL_CODE.load(Ordering::SeqCst); |
| 47 | + let name = if sig == libc::SIGINT { |
| 48 | + "SIGINT" |
| 49 | + } else { |
| 50 | + "SIGTERM" |
| 51 | + }; |
| 52 | + |
| 53 | + info!( |
| 54 | + pid, |
| 55 | + ppid, |
| 56 | + parent = %describe_pid(ppid), |
| 57 | + sender_pid = sender, |
| 58 | + sender = %describe_pid(sender), |
| 59 | + si_code = code, |
| 60 | + "[Signal] {name} — requesting exit" |
| 61 | + ); |
| 62 | + true |
| 63 | +} |
| 64 | + |
| 65 | +fn install() -> std::io::Result<UnixStream> { |
| 66 | + let (read, write) = UnixStream::pair()?; |
| 67 | + read.set_nonblocking(true)?; |
| 68 | + write.set_nonblocking(true)?; |
| 69 | + PIPE_WRITE_FD.store(write.as_raw_fd(), Ordering::SeqCst); |
| 70 | + // Handler writes to this fd for the life of the process. |
| 71 | + std::mem::forget(write); |
| 72 | + |
| 73 | + unsafe { |
| 74 | + let mut sa: libc::sigaction = std::mem::zeroed(); |
| 75 | + sa.sa_sigaction = record_sender as *const () as usize; |
| 76 | + sa.sa_flags = libc::SA_SIGINFO | libc::SA_RESTART; |
| 77 | + libc::sigemptyset(&mut sa.sa_mask); |
| 78 | + if libc::sigaction(libc::SIGTERM, &sa, std::ptr::null_mut()) != 0 { |
| 79 | + return Err(std::io::Error::last_os_error()); |
| 80 | + } |
| 81 | + if libc::sigaction(libc::SIGINT, &sa, std::ptr::null_mut()) != 0 { |
| 82 | + return Err(std::io::Error::last_os_error()); |
| 83 | + } |
| 84 | + } |
| 85 | + Ok(read) |
| 86 | +} |
| 87 | + |
| 88 | +unsafe extern "C" fn record_sender( |
| 89 | + sig: libc::c_int, |
| 90 | + info: *mut libc::siginfo_t, |
| 91 | + _ctx: *mut libc::c_void, |
| 92 | +) { |
| 93 | + if !info.is_null() { |
| 94 | + // SAFETY: kernel-filled siginfo for this delivery. |
| 95 | + let info = unsafe { &*info }; |
| 96 | + SENDER_PID.store(unsafe { info.si_pid() }, Ordering::SeqCst); |
| 97 | + SIGNAL_CODE.store(info.si_code, Ordering::SeqCst); |
| 98 | + } |
| 99 | + SIGNAL_NO.store(sig, Ordering::SeqCst); |
| 100 | + let fd = PIPE_WRITE_FD.load(Ordering::SeqCst); |
| 101 | + if fd >= 0 { |
| 102 | + let byte = [1u8]; |
| 103 | + unsafe { |
| 104 | + libc::write(fd, byte.as_ptr() as *const libc::c_void, 1); |
| 105 | + } |
| 106 | + } |
| 107 | +} |
| 108 | + |
| 109 | +/// Best-effort path/name for a live pid. `"unknown"` if pid is empty or gone. |
| 110 | +fn describe_pid(pid: libc::pid_t) -> String { |
| 111 | + if pid <= 0 { |
| 112 | + return "unknown".to_string(); |
| 113 | + } |
| 114 | + #[cfg(target_os = "macos")] |
| 115 | + { |
| 116 | + let mut buf = [0u8; 4096]; |
| 117 | + extern "C" { |
| 118 | + fn proc_pidpath(pid: i32, buffer: *mut libc::c_void, buffersize: u32) -> i32; |
| 119 | + } |
| 120 | + let n = unsafe { proc_pidpath(pid, buf.as_mut_ptr() as *mut _, buf.len() as u32) }; |
| 121 | + if n > 0 { |
| 122 | + return String::from_utf8_lossy(&buf[..n as usize]).into_owned(); |
| 123 | + } |
| 124 | + } |
| 125 | + #[cfg(target_os = "linux")] |
| 126 | + { |
| 127 | + if let Ok(path) = std::fs::read_link(format!("/proc/{pid}/exe")) { |
| 128 | + return path.display().to_string(); |
| 129 | + } |
| 130 | + if let Ok(comm) = std::fs::read_to_string(format!("/proc/{pid}/comm")) { |
| 131 | + return comm.trim().to_string(); |
| 132 | + } |
| 133 | + } |
| 134 | + format!("pid:{pid}") |
| 135 | +} |
| 136 | + |
| 137 | +#[cfg(test)] |
| 138 | +mod tests { |
| 139 | + use super::*; |
| 140 | + |
| 141 | + #[test] |
| 142 | + fn describe_pid_self_is_nonempty() { |
| 143 | + let me = std::process::id() as libc::pid_t; |
| 144 | + let desc = describe_pid(me); |
| 145 | + assert!(!desc.is_empty(), "self pid should resolve"); |
| 146 | + assert_ne!(desc, "unknown"); |
| 147 | + } |
| 148 | + |
| 149 | + #[test] |
| 150 | + fn describe_pid_zero_is_unknown() { |
| 151 | + assert_eq!(describe_pid(0), "unknown"); |
| 152 | + } |
| 153 | +} |
0 commit comments