Skip to main content

differential_engine/
subprocess.rs

1//! One child process, run to completion under a deadline and a cancel flag.
2//!
3//! The adapter that `llm` and `forgeio` share. Both spawn a program that
4//! carries its own credentials, feed it bytes, and want its bytes back before
5//! a deadline or the moment a reviewer gives up — and neither wants to own the
6//! watchdog that makes that safe. This is that watchdog, once.
7//!
8//! Threaded i/o on every pipe: a large stdin must not deadlock against a child
9//! that writes before it finishes reading, and a large stdout must not fill
10//! the pipe while the watchdog waits for exit.
11
12use std::io::{Read, Write};
13use std::path::Path;
14use std::process::{Command, ExitStatus, Stdio};
15use std::sync::Arc;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::time::{Duration, Instant};
18
19/// What to run and how long to wait for it.
20pub struct Run<'a> {
21    pub argv: &'a [String],
22    pub stdin: Option<&'a [u8]>,
23    /// `None` inherits this process's directory.
24    pub working_dir: Option<&'a Path>,
25    pub timeout: Duration,
26    /// Set from another thread to kill the child early.
27    pub cancel: Option<&'a Arc<AtomicBool>>,
28}
29
30/// The child's whole output. The status is not judged here: a non-zero exit
31/// means one thing to a completion and another to `is-ancestor`, and the
32/// caller knows which it asked for.
33pub struct Output {
34    pub status: ExitStatus,
35    pub stdout: Vec<u8>,
36    pub stderr: Vec<u8>,
37}
38
39/// Why there is no output. The command that failed is the caller's to name,
40/// so it is not repeated in here.
41#[derive(Debug)]
42pub enum Failure {
43    Spawn(std::io::Error),
44    Io(std::io::Error),
45    Timeout,
46    Cancelled,
47}
48
49pub fn run(spec: &Run<'_>) -> Result<Output, Failure> {
50    let mut cmd = Command::new(&spec.argv[0]);
51    if let Some(dir) = spec.working_dir {
52        cmd.current_dir(dir);
53    }
54    let mut child = cmd
55        .args(&spec.argv[1..])
56        .stdin(Stdio::piped())
57        .stdout(Stdio::piped())
58        .stderr(Stdio::piped())
59        .spawn()
60        .map_err(Failure::Spawn)?;
61
62    let mut stdin = child.stdin.take().expect("stdin piped");
63    let input = spec.stdin.map(<[u8]>::to_vec);
64    let writer = std::thread::spawn(move || {
65        if let Some(bytes) = input {
66            let _ = stdin.write_all(&bytes);
67        }
68        // stdin closes on drop, so a child that reads to EOF sees it.
69    });
70    let mut out_pipe = child.stdout.take().expect("stdout piped");
71    let stdout_reader = std::thread::spawn(move || {
72        let mut buf = Vec::new();
73        out_pipe.read_to_end(&mut buf).map(|_| buf)
74    });
75    let mut err_pipe = child.stderr.take().expect("stderr piped");
76    let stderr_reader = std::thread::spawn(move || {
77        let mut buf = Vec::new();
78        let _ = err_pipe.read_to_end(&mut buf);
79        buf
80    });
81
82    // Watchdog: poll for exit until the deadline, then kill. The poll is not
83    // just the deadline's — the cancel flag has to be read too, which is why
84    // this is a loop and not a `wait` with a timeout.
85    let deadline = Instant::now() + spec.timeout;
86    let cancelled = || spec.cancel.is_some_and(|c| c.load(Ordering::Relaxed));
87    let status = loop {
88        // Decide first, tear down once — a `try_wait` error included, or
89        // the child and its three threads would outlive this call.
90        let give_up = match child.try_wait() {
91            Ok(Some(status)) => break status,
92            Err(e) => Some(Failure::Io(e)),
93            Ok(None) if cancelled() => Some(Failure::Cancelled),
94            Ok(None) if Instant::now() >= deadline => Some(Failure::Timeout),
95            Ok(None) => None,
96        };
97        if let Some(err) = give_up {
98            let _ = child.kill();
99            let _ = child.wait();
100            let _ = writer.join();
101            let _ = stdout_reader.join();
102            let _ = stderr_reader.join();
103            return Err(err);
104        }
105        std::thread::sleep(Duration::from_millis(25));
106    };
107    let _ = writer.join();
108    let stdout = stdout_reader
109        .join()
110        .expect("stdout reader panicked")
111        .map_err(Failure::Io)?;
112    let stderr = stderr_reader.join().expect("stderr reader panicked");
113    Ok(Output {
114        status,
115        stdout,
116        stderr,
117    })
118}
119
120/// The first `limit` bytes of stderr as text, for an error message.
121pub fn stderr_excerpt(stderr: &[u8], limit: usize) -> String {
122    String::from_utf8_lossy(&stderr[..stderr.len().min(limit)])
123        .trim()
124        .to_string()
125}