rightkit_child_process/

native_capture.rs

//! Cancellable capture of native child pipes. Generic `Read` capture remains separate.

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();
    }
}

/// An owned reader of a native child stdout/stderr pipe, with bounded byte retention.
///
/// Unix read ends use nonblocking I/O; Windows read ends use owned overlapped I/O.
/// Cancellation wakes the private reader; no user callbacks run
/// on it. [`cancel`](Self::cancel), [`finish_within`](Self::finish_within) and `Drop` all join
/// the reader and close its pipe, even when a descendant keeps the write end open. Idle
/// readers check the pipe every 20 ms; cancellation also wakes that wait immediately.
///
/// This guarantee cannot be provided for arbitrary `Read`; use [`crate::OutputCapture`]
/// only when the reader's lifetime is controlled separately. Do not change the native
/// read end's blocking mode through another handle while this capture owns it.
/// Windows requires a byte pipe opened for overlapped I/O (including `OwnedCommand` pipes);
/// disconnected overlapped pipes are accepted as EOF without issuing a read.
/// Construction rejects unsupported handles before starting a reader.
pub struct NativeOutputCapture {
    shared: Arc<Shared>,
    stop: Arc<AtomicBool>,
    handle: Option<JoinHandle<()>>,
}

impl NativeOutputCapture {
    /// Own & drain a child's stdout. Errors before spawning a reader if the pipe cannot
    /// support cancellable native I/O or the operating system cannot start the reader thread.
    pub fn stdout(pipe: ChildStdout, max_bytes: usize) -> io::Result<Self> {
        Self::start(platform::stdout(pipe), max_bytes)
    }

    /// Own & drain a child's stderr, with the same guarantees as [`stdout`](Self::stdout).
    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);
                    }
                }
                // Close the read end before Done signals completion to finish_within.
                drop(pipe);
            })?;
        Ok(Self {
            shared,
            stop,
            handle: Some(handle),
        })
    }

    /// What has been captured so far. This does not stop the reader.
    pub fn snapshot(&self) -> CapturedOutput {
        self.shared
            .0
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .output
            .clone()
    }

    /// True only if the pipe reached EOF. A deadline/cancellation can return partial bytes
    /// without exceeding the retention cap; `CapturedOutput::truncated` describes cap loss.
    pub fn is_complete(&self) -> bool {
        self.shared.0.lock().unwrap_or_else(|e| e.into_inner()).eof
    }

    /// Take the native read error, if any. Cancellation is not an error.
    pub fn take_error(&self) -> Option<io::Error> {
        self.shared
            .0
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .error
            .take()
    }

    /// Wait up to `timeout` for EOF/error, then cancel & join the reader. Always closes the
    /// read end before returning; a writer retained by a descendant cannot extend the wait.
    /// Returned bytes may be partial: consult [`is_complete`](Self::is_complete).
    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()
    }

    /// Cancel immediately, join the reader & return retained bytes. Idempotent.
    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();
    }
}