use std::collections::HashMap;
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use super::protocol::{ScriptCommand, ScriptEvent};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScriptStatus {
Running,
Exited { success: bool },
}
pub struct ScriptProcess {
child: Option<Child>,
stdin_writer: Option<std::process::ChildStdin>,
command_buffer: Arc<Mutex<Vec<ScriptCommand>>>,
error_buffer: Arc<Mutex<Vec<String>>>,
_stdout_thread: Option<JoinHandle<()>>,
_stderr_thread: Option<JoinHandle<()>>,
exit: Option<ScriptStatus>,
}
impl ScriptProcess {
pub fn spawn(
command: &str,
args: &[&str],
env_vars: &HashMap<String, String>,
) -> Result<Self, String> {
let mut cmd = Command::new(command);
cmd.args(args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.envs(env_vars);
let mut child = cmd
.spawn()
.map_err(|e| format!("Failed to spawn '{}': {}", command, e))?;
let stdin_writer = child.stdin.take();
let stdout = child
.stdout
.take()
.ok_or_else(|| "Failed to capture stdout".to_string())?;
let stderr = child
.stderr
.take()
.ok_or_else(|| "Failed to capture stderr".to_string())?;
let command_buffer: Arc<Mutex<Vec<ScriptCommand>>> = Arc::new(Mutex::new(Vec::new()));
let error_buffer: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let cmd_buf = Arc::clone(&command_buffer);
let stdout_thread = std::thread::spawn(move || {
let reader = BufReader::new(stdout);
for line in reader.lines() {
match line {
Ok(text) => {
if text.is_empty() {
continue;
}
match serde_json::from_str::<ScriptCommand>(&text) {
Ok(cmd) => {
let mut buf = cmd_buf.lock().unwrap_or_else(|e| {
log::warn!("command_buffer mutex poisoned, recovering");
e.into_inner()
});
buf.push(cmd);
}
Err(e) => {
log::warn!(
"ScriptProcess: failed to parse stdout line as ScriptCommand: {}: {:?}",
e,
text
);
}
}
}
Err(e) => {
log::warn!("ScriptProcess: error reading stdout: {}", e);
break;
}
}
}
});
let err_buf = Arc::clone(&error_buffer);
let stderr_thread = std::thread::spawn(move || {
let reader = BufReader::new(stderr);
for line in reader.lines() {
match line {
Ok(text) => {
if text.is_empty() {
continue;
}
let mut buf = err_buf.lock().unwrap_or_else(|e| {
log::warn!("error_buffer mutex poisoned, recovering");
e.into_inner()
});
buf.push(text);
}
Err(e) => {
log::warn!("ScriptProcess: error reading stderr: {}", e);
break;
}
}
}
});
Ok(Self {
child: Some(child),
stdin_writer,
command_buffer,
error_buffer,
_stdout_thread: Some(stdout_thread),
_stderr_thread: Some(stderr_thread),
exit: None,
})
}
pub fn poll_status(&mut self) -> ScriptStatus {
if let Some(status) = self.exit {
return status;
}
let polled = match self.child.as_mut() {
Some(child) => match child.try_wait() {
Ok(Some(s)) => Some(ScriptStatus::Exited {
success: s.success(),
}),
Ok(None) => None,
Err(_) => Some(ScriptStatus::Exited { success: false }),
},
None => Some(ScriptStatus::Exited { success: true }),
};
match polled {
Some(status) => {
self.stdin_writer.take();
self.child.take();
self.exit = Some(status);
status
}
None => ScriptStatus::Running,
}
}
pub fn is_running(&mut self) -> bool {
matches!(self.poll_status(), ScriptStatus::Running)
}
pub fn send_event(&mut self, event: &ScriptEvent) -> Result<(), String> {
let stdin = self
.stdin_writer
.as_mut()
.ok_or_else(|| "stdin writer is not available".to_string())?;
let json = serde_json::to_string(event)
.map_err(|e| format!("Failed to serialize event: {}", e))?;
writeln!(stdin, "{}", json).map_err(|e| format!("Failed to write to stdin: {}", e))?;
stdin
.flush()
.map_err(|e| format!("Failed to flush stdin: {}", e))?;
Ok(())
}
pub fn read_commands(&self) -> Vec<ScriptCommand> {
let mut buf = self.command_buffer.lock().unwrap_or_else(|e| {
log::warn!("command_buffer mutex poisoned, recovering");
e.into_inner()
});
buf.drain(..).collect()
}
pub fn read_errors(&self) -> Vec<String> {
let mut buf = self.error_buffer.lock().unwrap_or_else(|e| {
log::warn!("error_buffer mutex poisoned, recovering");
e.into_inner()
});
buf.drain(..).collect()
}
pub fn stop(&mut self) {
self.stdin_writer.take();
if let Some(ref mut child) = self.child {
let _ = child.kill();
let _ = child.wait();
}
self.child.take();
}
}
impl Drop for ScriptProcess {
fn drop(&mut self) {
self.stop();
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::collections::HashMap;
use std::time::Duration;
fn no_env() -> HashMap<String, String> {
HashMap::new()
}
#[test]
fn poll_status_running_for_live_process() {
let mut p = ScriptProcess::spawn("sleep", &["2"], &no_env()).expect("spawn sleep");
assert_eq!(p.poll_status(), ScriptStatus::Running);
p.stop();
}
#[test]
fn poll_status_exited_clean_for_exit_zero() {
let mut p = ScriptProcess::spawn("true", &[], &no_env()).expect("spawn true");
std::thread::sleep(Duration::from_millis(50));
assert_eq!(p.poll_status(), ScriptStatus::Exited { success: true });
assert_eq!(p.poll_status(), ScriptStatus::Exited { success: true });
}
#[test]
fn poll_status_exited_failure_for_nonzero_exit() {
let mut p = ScriptProcess::spawn("false", &[], &no_env()).expect("spawn false");
std::thread::sleep(Duration::from_millis(50));
assert_eq!(p.poll_status(), ScriptStatus::Exited { success: false });
}
#[test]
fn is_running_false_after_exit() {
let mut p = ScriptProcess::spawn("true", &[], &no_env()).expect("spawn true");
std::thread::sleep(Duration::from_millis(50));
assert!(!p.is_running());
}
}