use std::io::Read;
use std::process::{Child, Output};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
pub struct LimitedProcessOutput {
pub output: Output,
pub completeness: crate::core::context_admission::OutputCompleteness,
}
#[cfg(test)]
pub(crate) fn wait_with_limits(
child: Child,
max_bytes: usize,
timeout: std::time::Duration,
kill_group: bool,
) -> Output {
wait_with_limits_captured(child, max_bytes, timeout, kill_group).output
}
pub fn wait_with_limits_captured(
mut child: Child,
max_bytes: usize,
timeout: std::time::Duration,
kill_group: bool,
) -> LimitedProcessOutput {
const STDERR_LIMIT: usize = 512 * 1024;
const DRAIN_GRACE: std::time::Duration = std::time::Duration::from_millis(250);
const READER_BUDGET: std::time::Duration = std::time::Duration::from_secs(2);
const REAP_BUDGET: std::time::Duration = std::time::Duration::from_secs(2);
let own_group = owned_group(&mut child, kill_group);
let stdout_pipe = child.stdout.take();
let stderr_pipe = child.stderr.take();
let start = std::time::Instant::now();
let truncated = Arc::new(AtomicBool::new(false));
let read_failed = Arc::new(AtomicBool::new(false));
let stdout_capture = Arc::new(StreamCapture::default());
let stderr_capture = Arc::new(StreamCapture::default());
let stdout_handle = std::thread::spawn({
let (capture, truncated, read_failed) = (
Arc::clone(&stdout_capture),
Arc::clone(&truncated),
Arc::clone(&read_failed),
);
move || read_stream(stdout_pipe, max_bytes, &capture, &truncated, &read_failed)
});
let stderr_handle = std::thread::spawn({
let (capture, truncated, read_failed) = (
Arc::clone(&stderr_capture),
Arc::clone(&truncated),
Arc::clone(&read_failed),
);
move || {
read_stream(
stderr_pipe,
STDERR_LIMIT,
&capture,
&truncated,
&read_failed,
);
}
});
let mut timed_out = false;
let mut wait_failed = false;
let mut kill_failed = false;
let mut forced_pipe_cleanup = false;
loop {
let hit_timeout = start.elapsed() > timeout;
if hit_timeout || truncated.load(Ordering::Relaxed) {
kill_failed |= !kill_child(&mut child, own_group);
timed_out = hit_timeout;
break;
}
match exited_unreaped(&mut child) {
Ok(true) => {
if own_group.is_some()
&& !readers_finish_within(&stdout_handle, &stderr_handle, DRAIN_GRACE)
{
kill_failed |= !kill_child(&mut child, own_group);
forced_pipe_cleanup = true;
}
break;
}
Err(_) => {
wait_failed = true;
kill_failed |= !kill_child(&mut child, own_group);
break;
}
Ok(false) => std::thread::sleep(std::time::Duration::from_millis(50)),
}
}
let reaped = reap_within(&mut child, REAP_BUDGET);
let readers_finished = readers_finish_within(&stdout_handle, &stderr_handle, READER_BUDGET);
let reader_join_failed =
!readers_finished || stdout_handle.join().is_err() || stderr_handle.join().is_err();
let (mut stdout_buf, stdout_truncated) = stdout_capture.take();
let (mut stderr_buf, stderr_truncated) = stderr_capture.take();
let decode_lossy =
std::str::from_utf8(&stdout_buf).is_err() || std::str::from_utf8(&stderr_buf).is_err();
let reader_incomplete = wait_failed
|| kill_failed
|| reaped.is_none()
|| forced_pipe_cleanup
|| read_failed.load(Ordering::Relaxed)
|| reader_join_failed;
if timed_out || stdout_truncated {
let notice = format!(
"\n[lean-ctx: output truncated at {} MB / {}s limit]\n",
max_bytes / (1024 * 1024),
timeout.as_secs()
);
stdout_buf.extend_from_slice(notice.as_bytes());
}
if stderr_truncated {
let notice = format!(
"\n[lean-ctx: stderr truncated at {} KB limit]\n",
STDERR_LIMIT / 1024
);
stderr_buf.extend_from_slice(notice.as_bytes());
}
if !readers_finished {
stdout_buf.extend_from_slice(
b"\n[lean-ctx: output not captured: a process outside this call kept it open]\n",
);
}
if reaped.is_none() {
stdout_buf.extend_from_slice(
b"\n[lean-ctx: the process could not be ended within the cleanup budget]\n",
);
}
let status = reaped.unwrap_or_else(synthetic_failure_status);
LimitedProcessOutput {
output: Output {
status,
stdout: stdout_buf,
stderr: stderr_buf,
},
completeness: crate::core::context_admission::OutputCompleteness {
truncated: stdout_truncated || stderr_truncated || truncated.load(Ordering::Relaxed),
timed_out,
reader_incomplete,
cancelled: false,
decode_lossy,
},
}
}
#[derive(Default)]
struct StreamCapture {
bytes: std::sync::Mutex<Vec<u8>>,
truncated: AtomicBool,
}
impl StreamCapture {
fn take(&self) -> (Vec<u8>, bool) {
let bytes = std::mem::take(
&mut *self
.bytes
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
);
(bytes, self.truncated.load(Ordering::Relaxed))
}
}
fn read_stream(
pipe: Option<impl Read>,
limit: usize,
capture: &StreamCapture,
any_truncated: &AtomicBool,
read_failed: &AtomicBool,
) {
let Some(mut pipe) = pipe else {
read_failed.store(true, Ordering::Relaxed);
return;
};
let mut chunk = [0u8; 8192];
loop {
match pipe.read(&mut chunk) {
Ok(0) => return,
Ok(n) => {
let mut bytes = capture
.bytes
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if bytes.len() + n > limit {
let remaining = limit.saturating_sub(bytes.len());
bytes.extend_from_slice(&chunk[..remaining]);
capture.truncated.store(true, Ordering::Relaxed);
any_truncated.store(true, Ordering::Relaxed);
return;
}
bytes.extend_from_slice(&chunk[..n]);
}
Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
Err(_) => {
read_failed.store(true, Ordering::Relaxed);
return;
}
}
}
}
#[cfg(unix)]
fn owned_group(child: &mut Child, kill_group: bool) -> Option<i32> {
if !kill_group {
return None;
}
let pid = child.id() as libc::pid_t;
let group = unsafe { libc::getpgid(pid) };
if group == pid {
return Some(pid);
}
(group == -1 && matches!(exited_unreaped(child), Ok(true))).then_some(pid)
}
#[cfg(not(unix))]
fn owned_group(_: &mut Child, _: bool) -> Option<i32> {
None
}
#[cfg(unix)]
pub(crate) fn exited_unreaped(child: &mut Child) -> std::io::Result<bool> {
let mut info: libc::siginfo_t = unsafe { std::mem::zeroed() };
let result = unsafe {
libc::waitid(
libc::P_PID,
child.id() as libc::id_t,
&raw mut info,
libc::WEXITED | libc::WNOHANG | libc::WNOWAIT,
)
};
if result == -1 {
let error = std::io::Error::last_os_error();
return if error.kind() == std::io::ErrorKind::Interrupted {
Ok(false)
} else {
Err(error)
};
}
Ok(unsafe { info.si_pid() } != 0)
}
#[cfg(not(unix))]
fn exited_unreaped(child: &mut Child) -> std::io::Result<bool> {
child.try_wait().map(|status| status.is_some())
}
fn readers_finish_within<A, B>(
first: &std::thread::JoinHandle<A>,
second: &std::thread::JoinHandle<B>,
budget: std::time::Duration,
) -> bool {
let deadline = std::time::Instant::now() + budget;
loop {
if first.is_finished() && second.is_finished() {
return true;
}
if std::time::Instant::now() >= deadline {
return false;
}
std::thread::sleep(std::time::Duration::from_millis(5));
}
}
fn kill_child(child: &mut Child, own_group: Option<i32>) -> bool {
#[cfg(unix)]
if let Some(group) = own_group {
unsafe { libc::killpg(group, libc::SIGKILL) };
}
#[cfg(not(unix))]
let _ = own_group;
#[cfg(windows)]
let tree_killed = {
const TREE_KILL_BUDGET: std::time::Duration = std::time::Duration::from_secs(5);
match std::process::Command::new("taskkill")
.args(["/F", "/T", "/PID", &child.id().to_string()])
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
{
Ok(mut tree_kill) => {
if let Some(status) = reap_within(&mut tree_kill, TREE_KILL_BUDGET) {
status.success()
} else {
let _ = tree_kill.kill();
let _ = reap_within(&mut tree_kill, TREE_KILL_BUDGET);
false
}
}
Err(_) => false,
}
};
#[cfg(not(windows))]
let tree_killed = true;
let _ = child.kill();
tree_killed
}
pub fn reap_within(
child: &mut Child,
budget: std::time::Duration,
) -> Option<std::process::ExitStatus> {
let deadline = std::time::Instant::now() + budget;
loop {
match child.try_wait() {
Ok(Some(status)) => return Some(status),
Ok(None) if std::time::Instant::now() < deadline => {
std::thread::sleep(std::time::Duration::from_millis(5));
}
Ok(None) | Err(_) => return None,
}
}
}
#[cfg(unix)]
fn synthetic_failure_status() -> std::process::ExitStatus {
use std::os::unix::process::ExitStatusExt;
std::process::ExitStatus::from_raw(1 << 8)
}
#[cfg(not(unix))]
fn synthetic_failure_status() -> std::process::ExitStatus {
use std::os::windows::process::ExitStatusExt;
std::process::ExitStatus::from_raw(1)
}
#[cfg(test)]
mod tests {
const OUTPUT_HELPER_ENV: &str = "LEAN_CTX_TEST_OUTPUT_HELPER";
const OUTPUT_HELPER_TEST: &str = "shell::exec::process::tests::wait_with_limits_output_helper";
#[test]
#[ignore = "subprocess helper for wait_with_limits tests"]
fn wait_with_limits_output_helper() {
use std::io::Write as _;
match std::env::var(OUTPUT_HELPER_ENV).as_deref() {
Ok("stdout-small") => {
std::io::stdout().write_all(b"hello\n").unwrap();
}
Ok("stdout-forever") => {
let chunk = [b'a'; 8192];
let mut stream = std::io::stdout().lock();
while stream.write_all(&chunk).is_ok() {}
std::thread::sleep(std::time::Duration::from_mins(1));
}
Ok("stderr-forever") => {
let chunk = [b'a'; 8192];
let mut stream = std::io::stderr().lock();
while stream.write_all(&chunk).is_ok() {}
std::thread::sleep(std::time::Duration::from_mins(1));
}
Ok("sleep") => {
std::thread::sleep(std::time::Duration::from_mins(1));
}
#[cfg(unix)]
Ok("guardian") => guardian(),
Ok(other) => panic!("unknown output helper mode: {other}"),
Err(error) => panic!("missing {OUTPUT_HELPER_ENV}: {error}"),
}
}
fn spawn_output_helper(mode: &str) -> std::process::Child {
std::process::Command::new(std::env::current_exe().unwrap())
.args([OUTPUT_HELPER_TEST, "--exact", "--ignored", "--nocapture"])
.env(OUTPUT_HELPER_ENV, mode)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.spawn()
.unwrap()
}
#[test]
fn wait_with_limits_captures_output() {
let child = spawn_output_helper("stdout-small");
let output = super::wait_with_limits(child, 1024, std::time::Duration::from_secs(5), false);
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
stdout.contains("hello"),
"expected 'hello' in output: {stdout}"
);
assert!(output.status.success());
}
#[test]
fn wait_with_limits_truncates_large_output() {
let child = spawn_output_helper("stdout-forever");
let output =
super::wait_with_limits(child, 1024, std::time::Duration::from_secs(10), false);
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
stdout.contains("[lean-ctx: output truncated"),
"expected truncation notice, got len={}: ...{}",
stdout.len(),
&stdout[stdout.len().saturating_sub(80)..]
);
}
#[test]
fn synthetic_failure_status_is_a_failure_without_spawning_anything() {
let status = super::synthetic_failure_status();
assert!(!status.success());
#[cfg(unix)]
{
use std::os::unix::process::ExitStatusExt;
assert_eq!(status.code(), Some(1));
assert_eq!(status.signal(), None);
}
}
#[test]
fn wait_with_limits_truncates_large_stderr() {
let child = spawn_output_helper("stderr-forever");
let output = super::wait_with_limits(
child,
1024 * 1024,
std::time::Duration::from_secs(10),
false,
);
let stderr = String::from_utf8_lossy(&output.stderr);
assert!(
stderr.contains("[lean-ctx: stderr truncated"),
"expected stderr truncation notice, got len={}: ...{}",
stderr.len(),
&stderr[stderr.len().saturating_sub(80)..]
);
}
#[test]
fn wait_with_limits_kills_promptly_on_truncation() {
let child = spawn_output_helper("stdout-forever");
let start = std::time::Instant::now();
let output =
super::wait_with_limits(child, 4096, std::time::Duration::from_secs(20), false);
let elapsed = start.elapsed();
assert!(
elapsed < std::time::Duration::from_secs(3),
"truncation should kill promptly, took {elapsed:?} (timeout was 20s)"
);
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(stdout.contains("[lean-ctx: output truncated"));
}
#[test]
fn wait_with_limits_timeout_kills_process() {
let child = spawn_output_helper("sleep");
let start = std::time::Instant::now();
let output =
super::wait_with_limits(child, 1024, std::time::Duration::from_millis(200), false);
let elapsed = start.elapsed();
assert!(
elapsed < std::time::Duration::from_secs(3),
"timeout should kill quickly, took {elapsed:?}"
);
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(stdout.contains("[lean-ctx: output truncated"));
}
#[cfg(unix)]
#[test]
fn wait_with_limits_group_kill_reaps_grandchildren() {
use std::os::unix::process::CommandExt as _;
let mut cmd = std::process::Command::new("sh");
cmd.args(["-c", "sleep 30 & sleep 30"])
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
cmd.process_group(0);
let child = cmd.spawn().unwrap();
let pgid = child.id() as libc::pid_t;
let start = std::time::Instant::now();
let _ = super::wait_with_limits(child, 1024, std::time::Duration::from_millis(200), true);
let elapsed = start.elapsed();
assert!(
elapsed < std::time::Duration::from_secs(5),
"group kill must unblock the reader threads, took {elapsed:?}"
);
let mut group_gone = false;
for _ in 0..50 {
if unsafe { libc::killpg(pgid, 0) } == -1 {
group_gone = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
assert!(group_gone, "process group {pgid} must be fully reaped");
}
#[cfg(unix)]
fn alive(pid: u32) -> bool {
unsafe { libc::kill(pid as libc::pid_t, 0) == 0 }
}
#[cfg(unix)]
fn guardian() {
use std::os::unix::process::CommandExt as _;
use std::process::{Command, Stdio};
let piped = |script: &str| {
let mut command = Command::new("sh");
command
.args(["-c", script])
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
command
};
let mut sentinel = Command::new("sleep")
.arg("60")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.unwrap();
let own = unsafe { libc::getpgrp() };
assert_eq!(own, std::process::id() as libc::pid_t, "guardian leads");
struct EndOwnGroup(libc::pid_t);
impl Drop for EndOwnGroup {
fn drop(&mut self) {
use std::io::Write as _;
let _ = std::io::stdout().flush();
unsafe { libc::killpg(self.0, libc::SIGKILL) };
}
}
let _end = EndOwnGroup(own);
assert_eq!(unsafe { libc::getpgid(sentinel.id() as libc::pid_t) }, own);
let child = piped("sleep 60 & sleep 60").spawn().unwrap();
let start = std::time::Instant::now();
let captured = super::wait_with_limits_captured(
child,
1024,
std::time::Duration::from_millis(200),
true,
);
assert!(start.elapsed() < std::time::Duration::from_secs(8));
assert!(captured.completeness.timed_out);
assert!(captured.completeness.reader_incomplete);
assert!(
alive(sentinel.id()),
"inherited group must not be signalled"
);
let child = piped("sleep 60 & yes").spawn().unwrap();
let captured =
super::wait_with_limits_captured(child, 1024, std::time::Duration::from_secs(30), true);
assert!(captured.completeness.truncated);
assert!(
alive(sentinel.id()),
"inherited group must not be signalled"
);
let mut command = piped("sleep 60 & echo started");
command.process_group(0);
let child = command.spawn().unwrap();
let group = child.id() as libc::pid_t;
let start = std::time::Instant::now();
let captured =
super::wait_with_limits_captured(child, 1024, std::time::Duration::from_secs(30), true);
assert!(
start.elapsed() < std::time::Duration::from_secs(5),
"leader gone, pipe holder alive: must not wait for it"
);
assert!(!captured.completeness.timed_out);
assert!(captured.completeness.reader_incomplete);
assert!(String::from_utf8_lossy(&captured.output.stdout).contains("started"));
let mut group_gone = false;
for _ in 0..100 {
if unsafe { libc::killpg(group, 0) } == -1 {
group_gone = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
assert!(group_gone, "the call's own group must be gone");
assert!(
alive(sentinel.id()),
"a foreign group must not be signalled"
);
let child = piped("echo written; echo problem >&2; sleep 60 &")
.spawn()
.unwrap();
let start = std::time::Instant::now();
let captured =
super::wait_with_limits_captured(child, 1024, std::time::Duration::from_secs(30), true);
assert!(start.elapsed() < std::time::Duration::from_secs(8));
assert!(captured.output.status.success());
assert!(captured.completeness.reader_incomplete);
assert!(!captured.completeness.timed_out && !captured.completeness.truncated);
assert_eq!(
String::from_utf8_lossy(&captured.output.stdout),
"written\n\n[lean-ctx: output not captured: a process outside this call kept it open]\n"
);
assert_eq!(captured.output.stderr, b"problem\n");
assert!(
alive(sentinel.id()),
"inherited group must not be signalled"
);
let child = piped("echo written; echo problem >&2").spawn().unwrap();
let captured =
super::wait_with_limits_captured(child, 1024, std::time::Duration::from_secs(30), true);
assert_eq!(captured.output.stdout, b"written\n");
assert_eq!(captured.output.stderr, b"problem\n");
assert!(captured.completeness.is_complete());
sentinel.kill().unwrap();
sentinel.wait().unwrap();
println!("GUARDIAN-COMPLETE");
}
#[cfg(unix)]
#[test]
fn limits_signal_only_a_group_the_call_owns() {
use std::os::unix::process::CommandExt as _;
let mut command = std::process::Command::new(std::env::current_exe().unwrap());
command
.args([OUTPUT_HELPER_TEST, "--exact", "--ignored", "--nocapture"])
.env(OUTPUT_HELPER_ENV, "guardian")
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
command.process_group(0);
let child = command.spawn().unwrap();
let captured = super::wait_with_limits_captured(
child,
64 * 1024,
std::time::Duration::from_mins(1),
true,
);
let stdout = String::from_utf8_lossy(&captured.output.stdout);
let stderr = String::from_utf8_lossy(&captured.output.stderr);
assert!(
stdout.contains("GUARDIAN-COMPLETE"),
"guardian status {:?}\n{stdout}\n{stderr}",
captured.output.status
);
}
#[cfg(unix)]
#[test]
fn stdin_relay_forwards_piped_data() {
use std::io::Write;
use std::os::unix::process::CommandExt as _;
let mut cmd = std::process::Command::new("cat");
cmd.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
cmd.process_group(0);
let mut child = cmd.spawn().expect("failed to spawn cat");
let mut child_stdin = child.stdin.take().unwrap();
std::thread::spawn(move || {
child_stdin.write_all(b"hello from pipe\n").unwrap();
drop(child_stdin);
});
let output = child.wait_with_output().expect("wait failed");
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
stdout.contains("hello from pipe"),
"#806: piped stdin must reach the child, got: {stdout}"
);
}
#[cfg(unix)]
#[test]
fn stdin_relay_no_data_does_not_hang() {
let start = std::time::Instant::now();
let mut cmd = std::process::Command::new("sh");
cmd.args(["-c", "echo ok"])
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
use std::os::unix::process::CommandExt as _;
cmd.process_group(0);
let mut child = cmd.spawn().unwrap();
drop(child.stdin.take());
let output = child.wait_with_output().unwrap();
let elapsed = start.elapsed();
assert!(
elapsed < std::time::Duration::from_secs(5),
"must not hang when stdin is closed immediately, took {elapsed:?}"
);
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(stdout.contains("ok"), "command output missing: {stdout}");
}
}