use anyhow::Result;
use nar_dev_utils::{debug_println, ResultBoost};
use std::{
error::Error,
ffi::OsStr,
fmt::{self, Debug, Display, Formatter},
io::{BufRead, BufReader, ErrorKind, Result as IoResult, Write},
process::{Child, ChildStdin, ChildStdout, Command, ExitStatus, Stdio},
sync::{
mpsc::{channel, Receiver, Sender},
Arc, Mutex,
},
thread::{self, JoinHandle},
};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct IoProcessError(String);
impl Display for IoProcessError {
fn fmt(&self, f: &mut Formatter) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl Error for IoProcessError {}
fn err(e: impl Debug) -> anyhow::Error {
IoProcessError(format!("{e:?}")).into()
}
type OutputListener = dyn FnMut(String) + Send + Sync;
type ArcMutex<T> = Arc<Mutex<T>>;
pub struct IoProcess {
command: Command,
out_listener: Option<Box<OutputListener>>,
}
impl IoProcess {
pub fn new(program_path: impl AsRef<OsStr>) -> Self {
let command = Command::new(program_path);
Self::from(command)
}
pub fn arg(mut self, arg: impl AsRef<OsStr>) -> Self {
self.command.arg(arg);
self
}
pub fn out_listener<F>(mut self, listener: F) -> Self
where
F: FnMut(String) + Send + Sync + 'static,
{
self.out_listener = Some(Box::new(listener));
self
}
pub fn launch(self) -> Result<IoProcessManager> {
Ok(self.try_launch()?)
}
pub fn try_launch(mut self) -> std::io::Result<IoProcessManager> {
let child =
self.command
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn()?;
debug_println!("Started process: {}", child.id());
let out_listener = self.out_listener;
Ok(IoProcessManager::new(child, out_listener))
}
}
impl From<Command> for IoProcess {
fn from(command: Command) -> Self {
Self {
command,
out_listener: None,
}
}
}
#[allow(dead_code)]
pub struct IoProcessManager {
process: Child,
thread_write_in: Option<JoinHandle<()>>,
thread_read_out: Option<JoinHandle<()>>,
termination_signal: ArcMutex<bool>,
child_out: Mutex<Receiver<String>>,
child_in: Mutex<Sender<String>>,
}
impl IoProcessManager {
pub fn new(mut child: Child, out_listener: Option<Box<OutputListener>>) -> Self {
let stdin = child.stdin.take().unwrap();
let stdout = child.stdout.take().unwrap();
let (child_out, out_sender) = channel();
let (in_receiver, child_in) = channel();
let termination_signal = Arc::new(Mutex::new(false));
let thread_write_in = Some(IoProcessManager::spawn_thread_write_in(
stdin,
child_in,
termination_signal.clone(),
));
let thread_read_out = Some(IoProcessManager::spawn_thread_read_out(
stdout,
child_out,
out_listener,
termination_signal.clone(),
));
Self {
process: child,
thread_read_out,
thread_write_in,
child_out: Mutex::new(out_sender),
child_in: Mutex::new(in_receiver),
termination_signal,
}
}
#[inline]
fn spawn_thread_write_in(
stdin: ChildStdin,
child_in_receiver: Receiver<String>,
termination_signal: ArcMutex<bool>,
) -> thread::JoinHandle<()> {
thread::spawn(move || {
let mut stdin = stdin;
for line in child_in_receiver {
if *termination_signal.lock().expect("无法锁定终止信号") {
break;
}
if let Err(e) = stdin.write_all(line.as_bytes()) {
match e.kind() {
ErrorKind::BrokenPipe => {
eprintln!("[IoProcessManager] 子进程已关闭");
break;
}
_ => eprintln!("[IoProcessManager] 子进程写入错误:{e}"),
}
}
}
})
}
#[inline]
fn spawn_thread_read_out(
stdout: ChildStdout,
child_out_sender: Sender<String>,
out_listener: Option<Box<dyn FnMut(String) + Send + Sync>>,
termination_signal: ArcMutex<bool>,
) -> thread::JoinHandle<()> {
let mut listener_code: Box<dyn FnMut(&String) + Send + Sync> = match out_listener {
Some(mut listener) => Box::new(move |s: &String| listener(s.clone())),
None => Box::new(move |_| {}),
};
thread::spawn(move || {
let mut stdout_reader = BufReader::new(stdout);
let mut buf = String::new();
loop {
match stdout_reader.read_line(&mut buf) {
Ok(0) => {
if *termination_signal.lock().expect("无法锁定终止信号") {
break;
}
}
Ok(_) => {
listener_code(&buf);
if let Err(e) = child_out_sender.send(buf.clone()) {
println!("无法向主进程发送消息:{e:?}");
break;
}
}
Err(e) => {
let message = e.to_string();
if message.contains("stream did not contain") {
} else {
println!("无法接收子进程输出:{e:?} in「{buf}」");
break;
}
}
}
buf.clear();
}
})
}
pub fn id(&self) -> u32 {
self.process.id()
}
pub fn fetch_output(&mut self) -> Result<String> {
self.child_out
.lock()
.transform_err(err)?
.recv()
.transform_err(err)
}
pub fn try_fetch_output(&mut self) -> Result<Option<String>> {
let out = self
.child_out
.lock()
.transform_err(err)?
.try_recv()
.ok();
Ok(out)
}
pub fn put(&self, input_line: impl ToString) -> Result<()> {
self.child_in
.lock()
.transform_err(err)?
.send(input_line.to_string())
.transform_err(err)
}
pub fn put_line(&self, input: impl ToString) -> Result<()> {
self.put(format!("{}\n", input.to_string()))
}
pub fn wait(&mut self) -> IoResult<ExitStatus> {
self.process.wait()
}
pub fn kill(&mut self) -> Result<()> {
let mut signal = self.termination_signal.lock().transform_err(err)?;
*signal = true;
drop(signal);
let _ = self
.put("\n")
.inspect_err(|e| println!("向「进程读取」子线程发送消息失败!{e}"));
drop(
self.thread_write_in
.take()
.map(|t| t.join().transform_err(err)),
); drop(self.thread_read_out.take());
if let Ok(child) = Command::new("taskkill")
.args(["-F", "-PID", &self.process.id().to_string()])
.spawn()
{
if let Err(err) = child.wait_with_output() {
println!("指令执行失败!{err:?}");
}
}
self.process.kill().transform_err(err)
}
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
use crate::tests::cin_paths::ONA as PATH_ONA;
use std::{
process::exit,
sync::{Arc, Mutex},
};
pub fn await_fetch_until(process: &mut IoProcessManager, criterion: impl Fn(String) -> bool) {
loop {
let out = process.fetch_output().expect("无法拉取输出");
println!("fetch到其中一个输出: {out:?}");
if criterion(out) {
break;
}
}
}
fn launch_ona() -> (IoProcessManager, ArcMutex<Vec<String>>) {
let outputs = Arc::new(Mutex::new(vec![]));
let outputs_inner = outputs.clone();
let process = IoProcess::new(PATH_ONA)
.arg("shell")
.out_listener(move |output: String| {
outputs_inner
.lock()
.expect("无法锁定 outputs_inner")
.push(output.clone());
print!("[OUT] {}", output);
});
(process.launch().expect("ONA启动失败"), outputs)
}
#[test]
fn test_ona() {
let (mut process, outputs) = launch_ona();
let output_must_contains = |s: &str| {
let outputs = outputs.lock().expect("无法锁定 outputs");
let line = outputs
.iter()
.find(|line| line.contains(s))
.expect("没有指定的输出!");
println!("检验「{s:?}」成功!所在之处:{line:?}");
};
let input = "<A --> B>.";
process.put_line(input).expect("无法放置输入");
await_fetch_until(&mut process, |s| s.contains(input));
output_must_contains("<A --> B>.");
process.put("<B --> C>.\n").expect("无法放置输入");
await_fetch_until(&mut process, |s| s.contains("<B --> C>."));
process.put("<A --> C>?\n").expect("无法放置输入");
await_fetch_until(&mut process, |s| s.contains("<A --> C>?"));
const EXPECTED_ANSWER: &str = "Answer: <A --> C>.";
await_fetch_until(&mut process, |s| s.contains(EXPECTED_ANSWER));
output_must_contains(EXPECTED_ANSWER);
{
let r = process.child_out.lock().unwrap();
for _ in r.try_iter() {
let line = r.recv().expect("接收失败!");
print!("从输出中读取到的一行(多了会阻塞!):{line}");
}
}
process.kill().expect("无法杀死进程");
println!("Process killed.");
dbg!(&*outputs);
exit(0);
}
}