use std::collections::{BTreeMap, BTreeSet};
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use anyhow::{Context, Result};
use tokio::io::AsyncReadExt;
use tokio::sync::{mpsc, watch};
use tokio::task::{JoinHandle, JoinSet};
use crate::hel_acp::RuntimeEvent;
pub const DEFAULT_TERMINAL_OUTPUT_BYTES: usize = 4 * 1024 * 1024;
const TERMINAL_BUFFER_SLACK_BYTES: usize = 64 * 1024;
const TERMINAL_READ_CHUNK_BYTES: usize = 16 * 1024;
const TERMINAL_REAP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
#[derive(Debug, Clone)]
pub struct TerminalSpawn {
pub command: String,
pub args: Vec<String>,
pub env: Vec<(String, String)>,
pub cwd: PathBuf,
pub output_byte_limit: usize,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TerminalExit {
pub exit_code: Option<u32>,
pub signal: Option<String>,
}
#[derive(Debug, Clone)]
pub struct TerminalSnapshot {
pub output: String,
pub truncated: bool,
pub exit: Option<TerminalExit>,
}
#[derive(Debug)]
pub struct TerminalBuffer {
bytes: Vec<u8>,
limit: usize,
dropped: bool,
}
impl TerminalBuffer {
#[must_use]
pub fn new(limit: usize) -> Self {
Self {
bytes: Vec::new(),
limit: limit.max(1),
dropped: false,
}
}
pub fn append(&mut self, chunk: &[u8]) {
self.bytes.extend_from_slice(chunk);
if self.bytes.len() > self.limit.saturating_add(TERMINAL_BUFFER_SLACK_BYTES) {
let excess = self.bytes.len() - self.limit;
self.bytes.drain(..excess);
self.dropped = true;
}
}
#[must_use]
pub fn read(&self) -> (String, bool) {
let mut start = self.bytes.len().saturating_sub(self.limit);
let truncated = self.dropped || start > 0;
if truncated {
while start < self.bytes.len() && self.bytes[start] & 0b1100_0000 == 0b1000_0000 {
start += 1;
}
}
(
String::from_utf8_lossy(&self.bytes[start..]).into_owned(),
truncated,
)
}
#[must_use]
pub fn buffered_len(&self) -> usize {
self.bytes.len()
}
}
struct ProcessGroup {
pid: i32,
exit: watch::Receiver<Option<TerminalExit>>,
}
impl ProcessGroup {
fn kill_if_live(&self) {
if self.exit.borrow().is_none() {
kill_process_group(self.pid);
}
}
}
impl Drop for ProcessGroup {
fn drop(&mut self) {
self.kill_if_live();
}
}
struct TerminalEntry {
group: ProcessGroup,
buffer: Arc<Mutex<TerminalBuffer>>,
exit: watch::Receiver<Option<TerminalExit>>,
supervisor: JoinHandle<()>,
}
#[derive(Clone, Default)]
pub struct TerminalRegistry {
terminals: Arc<Mutex<BTreeMap<String, TerminalEntry>>>,
next_id: Arc<AtomicU64>,
}
impl TerminalRegistry {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn create(
&self,
spawn: TerminalSpawn,
events: mpsc::Sender<RuntimeEvent>,
) -> Result<String> {
let line = shell_line(&spawn.command, &spawn.args);
let mut child = spawn_shell(&line, &spawn)?;
let pid = child
.id()
.and_then(|pid| i32::try_from(pid).ok())
.context("client terminal has no usable process ID")?;
let stdout = child
.stdout
.take()
.context("client terminal stdout unavailable")?;
let stderr = child
.stderr
.take()
.context("client terminal stderr unavailable")?;
let buffer = Arc::new(Mutex::new(TerminalBuffer::new(spawn.output_byte_limit)));
let (exit_tx, exit_rx) = watch::channel(None);
let terminal_id = format!("term-{}", self.next_id.fetch_add(1, Ordering::Relaxed) + 1);
let supervisor = tokio::spawn(supervise(
terminal_id.clone(),
child,
stdout,
stderr,
buffer.clone(),
exit_tx,
events,
));
self.terminals
.lock()
.expect("terminal registry lock poisoned")
.insert(
terminal_id.clone(),
TerminalEntry {
group: ProcessGroup {
pid,
exit: exit_rx.clone(),
},
buffer,
exit: exit_rx,
supervisor,
},
);
Ok(terminal_id)
}
#[must_use]
pub fn output(&self, terminal_id: &str) -> Option<TerminalSnapshot> {
let terminals = self
.terminals
.lock()
.expect("terminal registry lock poisoned");
let entry = terminals.get(terminal_id)?;
let (output, truncated) = entry
.buffer
.lock()
.expect("terminal buffer lock poisoned")
.read();
Some(TerminalSnapshot {
output,
truncated,
exit: entry.exit.borrow().clone(),
})
}
#[must_use]
pub fn exit_receiver(
&self,
terminal_id: &str,
) -> Option<watch::Receiver<Option<TerminalExit>>> {
Some(
self.terminals
.lock()
.expect("terminal registry lock poisoned")
.get(terminal_id)?
.exit
.clone(),
)
}
pub fn kill(&self, terminal_id: &str) -> bool {
let terminals = self
.terminals
.lock()
.expect("terminal registry lock poisoned");
let Some(entry) = terminals.get(terminal_id) else {
return false;
};
entry.group.kill_if_live();
true
}
pub fn kill_live(&self) {
let terminals = self
.terminals
.lock()
.expect("terminal registry lock poisoned");
for entry in terminals.values() {
entry.group.kill_if_live();
}
}
pub fn release(&self, terminal_id: &str) -> Option<JoinHandle<()>> {
let entry = self
.terminals
.lock()
.expect("terminal registry lock poisoned")
.remove(terminal_id)?;
Some(entry.supervisor)
}
pub async fn shutdown(&self, events: &mpsc::Sender<RuntimeEvent>) {
let entries = std::mem::take(
&mut *self
.terminals
.lock()
.expect("terminal registry lock poisoned"),
);
let mut unreaped = BTreeSet::new();
let mut supervisors = JoinSet::new();
for (terminal_id, entry) in entries {
entry.group.kill_if_live();
let TerminalEntry { supervisor, .. } = entry;
unreaped.insert(terminal_id.clone());
supervisors.spawn(async move { (terminal_id, supervisor.await) });
}
let deadline = tokio::time::Instant::now() + TERMINAL_REAP_TIMEOUT;
let mut failures = Vec::new();
loop {
match tokio::time::timeout_at(deadline, supervisors.join_next()).await {
Ok(Some(Ok((terminal_id, reaped)))) => {
if let Err(error) = reaped {
failures.push(format!(
"client terminal {terminal_id} supervisor failed: {error}"
));
}
unreaped.remove(&terminal_id);
}
Ok(Some(Err(_))) => {}
Ok(None) => break,
Err(_) => break,
}
}
for terminal_id in unreaped {
failures.push(format!(
"client terminal {terminal_id} did not finish within \
{TERMINAL_REAP_TIMEOUT:?} of being killed"
));
}
for message in failures {
report(events, message).await;
}
}
}
impl TerminalSpawn {
#[must_use]
pub fn display_command(&self) -> String {
let interpreter = PathBuf::from(&self.command)
.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| matches!(name, "sh" | "bash" | "dash" | "zsh"));
if interpreter
&& let Some(index) = self.args.iter().position(|arg| arg == "-c")
&& let Some(script) = self.args.get(index + 1)
{
return script.clone();
}
shell_line(&self.command, &self.args)
}
}
pub async fn wait_for_exit(mut exit: watch::Receiver<Option<TerminalExit>>) -> TerminalExit {
loop {
if let Some(exit) = exit.borrow_and_update().clone() {
return exit;
}
if exit.changed().await.is_err() {
return TerminalExit::default();
}
}
}
#[must_use]
pub fn shell_line(command: &str, args: &[String]) -> String {
let mut line = command.to_owned();
for arg in args {
line.push(' ');
line.push_str(&crate::hel_targets::posix_quote(arg));
}
line
}
#[cfg(unix)]
fn spawn_shell(line: &str, spawn: &TerminalSpawn) -> Result<tokio::process::Child> {
let mut command = tokio::process::Command::new("sh");
command
.arg("-c")
.arg(line)
.current_dir(&spawn.cwd)
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
for (name, value) in &spawn.env {
command.env(name, value);
}
command.process_group(0);
command
.spawn()
.with_context(|| format!("spawn client terminal: {line}"))
}
#[cfg(not(unix))]
fn spawn_shell(_line: &str, _spawn: &TerminalSpawn) -> Result<tokio::process::Child> {
anyhow::bail!("client terminals need Unix process groups")
}
#[cfg(unix)]
fn kill_process_group(pid: i32) {
crate::hel_subprocess::terminate_process_group(pid, libc::SIGKILL);
}
#[cfg(not(unix))]
fn kill_process_group(_pid: i32) {}
#[cfg(unix)]
fn exit_from_status(status: std::process::ExitStatus) -> TerminalExit {
use std::os::unix::process::ExitStatusExt;
TerminalExit {
exit_code: status.code().and_then(|code| u32::try_from(code).ok()),
signal: status.signal().map(signal_name),
}
}
#[cfg(not(unix))]
fn exit_from_status(status: std::process::ExitStatus) -> TerminalExit {
TerminalExit {
exit_code: status.code().and_then(|code| u32::try_from(code).ok()),
signal: None,
}
}
#[cfg(unix)]
fn signal_name(signal: i32) -> String {
let name = match signal {
libc::SIGHUP => "SIGHUP",
libc::SIGINT => "SIGINT",
libc::SIGQUIT => "SIGQUIT",
libc::SIGILL => "SIGILL",
libc::SIGABRT => "SIGABRT",
libc::SIGBUS => "SIGBUS",
libc::SIGFPE => "SIGFPE",
libc::SIGKILL => "SIGKILL",
libc::SIGUSR1 => "SIGUSR1",
libc::SIGSEGV => "SIGSEGV",
libc::SIGUSR2 => "SIGUSR2",
libc::SIGPIPE => "SIGPIPE",
libc::SIGALRM => "SIGALRM",
libc::SIGTERM => "SIGTERM",
libc::SIGSTOP => "SIGSTOP",
libc::SIGTSTP => "SIGTSTP",
libc::SIGCONT => "SIGCONT",
_ => return format!("SIG{signal}"),
};
name.to_owned()
}
async fn supervise(
terminal_id: String,
mut child: tokio::process::Child,
stdout: tokio::process::ChildStdout,
stderr: tokio::process::ChildStderr,
buffer: Arc<Mutex<TerminalBuffer>>,
exit: watch::Sender<Option<TerminalExit>>,
events: mpsc::Sender<RuntimeEvent>,
) {
let (from_stdout, from_stderr, waited) = tokio::join!(
drain_pipe(stdout, buffer.clone()),
drain_pipe(stderr, buffer.clone()),
child.wait(),
);
let (status, failed_reap) = match waited {
Ok(status) => (exit_from_status(status), None),
Err(error) => (
TerminalExit::default(),
Some(format!("reap client terminal {terminal_id}: {error}")),
),
};
if exit.send(Some(status.clone())).is_err() {
tracing::debug!(
%terminal_id,
operation = "terminal_exit",
"terminal exit watcher was closed before publication"
);
}
if let Some(message) = failed_reap {
report(&events, message).await;
}
for (stream, result) in [("stdout", from_stdout), ("stderr", from_stderr)] {
if let Err(error) = result {
report(
&events,
format!("read client terminal {terminal_id} {stream}: {error:#}"),
)
.await;
}
}
let (output, truncated) = buffer.lock().expect("terminal buffer lock poisoned").read();
report_event(
&events,
RuntimeEvent::TerminalClosed {
terminal_id,
output,
truncated,
exit_code: status.exit_code,
signal: status.signal,
},
)
.await;
}
async fn drain_pipe<R>(mut pipe: R, buffer: Arc<Mutex<TerminalBuffer>>) -> Result<()>
where
R: tokio::io::AsyncRead + Unpin,
{
let mut chunk = vec![0_u8; TERMINAL_READ_CHUNK_BYTES];
loop {
match pipe.read(&mut chunk).await {
Ok(0) => return Ok(()),
Ok(read) => buffer
.lock()
.expect("terminal buffer lock poisoned")
.append(&chunk[..read]),
Err(error) => return Err(error).context("read client terminal pipe"),
}
}
}
async fn report(events: &mpsc::Sender<RuntimeEvent>, message: String) {
report_event(events, RuntimeEvent::Warning { message }).await;
}
async fn report_event(events: &mpsc::Sender<RuntimeEvent>, event: RuntimeEvent) {
if let Err(error) = events.send(event).await {
tracing::debug!(
operation = "terminal_event",
%error,
"terminal event could not reach the relay coordinator"
);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_buffer_keeps_the_tail_and_latches_truncated() {
let mut buffer = TerminalBuffer::new(8);
buffer.append(b"abcd");
assert_eq!(buffer.read(), ("abcd".to_owned(), false));
buffer.append(b"efghijkl");
let (output, truncated) = buffer.read();
assert_eq!(output, "efghijkl", "the buffer must serve the last bytes");
assert!(truncated);
let mut latched = TerminalBuffer::new(4);
latched.append(&vec![b'x'; 4 + TERMINAL_BUFFER_SLACK_BYTES + 1]);
assert_eq!(latched.buffered_len(), 4);
latched.append(b"");
assert!(latched.read().1);
}
#[test]
fn the_buffer_serves_output_starting_on_a_character_boundary() {
let mut buffer = TerminalBuffer::new(2);
buffer.append("aéb".as_bytes());
let (output, truncated) = buffer.read();
assert!(truncated);
assert_eq!(output, "b", "a partial character must not be served");
assert!(output.is_char_boundary(0));
}
#[test]
fn the_buffer_stays_bounded_while_a_child_floods_it() {
let limit = 4 * 1024;
let mut buffer = TerminalBuffer::new(limit);
for index in 0..1024 {
buffer.append(&vec![b'0' + u8::try_from(index % 10).unwrap(); 1024]);
assert!(
buffer.buffered_len() <= limit + TERMINAL_BUFFER_SLACK_BYTES,
"a flooded buffer must stay bounded, held {}",
buffer.buffered_len()
);
}
let (output, truncated) = buffer.read();
assert!(truncated);
assert_eq!(output.len(), limit, "a read must never exceed the limit");
assert!(
output.ends_with(&"3".repeat(1024)),
"the retained bytes must be the tail"
);
}
#[test]
fn a_shell_line_carries_the_command_verbatim_and_quotes_the_arguments() {
assert_eq!(
shell_line(
"/bin/bash",
&["-c".to_owned(), "echo 'hi'; rm -rf /".to_owned()]
),
"/bin/bash '-c' 'echo '\\''hi'\\''; rm -rf /'"
);
assert_eq!(
shell_line("/bin/bash -lc 'echo hi'", &[]),
"/bin/bash -lc 'echo hi'"
);
}
#[test]
fn display_command_unwraps_an_interpreter_script() {
let spawn = TerminalSpawn {
command: "/bin/bash".into(),
args: vec!["-c".into(), "cargo mutants --in-diff diff".into()],
env: Vec::new(),
cwd: PathBuf::from("/workspace"),
output_byte_limit: 1024,
};
assert_eq!(spawn.display_command(), "cargo mutants --in-diff diff");
}
#[test]
fn display_command_preserves_a_non_interpreter_invocation() {
let spawn = TerminalSpawn {
command: "/usr/bin/cargo".into(),
args: vec!["test".into(), "--workspace".into()],
env: Vec::new(),
cwd: PathBuf::from("/workspace"),
output_byte_limit: 1024,
};
assert_eq!(
spawn.display_command(),
"/usr/bin/cargo 'test' '--workspace'"
);
}
fn register_stuck_terminal(registry: &TerminalRegistry, terminal_id: &str) {
let (_exit, exit_rx) = watch::channel(Some(TerminalExit::default()));
registry
.terminals
.lock()
.expect("terminal registry lock poisoned")
.insert(
terminal_id.to_owned(),
TerminalEntry {
group: ProcessGroup {
pid: i32::MAX,
exit: exit_rx.clone(),
},
buffer: Arc::new(Mutex::new(TerminalBuffer::new(1024))),
exit: exit_rx,
supervisor: tokio::spawn(std::future::pending::<()>()),
},
);
}
#[tokio::test(start_paused = true)]
async fn shutdown_reaps_stuck_terminals_under_one_shared_deadline() {
let registry = TerminalRegistry::new();
let (events, mut reports) = mpsc::channel(16);
for index in 0..4 {
register_stuck_terminal(®istry, &format!("term-{index}"));
}
let started = tokio::time::Instant::now();
registry.shutdown(&events).await;
let elapsed = started.elapsed();
assert!(
elapsed < TERMINAL_REAP_TIMEOUT * 2,
"teardown must be bounded by one reap timeout, took {elapsed:?}"
);
let mut reported = Vec::new();
while let Ok(RuntimeEvent::Warning { message }) = reports.try_recv() {
reported.push(message);
}
assert_eq!(reported.len(), 4, "{reported:?}");
for index in 0..4 {
assert!(
reported
.iter()
.any(|message| message.contains(&format!("term-{index} did not finish"))),
"every stuck terminal must still be reported: {reported:?}"
);
}
}
#[tokio::test(start_paused = true)]
async fn shutdown_reports_a_supervisor_that_finished() {
let registry = TerminalRegistry::new();
let (events, mut reports) = mpsc::channel(16);
register_stuck_terminal(®istry, "stuck");
let (_exit, exit_rx) = watch::channel(Some(TerminalExit::default()));
registry
.terminals
.lock()
.expect("terminal registry lock poisoned")
.insert(
"finished".to_owned(),
TerminalEntry {
group: ProcessGroup {
pid: i32::MAX,
exit: exit_rx.clone(),
},
buffer: Arc::new(Mutex::new(TerminalBuffer::new(1024))),
exit: exit_rx,
supervisor: tokio::spawn(async {}),
},
);
registry.shutdown(&events).await;
let mut reported = Vec::new();
while let Ok(RuntimeEvent::Warning { message }) = reports.try_recv() {
reported.push(message);
}
assert_eq!(
reported.len(),
1,
"a supervisor that finished is not a failure: {reported:?}"
);
assert!(reported[0].contains("stuck did not finish"), "{reported:?}");
}
#[cfg(unix)]
#[test]
fn a_killed_child_reports_its_signal_by_name() {
assert_eq!(signal_name(libc::SIGKILL), "SIGKILL");
assert_eq!(signal_name(libc::SIGTERM), "SIGTERM");
assert_eq!(signal_name(64), "SIG64");
}
}