Skip to main content

rightkit_qa/
process.rs

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