use crate::CapturedOutput;
use std::{
fs::File,
io,
process::{ChildStderr, ChildStdout},
sync::{
atomic::{AtomicBool, Ordering},
Arc, Condvar, Mutex,
},
thread::{self, JoinHandle},
time::Duration,
};
#[cfg(unix)]
#[path = "native_capture/unix.rs"]
mod platform;
#[cfg(windows)]
#[path = "native_capture/windows.rs"]
pub(crate) mod platform;
const CHUNK: usize = 8 * 1024;
const IDLE_WAIT: Duration = Duration::from_millis(20);
#[derive(Default)]
struct State {
output: CapturedOutput,
error: Option<io::Error>,
eof: bool,
finished: bool,
}
type Shared = (Mutex<State>, Condvar);
struct Done(Arc<Shared>);
impl Drop for Done {
fn drop(&mut self) {
self.0 .0.lock().unwrap_or_else(|e| e.into_inner()).finished = true;
self.0 .1.notify_all();
}
}
pub struct NativeOutputCapture {
shared: Arc<Shared>,
stop: Arc<AtomicBool>,
handle: Option<JoinHandle<()>>,
}
impl NativeOutputCapture {
pub fn stdout(pipe: ChildStdout, max_bytes: usize) -> io::Result<Self> {
Self::start(platform::stdout(pipe), max_bytes)
}
pub fn stderr(pipe: ChildStderr, max_bytes: usize) -> io::Result<Self> {
Self::start(platform::stderr(pipe), max_bytes)
}
fn start(pipe: File, max_bytes: usize) -> io::Result<Self> {
let prepared = platform::prepare(pipe)?;
let shared: Arc<Shared> = Arc::default();
let stop = Arc::new(AtomicBool::new(false));
let state = Arc::clone(&shared);
let cancelled = Arc::clone(&stop);
let handle = thread::Builder::new()
.name("rightkit-native-output-capture".into())
.spawn(move || {
let _done = Done(Arc::clone(&state));
let mut pipe = platform::reader(prepared);
let mut buffer = [0; CHUNK];
while !cancelled.load(Ordering::Acquire) {
match platform::read(&mut pipe, &mut buffer) {
Ok(0) => {
state.0.lock().unwrap_or_else(|e| e.into_inner()).eof = true;
break;
}
Ok(count) => state
.0
.lock()
.unwrap_or_else(|e| e.into_inner())
.output
.push(&buffer[..count], max_bytes),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
thread::park_timeout(IDLE_WAIT);
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => {
state.0.lock().unwrap_or_else(|e| e.into_inner()).error = Some(error);
break;
}
}
}
if let Err(error) = platform::shutdown(&mut pipe) {
let mut state = state.0.lock().unwrap_or_else(|e| e.into_inner());
if state.error.is_none() {
state.error = Some(error);
}
}
drop(pipe);
})?;
Ok(Self {
shared,
stop,
handle: Some(handle),
})
}
pub fn snapshot(&self) -> CapturedOutput {
self.shared
.0
.lock()
.unwrap_or_else(|e| e.into_inner())
.output
.clone()
}
pub fn is_complete(&self) -> bool {
self.shared.0.lock().unwrap_or_else(|e| e.into_inner()).eof
}
pub fn take_error(&self) -> Option<io::Error> {
self.shared
.0
.lock()
.unwrap_or_else(|e| e.into_inner())
.error
.take()
}
pub fn finish_within(&mut self, timeout: Duration) -> CapturedOutput {
{
let state = self.shared.0.lock().unwrap_or_else(|e| e.into_inner());
drop(
self.shared
.1
.wait_timeout_while(state, timeout, |state| !state.finished)
.unwrap_or_else(|e| e.into_inner()),
);
}
self.stop_and_join();
self.snapshot()
}
pub fn cancel(&mut self) -> CapturedOutput {
self.stop_and_join();
self.snapshot()
}
fn stop_and_join(&mut self) {
self.stop.store(true, Ordering::Release);
if let Some(handle) = self.handle.take() {
handle.thread().unpark();
if handle.join().is_err() {
self.shared
.0
.lock()
.unwrap_or_else(|e| e.into_inner())
.error = Some(io::Error::other("native pipe reader panicked"));
}
}
}
}
impl Drop for NativeOutputCapture {
fn drop(&mut self) {
self.stop_and_join();
}
}