use std::ffi::{OsStr, OsString};
use std::io::{self, Read, Write as _};
use std::path::PathBuf;
use std::process::{Child, Command, ExitStatus, Output, Stdio};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel};
use std::thread;
use std::time::{Duration, Instant};
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Launcher {
pub prefix: Vec<OsString>,
pub cwd: Option<PathBuf>,
pub env: Vec<(OsString, OsString)>,
pub unset_env: Vec<OsString>,
}
impl Launcher {
pub fn wrapped(mut self, prefix: impl IntoIterator<Item = impl Into<OsString>>) -> Self {
self.prefix = prefix.into_iter().map(Into::into).collect();
self
}
pub fn in_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.cwd = Some(dir.into());
self
}
pub fn with_env(mut self, name: impl Into<OsString>, value: impl Into<OsString>) -> Self {
self.env.push((name.into(), value.into()));
self
}
pub fn without_env(mut self, names: impl IntoIterator<Item = impl Into<OsString>>) -> Self {
self.unset_env.extend(names.into_iter().map(Into::into));
self
}
pub fn command(&self, program: impl AsRef<OsStr>) -> Command {
let mut command = match self.prefix.split_first() {
Some((wrapper, arguments)) => {
let mut command = Command::new(wrapper);
command.args(arguments).arg(program);
command
}
None => Command::new(program),
};
if let Some(dir) = &self.cwd {
command.current_dir(dir);
}
for name in &self.unset_env {
command.env_remove(name);
}
command.envs(self.env.iter().map(|(name, value)| (name, value)));
detach(&mut command);
command
}
pub fn spawn(
&self,
program: impl AsRef<OsStr>,
args: impl IntoIterator<Item = impl AsRef<OsStr>>,
how: Spawn,
) -> io::Result<Process> {
let mut command = self.command(program);
command.args(args).stdin(match how.stdin {
Some(_) => Stdio::piped(),
None => Stdio::null(),
});
let (sender, lines) = channel();
let mut pipes: Vec<(Box<dyn Read + Send>, Tag)> = Vec::new();
if how.merge_stderr {
let (reader, writer) = io::pipe()?;
command.stdout(writer.try_clone()?).stderr(writer);
pipes.push((Box::new(reader), OutputLine::Out));
} else {
command.stdout(Stdio::piped()).stderr(Stdio::piped());
}
let spawned = command.spawn();
drop(command);
let child = spawned?;
let group = Group(Arc::new(AtomicU32::new(child.id())));
let mut process = Process {
child,
lines,
group,
remote: None,
forced: false,
};
let (stdout, stderr) = (process.child.stdout.take(), process.child.stderr.take());
#[cfg(unix)]
let (stdout, stderr) = match (stdout, stderr) {
(Some(stdout), Some(stderr)) => {
let sender = sender.clone();
thread::Builder::new()
.name("rig-process-output".to_owned())
.spawn(move || forward_both(stdout, stderr, &sender))?;
(None, None)
}
apart => apart,
};
if let Some(stdout) = stdout {
pipes.push((Box::new(stdout), OutputLine::Out));
}
if let Some(stderr) = stderr {
pipes.push((Box::new(stderr), OutputLine::Err));
}
for (pipe, tag) in pipes {
let sender = sender.clone();
thread::Builder::new()
.name("rig-process-output".to_owned())
.spawn(move || forward(pipe, tag, &sender))?;
}
if let (Some(mut stdin), Some(bytes)) = (process.child.stdin.take(), how.stdin) {
thread::Builder::new()
.name("rig-process-input".to_owned())
.spawn(move || {
stdin.write_all(&bytes).ok();
})?;
}
Ok(process)
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Spawn {
pub stdin: Option<Vec<u8>>,
pub merge_stderr: bool,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OutputLine {
Out(String),
Err(String),
}
impl OutputLine {
pub fn text(&self) -> &str {
match self {
Self::Out(text) | Self::Err(text) => text,
}
}
}
type Tag = fn(String) -> OutputLine;
const POLL: Duration = Duration::from_millis(10);
const PIPE_GRACE: Duration = Duration::from_millis(100);
pub const STOP_GRACE: Duration = Duration::from_secs(2);
const MAX_PIECE: usize = 64 * 1024;
fn forward(mut pipe: impl Read, tag: Tag, lines: &Sender<OutputLine>) {
let mut pending = Vec::new();
let mut buffer = [0; 8192];
loop {
let read = match pipe.read(&mut buffer) {
Ok(0) => break,
Ok(read) => read,
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
Err(_) => break,
};
pending.extend_from_slice(buffer.get(..read).unwrap_or_default());
let mut start = 0;
while let Some(end) = pending.get(start..).and_then(piece_end) {
let piece = pending.get(start..start + end).unwrap_or_default();
if lines
.send(tag(String::from_utf8_lossy(piece).into_owned()))
.is_err()
{
return;
}
start += end;
}
pending.drain(..start);
}
if !pending.is_empty() {
lines
.send(tag(String::from_utf8_lossy(&pending).into_owned()))
.ok();
}
}
#[cfg(unix)]
fn forward_both(
stdout: std::process::ChildStdout,
stderr: std::process::ChildStderr,
lines: &Sender<OutputLine>,
) {
use std::os::fd::AsRawFd;
let fds = [stdout.as_raw_fd(), stderr.as_raw_fd()];
let mut pipes: [(Box<dyn Read>, Tag, Vec<u8>, bool); 2] = [
(Box::new(stdout), OutputLine::Out, Vec::new(), true),
(Box::new(stderr), OutputLine::Err, Vec::new(), true),
];
let mut buffer = [0; 8192];
while pipes.iter().any(|(.., open)| *open) {
let mut polled = fds.map(|fd| libc::pollfd {
fd,
events: libc::POLLIN,
revents: 0,
});
for (poll, (.., open)) in polled.iter_mut().zip(&pipes) {
if !*open {
poll.fd = -1;
}
}
let ready = unsafe { libc::poll(polled.as_mut_ptr(), 2, -1) };
if ready < 0 {
if io::Error::last_os_error().kind() == io::ErrorKind::Interrupted {
continue;
}
break;
}
for (poll, (pipe, tag, pending, open)) in polled.iter().zip(&mut pipes) {
if !*open || poll.revents == 0 {
continue;
}
let read = match pipe.read(&mut buffer) {
Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
Ok(0) | Err(_) => None,
Ok(read) => Some(read),
};
let Some(read) = read else {
*open = false;
if !pending.is_empty() {
let text = String::from_utf8_lossy(pending).into_owned();
pending.clear();
if lines.send(tag(text)).is_err() {
return;
}
}
continue;
};
pending.extend_from_slice(buffer.get(..read).unwrap_or_default());
let mut start = 0;
while let Some(end) = pending.get(start..).and_then(piece_end) {
let piece = pending.get(start..start + end).unwrap_or_default();
if lines
.send(tag(String::from_utf8_lossy(piece).into_owned()))
.is_err()
{
return;
}
start += end;
}
pending.drain(..start);
}
}
}
fn piece_end(bytes: &[u8]) -> Option<usize> {
if let Some(at) = bytes.iter().position(|&byte| byte == b'\n') {
return Some(at + 1);
}
let head = bytes.get(..MAX_PIECE)?;
match std::str::from_utf8(head) {
Err(error) if error.error_len().is_none() && error.valid_up_to() > 0 => {
Some(error.valid_up_to())
}
_ => Some(head.len()),
}
}
pub struct Process {
child: Child,
lines: Receiver<OutputLine>,
group: Group,
remote: Option<String>,
forced: bool,
}
impl Process {
pub fn id(&self) -> u32 {
self.child.id()
}
pub fn through(mut self, host: impl Into<String>) -> Self {
self.remote = Some(host.into());
self
}
pub fn remote(&self) -> Option<&str> {
self.remote.as_deref()
}
pub fn forced(&self) -> bool {
self.forced
}
pub fn lines(&self) -> &Receiver<OutputLine> {
&self.lines
}
pub fn group(&self) -> Group {
self.group.clone()
}
pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
let status = self.child.try_wait();
if !matches!(status, Ok(None)) {
self.group.0.store(0, Ordering::SeqCst);
}
status
}
pub fn wait_for(&mut self, timeout: Duration) -> io::Result<Option<ExitStatus>> {
let deadline = Instant::now().checked_add(timeout);
let mut nap = Duration::from_millis(1);
loop {
if let Some(status) = self.try_wait()? {
return Ok(Some(status));
}
let left = deadline.map_or(POLL, |deadline| {
deadline.saturating_duration_since(Instant::now())
});
if left.is_zero() {
return Ok(None);
}
thread::sleep(left.min(nap));
nap = (nap * 2).min(POLL);
}
}
pub fn output(mut self, timeout: Option<Duration>) -> io::Result<Output> {
let deadline = timeout.and_then(|timeout| Instant::now().checked_add(timeout));
let (mut stdout, mut stderr) = (Vec::new(), Vec::new());
let ended = self.run(deadline, None, |line| {
match line {
OutputLine::Out(text) => stdout.extend(text.into_bytes()),
OutputLine::Err(text) => stderr.extend(text.into_bytes()),
}
true
})?;
match ended {
Ended::Exited(status) => Ok(Output {
status,
stdout,
stderr,
}),
Ended::TimedOut | Ended::Quiet | Ended::Stopped => Err(io::Error::new(
io::ErrorKind::TimedOut,
"the process ran past its timeout and was killed",
)),
}
}
pub fn run(
&mut self,
deadline: Option<Instant>,
quiet: Option<Duration>,
mut sink: impl FnMut(OutputLine) -> bool,
) -> io::Result<Ended> {
let mut more = true;
let mut heard = Instant::now();
let ended = loop {
if let Some(status) = self.try_wait()? {
break Ended::Exited(status);
}
if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
break Ended::TimedOut;
}
if quiet.is_some_and(|quiet| heard.elapsed() >= quiet) {
break Ended::Quiet;
}
if !more {
break Ended::Stopped;
}
match self.lines.recv_timeout(POLL) {
Ok(line) => {
heard = Instant::now();
more = sink(line);
}
Err(RecvTimeoutError::Timeout) => {}
Err(RecvTimeoutError::Disconnected) => {
self.wait_for(POLL)?;
}
}
};
if !matches!(ended, Ended::Exited(_)) {
self.stop(STOP_GRACE).ok();
self.read_rest(&mut sink);
} else if !self.read_rest(&mut sink) {
self.kill().ok();
self.read_rest(&mut sink);
}
Ok(ended)
}
fn read_rest(&self, sink: &mut impl FnMut(OutputLine) -> bool) -> bool {
let deadline = Instant::now() + PIPE_GRACE;
loop {
match self
.lines
.recv_timeout(deadline.saturating_duration_since(Instant::now()))
{
Ok(line) => {
sink(line);
}
Err(RecvTimeoutError::Disconnected) => return true,
Err(RecvTimeoutError::Timeout) => return false,
}
}
}
pub fn stop(&mut self, grace: Duration) -> io::Result<ExitStatus> {
self.forced = false;
if self.try_wait()?.is_none() {
self.group.terminate();
self.forced = self.wait_for(grace)?.is_none();
}
self.kill()
}
pub fn kill(&mut self) -> io::Result<ExitStatus> {
self.group.0.store(0, Ordering::SeqCst);
#[cfg(unix)]
kill_group_of(self.child.id());
#[cfg(not(unix))]
if self.child.try_wait()?.is_none() {
self.child.kill()?;
}
self.child.wait()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Ended {
Exited(ExitStatus),
TimedOut,
Quiet,
Stopped,
}
impl Drop for Process {
fn drop(&mut self) {
if self.group.0.load(Ordering::SeqCst) != 0 {
self.kill().ok();
}
}
}
#[derive(Clone, Debug)]
pub struct Group(Arc<AtomicU32>);
impl Group {
pub fn kill(&self) {
self.signal(Signal::Kill);
}
pub fn terminate(&self) {
self.signal(Signal::Terminate);
}
fn signal(&self, signal: Signal) {
let leader = self.0.load(Ordering::SeqCst);
#[cfg(unix)]
if leader != 0 {
signal_group(leader, signal);
}
#[cfg(not(unix))]
let _ = (leader, signal);
}
}
pub struct KillOnDrop(pub Group);
impl Drop for KillOnDrop {
fn drop(&mut self) {
self.0.terminate();
let group = self.0.clone();
let later = thread::Builder::new()
.name("rig-process-stop".to_owned())
.spawn(move || {
thread::sleep(STOP_GRACE);
group.kill();
});
if later.is_err() {
self.0.kill();
}
}
}
#[derive(Clone, Copy)]
enum Signal {
Terminate,
Kill,
}
pub fn detach(command: &mut Command) {
command.env("GIT_TERMINAL_PROMPT", "0");
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
unsafe {
command.pre_exec(|| {
if libc::setsid() == -1 {
Err(std::io::Error::last_os_error())
} else {
Ok(())
}
});
}
}
}
#[cfg(unix)]
pub fn kill_group(child: &mut Child) {
kill_group_of(child.id());
}
#[cfg(unix)]
pub fn kill_group_of(leader: u32) {
signal_group(leader, Signal::Kill);
}
#[cfg(unix)]
fn signal_group(leader: u32, signal: Signal) {
let signal = match signal {
Signal::Terminate => libc::SIGTERM,
Signal::Kill => libc::SIGKILL,
};
if let Ok(group) = i32::try_from(leader) {
unsafe {
libc::kill(-group, signal);
}
}
}
#[cfg(not(unix))]
pub fn kill_group(child: &mut Child) {
child.kill().ok();
}
#[cfg(test)]
mod tests;