use std::{
ffi::OsString,
os::unix::process::CommandExt as _,
path::PathBuf,
sync::{
Arc, Mutex as StdMutex,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use scv_core::ToolError;
use scv_protocol::{Frame, FrameDecoder, Overflow};
use serde::Serialize;
use tokio::{
io::{AsyncWriteExt as _, BufReader},
process::{Child, ChildStdin, ChildStdout, Command},
sync::{Mutex, mpsc, watch},
};
use crate::{
delegate::{
output::TailBuffer,
records::{
self, DelegationGuard, DelegationRegistry, PendingDelegation, STOP_GRACE, group_exists,
},
},
process::{ProcessGroup, apply_agent_environment, drain_output},
sync::lock,
};
const STDERR_TAIL_BYTES: usize = 4096;
const LINE_QUEUE: usize = 256;
pub(crate) struct LiveSpec {
pub(crate) executable: OsString,
pub(crate) args: Vec<OsString>,
pub(crate) cwd: PathBuf,
pub(crate) environment: Vec<(OsString, OsString)>,
pub(crate) max_line_bytes: usize,
}
#[derive(Debug)]
pub(crate) enum LiveLine {
Line(Vec<u8>),
TooLong,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Life {
Running,
Exited,
Finished,
}
pub(crate) struct LiveChild {
pid: u32,
stdin: Mutex<Option<ChildStdin>>,
lines: Mutex<mpsc::Receiver<LiveLine>>,
stderr: Arc<Mutex<TailBuffer>>,
guard: Arc<StdMutex<Option<DelegationGuard>>>,
life: watch::Receiver<Life>,
closed: AtomicBool,
}
impl std::fmt::Debug for LiveChild {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("LiveChild")
.field("pid", &self.pid)
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish()
}
}
impl LiveChild {
pub(crate) fn spawn(
spec: LiveSpec,
registration: Option<(Arc<DelegationRegistry>, PendingDelegation)>,
) -> Result<Arc<Self>, ToolError> {
let mut command = Command::new(&spec.executable);
apply_agent_environment(command.as_std_mut(), &spec.environment);
command
.args(&spec.args)
.current_dir(&spec.cwd)
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
command.as_std_mut().process_group(0);
let mut child = command.spawn().map_err(|error| {
ToolError::unavailable(format!("launch {:?}: {error}", spec.executable))
})?;
let pid = child
.id()
.ok_or_else(|| ToolError::failed("child process has no pid"))?;
records::track_spawned(pid);
let guard = Arc::new(StdMutex::new(
registration.and_then(|(registry, pending)| registry.register(pending, pid).ok()),
));
let stdin = child.stdin.take();
let stdout = child.stdout.take();
let stderr = Arc::new(Mutex::new(TailBuffer::new(STDERR_TAIL_BYTES)));
if let Some(reader) = child.stderr.take() {
tokio::spawn(drain_output(reader, Arc::clone(&stderr)));
}
let (sender, receiver) = mpsc::channel(LINE_QUEUE);
if let Some(stdout) = stdout {
tokio::spawn(read_lines(stdout, sender, spec.max_line_bytes));
}
let (life_tx, life) = watch::channel(Life::Running);
tokio::spawn(reap(pid, child, Arc::clone(&guard), life_tx));
Ok(Arc::new(Self {
pid,
stdin: Mutex::new(stdin),
lines: Mutex::new(receiver),
stderr,
guard,
life,
closed: AtomicBool::new(false),
}))
}
pub(crate) async fn send(&self, message: &impl Serialize) -> Result<(), ToolError> {
let mut bytes = serde_json::to_vec(message)
.map_err(|error| ToolError::failed(format!("encode message: {error}")))?;
bytes.push(b'\n');
let mut stdin = self.stdin.lock().await;
let writer = stdin
.as_mut()
.ok_or_else(|| ToolError::failed("the child's input is closed"))?;
writer
.write_all(&bytes)
.await
.map_err(|error| ToolError::failed(format!("write to child: {error}")))?;
writer
.flush()
.await
.map_err(|error| ToolError::failed(format!("write to child: {error}")))
}
pub(crate) async fn recv(&self) -> Option<LiveLine> {
self.lines.lock().await.recv().await
}
pub(crate) async fn stderr_tail(&self) -> String {
self.stderr.lock().await.text()
}
pub(crate) fn is_running(&self) -> bool {
!self.closed.load(Ordering::Acquire) && *self.life.borrow() == Life::Running
}
pub(crate) fn begin_turn(&self, turn: u32) -> LiveTurn<'_> {
if let Some(guard) = lock(&self.guard).as_ref() {
guard.set_turn(turn);
}
LiveTurn(self)
}
pub(crate) fn set_background_jobs(&self, jobs: usize) {
if let Some(guard) = lock(&self.guard).as_ref() {
guard.set_background_jobs(jobs);
}
}
pub(crate) async fn close(&self) {
if self.closed.swap(true, Ordering::AcqRel) {
return;
}
let stdin = self.stdin.lock().await.take();
shut_down(self.pid, stdin, self.life.clone()).await;
}
}
#[must_use = "the turn ends when this guard drops"]
pub(crate) struct LiveTurn<'a>(&'a LiveChild);
impl Drop for LiveTurn<'_> {
fn drop(&mut self) {
if let Some(guard) = lock(&self.0.guard).as_ref() {
guard.set_idle();
}
}
}
impl Drop for LiveChild {
fn drop(&mut self) {
if self.closed.swap(true, Ordering::AcqRel) {
return;
}
let stdin = self.stdin.get_mut().take();
let pid = self.pid;
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(shut_down(pid, stdin, self.life.clone()));
} else {
kill_group(pid);
records::untrack_spawned(pid);
drop(lock(&self.guard).take());
}
}
}
async fn shut_down(pid: u32, stdin: Option<ChildStdin>, mut life: watch::Receiver<Life>) {
drop(stdin);
if tokio::time::timeout(STOP_GRACE, life.wait_for(|state| *state != Life::Running))
.await
.is_err()
{
kill_group(pid);
}
let _ = tokio::time::timeout(
STOP_GRACE * 2 + Duration::from_millis(500),
life.wait_for(|state| *state == Life::Finished),
)
.await;
}
async fn reap(
pid: u32,
mut child: Child,
guard: Arc<StdMutex<Option<DelegationGuard>>>,
life: watch::Sender<Life>,
) {
let _ = child.wait().await;
let _ = life.send(Life::Exited);
if group_exists(pid) {
kill_group(pid);
}
records::untrack_spawned(pid);
let guard = lock(&guard).take();
if let Some(guard) = guard {
guard.finish().await;
}
let _ = life.send(Life::Finished);
}
async fn read_lines(stdout: ChildStdout, sender: mpsc::Sender<LiveLine>, max_bytes: usize) {
let mut reader = BufReader::new(stdout);
let mut decoder = FrameDecoder::new(max_bytes, Overflow::Stop);
loop {
let line = match scv_client::read_frame(&mut reader, &mut decoder).await {
Ok(Frame::Line(line)) => LiveLine::Line(line),
Ok(Frame::TooLarge) => {
let _ = sender.send(LiveLine::TooLong).await;
return;
}
Ok(Frame::End | Frame::Truncated(_)) | Err(_) => return,
};
if sender.send(line).await.is_err() {
return;
}
}
}
fn kill_group(pid: u32) {
if let Some(group) = ProcessGroup::new(pid) {
group.signal(libc::SIGKILL);
}
}
#[cfg(test)]
mod tests;