1use 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
57pub 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
97fn 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#[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 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 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
189pub 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); }
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); (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 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
313pub 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 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
365pub 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}