Skip to main content

rightkit_qa/
process.rs

1//! Owned processes: every child is spawned through `rightkit-child-process` (process
2//! group / Job Object), registered with a [`Tracker`], and killed by recorded pid
3//! only. Nothing here ever matches processes by name for killing.
4use crate::util::{err, Result};
5use rightkit_child_process::OwnedCommand;
6use std::collections::BTreeMap;
7use std::io::Read;
8use std::path::PathBuf;
9use std::process::{Command, Stdio};
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, Instant};
12
13pub fn is_alive(pid: u32) -> bool {
14    if pid == 0 {
15        return false;
16    }
17    #[cfg(unix)]
18    {
19        Command::new("kill")
20            .args(["-0", &pid.to_string()])
21            .stdout(Stdio::null())
22            .stderr(Stdio::null())
23            .status()
24            .map(|s| s.success())
25            .unwrap_or(false)
26    }
27    #[cfg(windows)]
28    {
29        // Native query: spawning `tasklist.exe` cost several hundred ms per call and was made a
30        // dozen times per `stop` (2.8 s of a 3 s stop). Fail closed: a process that exists but cannot
31        // be opened (access denied) counts as alive.
32        use windows_sys::Win32::Foundation::{CloseHandle, GetLastError, ERROR_ACCESS_DENIED};
33        use windows_sys::Win32::System::Threading::{
34            GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
35        };
36        const STILL_ACTIVE: u32 = 259;
37        unsafe {
38            let h = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
39            if h.is_null() {
40                return GetLastError() == ERROR_ACCESS_DENIED;
41            }
42            let mut code = 0u32;
43            let ok = GetExitCodeProcess(h, &mut code) != 0;
44            CloseHandle(h);
45            !ok || code == STILL_ACTIVE
46        }
47    }
48}
49
50#[cfg(windows)]
51fn hidden(cmd: &mut Command) -> &mut Command {
52    use std::os::windows::process::CommandExt;
53    cmd.creation_flags(0x0800_0000)
54}
55
56#[cfg(unix)]
57fn children_of(pid: u32) -> Vec<u32> {
58    let out = Command::new("pgrep")
59        .args(["-P", &pid.to_string()])
60        .output();
61    out.map(|o| {
62        String::from_utf8_lossy(&o.stdout)
63            .lines()
64            .filter_map(|l| l.trim().parse().ok())
65            .collect()
66    })
67    .unwrap_or_default()
68}
69
70/// Terminate `pid` and its descendants. TERM first, KILL after a short grace.
71pub fn kill_tree(pid: u32) {
72    if !is_alive(pid) {
73        return;
74    }
75    #[cfg(windows)]
76    {
77        let _ = hidden(Command::new("taskkill.exe").args(["/PID", &pid.to_string(), "/T", "/F"]))
78            .stdout(Stdio::null())
79            .stderr(Stdio::null())
80            .status();
81    }
82    #[cfg(unix)]
83    {
84        let mut all = vec![pid];
85        let mut i = 0;
86        while i < all.len() {
87            let kids = children_of(all[i]);
88            all.extend(kids);
89            i += 1;
90        }
91        for sig in ["-TERM", "-KILL"] {
92            for p in all.iter().rev() {
93                let _ = Command::new("kill")
94                    .args([sig, &p.to_string()])
95                    .stdout(Stdio::null())
96                    .stderr(Stdio::null())
97                    .status();
98            }
99            let deadline = std::time::Instant::now() + Duration::from_millis(1500);
100            while std::time::Instant::now() < deadline && all.iter().any(|p| is_alive(*p)) {
101                crate::util::sleep_ms(25);
102            }
103            if !all.iter().any(|p| is_alive(*p)) {
104                break;
105            }
106        }
107    }
108}
109
110/// Kill the process group led by `pid` (the owned-command leader) so descendants that
111/// outlived the leader are cleaned up too. On Windows the Job Object dies with the child handle.
112fn kill_group(pid: u32) {
113    #[cfg(unix)]
114    {
115        let _ = Command::new("kill")
116            .args(["-KILL", "--", &format!("-{pid}")])
117            .stdout(Stdio::null())
118            .stderr(Stdio::null())
119            .status();
120    }
121    #[cfg(windows)]
122    {
123        let _ = hidden(Command::new("taskkill.exe").args(["/PID", &pid.to_string(), "/T", "/F"]))
124            .stdout(Stdio::null())
125            .stderr(Stdio::null())
126            .status();
127    }
128}
129
130/// Registry of every pid a run started, so the run can assert nothing leaked.
131#[derive(Clone, Default)]
132pub struct Tracker {
133    inner: Arc<Mutex<BTreeMap<u32, String>>>,
134    registry: Arc<Mutex<Option<PathBuf>>>,
135}
136
137impl Tracker {
138    pub fn new() -> Self {
139        Self::default()
140    }
141    pub fn register(&self, pid: u32, label: &str) {
142        self.inner.lock().unwrap().insert(pid, label.to_string());
143        self.persist();
144    }
145    pub fn forget(&self, pid: u32) {
146        self.inner.lock().unwrap().remove(&pid);
147        self.persist();
148    }
149    pub fn alive(&self) -> Vec<(u32, String)> {
150        self.inner
151            .lock()
152            .unwrap()
153            .iter()
154            .filter(|(p, _)| is_alive(**p))
155            .map(|(p, l)| (*p, l.clone()))
156            .collect()
157    }
158    pub fn pids(&self) -> Vec<u32> {
159        self.inner.lock().unwrap().keys().copied().collect()
160    }
161    pub fn kill_all(&self) {
162        for (pid, _) in self.alive() {
163            kill_tree(pid);
164        }
165    }
166    /// Fail when any tracked process survived; survivors are killed first so a
167    /// failing assertion never leaves its own litter.
168    pub fn assert_no_orphans(&self) -> Result<()> {
169        let alive = self.alive();
170        if alive.is_empty() {
171            return Ok(());
172        }
173        self.kill_all();
174        let list = alive
175            .iter()
176            .map(|(p, l)| format!("{p} ({l})"))
177            .collect::<Vec<_>>()
178            .join(", ");
179        err(format!("orphan processes still running: {list}"))
180    }
181}
182
183#[derive(Debug, Clone, Default)]
184pub struct RunOptions {
185    pub cwd: Option<PathBuf>,
186    pub env: Vec<(String, String)>,
187    /// When true the inherited environment is dropped first.
188    pub clear_env: bool,
189    pub timeout: Option<Duration>,
190    pub label: String,
191}
192
193#[derive(Debug, Clone)]
194pub struct Output {
195    pub pid: u32,
196    pub code: Option<i32>,
197    pub stdout: String,
198    pub stderr: String,
199    pub timed_out: bool,
200}
201
202/// Run to completion with a hard deadline; on timeout the whole tree is killed.
203pub fn run_owned(
204    program: &str,
205    args: &[String],
206    opts: &RunOptions,
207    tracker: &Tracker,
208) -> Result<Output> {
209    let mut cmd = Command::new(program);
210    cmd.args(args)
211        .stdin(Stdio::null())
212        .stdout(Stdio::piped())
213        .stderr(Stdio::piped());
214    #[cfg(windows)]
215    {
216        use std::os::windows::process::CommandExt;
217        cmd.creation_flags(0x0800_0000); // CREATE_NO_WINDOW: never flash a console from a GUI host
218    }
219    if let Some(cwd) = &opts.cwd {
220        cmd.current_dir(cwd);
221    }
222    if opts.clear_env {
223        cmd.env_clear();
224    }
225    for (k, v) in &opts.env {
226        cmd.env(k, v);
227    }
228    let mut owned = OwnedCommand::from_command(cmd);
229    owned.windows_hide();
230    let mut child = owned
231        .spawn()
232        .map_err(|e| crate::util::Error(format!("failed to start {program}: {e}")))?;
233    let pid = child.id();
234    tracker.register(
235        pid,
236        if opts.label.is_empty() {
237            program
238        } else {
239            &opts.label
240        },
241    );
242    let out_buf: Arc<Mutex<Vec<u8>>> = Arc::default();
243    let err_buf: Arc<Mutex<Vec<u8>>> = Arc::default();
244    let (done_tx, done_rx) = std::sync::mpsc::channel::<()>();
245    for (stream, buf) in [
246        (
247            child
248                .take_stdout()
249                .map(|s| Box::new(s) as Box<dyn Read + Send>),
250            out_buf.clone(),
251        ),
252        (
253            child
254                .take_stderr()
255                .map(|s| Box::new(s) as Box<dyn Read + Send>),
256            err_buf.clone(),
257        ),
258    ] {
259        let tx = done_tx.clone();
260        std::thread::spawn(move || {
261            if let Some(mut s) = stream {
262                let mut chunk = [0u8; 8192];
263                while let Ok(n) = s.read(&mut chunk) {
264                    if n == 0 {
265                        break;
266                    }
267                    buf.lock().unwrap().extend_from_slice(&chunk[..n]);
268                }
269            }
270            let _ = tx.send(());
271        });
272    }
273    drop(done_tx);
274    let timeout = opts.timeout.unwrap_or(Duration::from_secs(120));
275    let status = child.wait_timeout(timeout)?;
276    let (code, timed_out) = match status {
277        Some(s) => {
278            kill_group(pid); // owned descendants never outlive their command
279            (s.code(), false)
280        }
281        None => {
282            let st = child.terminate_tree().ok();
283            kill_group(pid);
284            (st.and_then(|s| s.code()), true)
285        }
286    };
287    // Stream collection is bounded: a descendant that outlives the direct child and
288    // keeps the pipes open must not hang the run. Give readers a short grace, then
289    // kill the child's whole process group and give them a final bounded wait.
290    let mut finished = 0;
291    let grace = Instant::now() + Duration::from_secs(2);
292    while finished < 2 {
293        match done_rx.recv_timeout(grace.saturating_duration_since(Instant::now())) {
294            Ok(()) => finished += 1,
295            Err(_) => break,
296        }
297    }
298    let mut leaked_descendants = false;
299    if finished < 2 {
300        leaked_descendants = true;
301        kill_group(pid);
302        let last = Instant::now() + Duration::from_secs(3);
303        while finished < 2 {
304            match done_rx.recv_timeout(last.saturating_duration_since(Instant::now())) {
305                Ok(()) => finished += 1,
306                Err(_) => break,
307            }
308        }
309    }
310    drop(child);
311    tracker.forget(pid);
312    let mut stderr = String::from_utf8_lossy(&err_buf.lock().unwrap()).into_owned();
313    if leaked_descendants {
314        stderr.push_str("\n[rightkit-qa] descendant processes outlived the command and held its output open; the process group was killed");
315    }
316    let stdout = String::from_utf8_lossy(&out_buf.lock().unwrap()).into_owned();
317    Ok(Output {
318        pid,
319        code,
320        stdout,
321        stderr,
322        timed_out,
323    })
324}
325
326// ---------------------------------------------------------------- crash sweep
327
328/// `ps` command line for a pid (unix); the fingerprint that makes a recorded pid
329/// safe to kill later without risking an unrelated process that reused it.
330pub fn fingerprint(pid: u32) -> Option<String> {
331    #[cfg(unix)]
332    {
333        let out = Command::new("ps")
334            .args(["-o", "command=", "-p", &pid.to_string()])
335            .output()
336            .ok()?;
337        let s = String::from_utf8_lossy(&out.stdout).trim().to_string();
338        if s.is_empty() {
339            None
340        } else {
341            Some(s)
342        }
343    }
344    #[cfg(windows)]
345    {
346        let _ = pid;
347        None
348    }
349}
350
351impl Tracker {
352    /// Persist every registration to `path` as `[{pid, label, command}]`, so a
353    /// crashed run's survivors can be swept by the next run.
354    pub fn with_registry(path: PathBuf) -> Self {
355        let t = Self::default();
356        *t.registry.lock().unwrap() = Some(path);
357        t
358    }
359
360    pub(crate) fn persist(&self) {
361        let Some(path) = self.registry.lock().unwrap().clone() else {
362            return;
363        };
364        let rows: Vec<serde_json::Value> = self
365            .inner
366            .lock()
367            .unwrap()
368            .iter()
369            .map(|(pid, label)| serde_json::json!({"pid": pid, "label": label, "command": fingerprint(*pid)}))
370            .collect();
371        if let Some(p) = path.parent() {
372            let _ = std::fs::create_dir_all(p);
373        }
374        let _ = std::fs::write(path, serde_json::to_vec_pretty(&rows).unwrap_or_default());
375    }
376}
377
378/// Kill processes recorded by earlier runs that are still alive AND still have
379/// the exact command line that was recorded. Returns the pids killed.
380pub fn sweep_registries(runs_root: &std::path::Path, skip_run: &str) -> Vec<u32> {
381    let mut killed = vec![];
382    let Ok(entries) = std::fs::read_dir(runs_root) else {
383        return killed;
384    };
385    for e in entries.flatten() {
386        if e.file_name().to_string_lossy() == skip_run {
387            continue;
388        }
389        let Ok(ft) = e.file_type() else { continue };
390        if !ft.is_dir() || ft.is_symlink() {
391            continue;
392        }
393        let Some(marker) = crate::storage::read_marker(&e.path().join(crate::storage::RUN_MARKER))
394        else {
395            continue;
396        };
397        if marker.run_id != e.file_name().to_string_lossy()
398            || crate::storage::marker_active(&marker, crate::util::now_ms())
399        {
400            continue;
401        }
402        let file = e.path().join("pids.json");
403        let Ok(bytes) = std::fs::read(&file) else {
404            continue;
405        };
406        let rows: Vec<serde_json::Value> = serde_json::from_slice(&bytes).unwrap_or_default();
407        for r in rows {
408            let (Some(pid), Some(cmd)) =
409                (r["pid"].as_u64().map(|p| p as u32), r["command"].as_str())
410            else {
411                continue;
412            };
413            if is_alive(pid) && fingerprint(pid).as_deref() == Some(cmd) {
414                kill_tree(pid);
415                killed.push(pid);
416            }
417        }
418        let _ = std::fs::remove_file(&file);
419    }
420    killed
421}