rightkit_child_process/
native_capture.rs

1//! Cancellable capture of native child pipes. Generic `Read` capture remains separate.
2
3use crate::CapturedOutput;
4use std::{
5    fs::File,
6    io,
7    process::{ChildStderr, ChildStdout},
8    sync::{
9        atomic::{AtomicBool, Ordering},
10        Arc, Condvar, Mutex,
11    },
12    thread::{self, JoinHandle},
13    time::Duration,
14};
15
16#[cfg(unix)]
17#[path = "native_capture/unix.rs"]
18mod platform;
19#[cfg(windows)]
20#[path = "native_capture/windows.rs"]
21pub(crate) mod platform;
22
23const CHUNK: usize = 8 * 1024;
24const IDLE_WAIT: Duration = Duration::from_millis(20);
25
26#[derive(Default)]
27struct State {
28    output: CapturedOutput,
29    error: Option<io::Error>,
30    eof: bool,
31    finished: bool,
32}
33
34type Shared = (Mutex<State>, Condvar);
35
36/// Which of a child's pipes a capture reads. Names the stream for the test seam only.
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub(crate) enum Stream {
39    Stdout,
40    Stderr,
41}
42
43/// Test seam: make the next capture of `stream` that the calling thread starts fail its first
44/// read with `kind`, which the reader stores exactly as it stores a real read failure.
45#[cfg(test)]
46pub(crate) fn inject_read_fault(stream: Stream, kind: io::ErrorKind) {
47    test_fault::arm(stream, kind);
48}
49
50/// Receives each chunk the reader thread reads, on that thread, together with the capture's
51/// stop flag. A sink that can wait (a full channel) must give up once the flag is set, so a
52/// join never waits on it. Only [`crate::run_streaming`] installs one.
53pub(crate) type ChunkSink = Box<dyn FnMut(&[u8], &AtomicBool) + Send + 'static>;
54
55struct Done(Arc<Shared>);
56
57impl Drop for Done {
58    fn drop(&mut self) {
59        self.0 .0.lock().unwrap_or_else(|e| e.into_inner()).finished = true;
60        self.0 .1.notify_all();
61    }
62}
63
64/// An owned reader of a native child stdout/stderr pipe, with bounded byte retention.
65///
66/// Unix read ends use nonblocking I/O; Windows read ends use owned overlapped I/O.
67/// Cancellation wakes the private reader; no user callbacks run
68/// on it. [`cancel`](Self::cancel), [`finish_within`](Self::finish_within) and `Drop` all join
69/// the reader and close its pipe, even when a descendant keeps the write end open. Idle
70/// readers check the pipe every 20 ms; cancellation also wakes that wait immediately.
71///
72/// This guarantee cannot be provided for arbitrary `Read`; use [`crate::OutputCapture`]
73/// only when the reader's lifetime is controlled separately. Do not change the native
74/// read end's blocking mode through another handle while this capture owns it.
75/// Windows requires a byte pipe opened for overlapped I/O (including `OwnedCommand` pipes);
76/// disconnected overlapped pipes are accepted as EOF without issuing a read.
77/// Construction rejects unsupported handles before starting a reader.
78pub struct NativeOutputCapture {
79    shared: Arc<Shared>,
80    stop: Arc<AtomicBool>,
81    handle: Option<JoinHandle<()>>,
82}
83
84impl NativeOutputCapture {
85    /// Own & drain a child's stdout. Errors before spawning a reader if the pipe cannot
86    /// support cancellable native I/O or the operating system cannot start the reader thread.
87    pub fn stdout(pipe: ChildStdout, max_bytes: usize) -> io::Result<Self> {
88        Self::start(platform::stdout(pipe), max_bytes, None, Stream::Stdout)
89    }
90
91    /// Own & drain a child's stderr, with the same guarantees as [`stdout`](Self::stdout).
92    pub fn stderr(pipe: ChildStderr, max_bytes: usize) -> io::Result<Self> {
93        Self::start(platform::stderr(pipe), max_bytes, None, Stream::Stderr)
94    }
95
96    /// Like [`stdout`](Self::stdout), but every chunk is also handed to `sink` on the reader
97    /// thread and nothing is retained.
98    pub(crate) fn stdout_streaming(pipe: ChildStdout, sink: ChunkSink) -> io::Result<Self> {
99        Self::start(platform::stdout(pipe), 0, Some(sink), Stream::Stdout)
100    }
101
102    /// The stderr counterpart of [`stdout_streaming`](Self::stdout_streaming).
103    pub(crate) fn stderr_streaming(pipe: ChildStderr, sink: ChunkSink) -> io::Result<Self> {
104        Self::start(platform::stderr(pipe), 0, Some(sink), Stream::Stderr)
105    }
106
107    fn start(
108        pipe: File,
109        max_bytes: usize,
110        mut sink: Option<ChunkSink>,
111        stream: Stream,
112    ) -> io::Result<Self> {
113        let prepared = platform::prepare(pipe)?;
114        // Test seam: taken here, on the thread that starts the capture, so a fault armed by one
115        // test can only reach a capture that same test starts (see `inject_read_fault`).
116        #[cfg(test)]
117        let mut fault = test_fault::take(stream);
118        #[cfg(not(test))]
119        let _ = stream;
120        let shared: Arc<Shared> = Arc::default();
121        let stop = Arc::new(AtomicBool::new(false));
122        let state = Arc::clone(&shared);
123        let cancelled = Arc::clone(&stop);
124        let handle = thread::Builder::new()
125            .name("rightkit-native-output-capture".into())
126            .spawn(move || {
127                let _done = Done(Arc::clone(&state));
128                let mut pipe = platform::reader(prepared);
129                let mut buffer = [0; CHUNK];
130                while !cancelled.load(Ordering::Acquire) {
131                    #[cfg(test)]
132                    let read = match fault.take() {
133                        Some(kind) => Err(io::Error::new(kind, "injected pipe read failure")),
134                        None => platform::read(&mut pipe, &mut buffer),
135                    };
136                    #[cfg(not(test))]
137                    let read = platform::read(&mut pipe, &mut buffer);
138                    match read {
139                        Ok(0) => {
140                            state.0.lock().unwrap_or_else(|e| e.into_inner()).eof = true;
141                            break;
142                        }
143                        Ok(count) => {
144                            if let Some(sink) = sink.as_mut() {
145                                sink(&buffer[..count], &cancelled);
146                            }
147                            state
148                                .0
149                                .lock()
150                                .unwrap_or_else(|e| e.into_inner())
151                                .output
152                                .push(&buffer[..count], max_bytes);
153                        }
154                        Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
155                            thread::park_timeout(IDLE_WAIT);
156                        }
157                        Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
158                        Err(error) => {
159                            state.0.lock().unwrap_or_else(|e| e.into_inner()).error = Some(error);
160                            break;
161                        }
162                    }
163                }
164                if let Err(error) = platform::shutdown(&mut pipe) {
165                    let mut state = state.0.lock().unwrap_or_else(|e| e.into_inner());
166                    if state.error.is_none() {
167                        state.error = Some(error);
168                    }
169                }
170                // Close the read end before Done signals completion to finish_within.
171                drop(pipe);
172            })?;
173        Ok(Self {
174            shared,
175            stop,
176            handle: Some(handle),
177        })
178    }
179
180    /// What has been captured so far. This does not stop the reader.
181    pub fn snapshot(&self) -> CapturedOutput {
182        self.shared
183            .0
184            .lock()
185            .unwrap_or_else(|e| e.into_inner())
186            .output
187            .clone()
188    }
189
190    /// True only if the pipe reached EOF. A deadline/cancellation can return partial bytes
191    /// without exceeding the retention cap; `CapturedOutput::truncated` describes cap loss.
192    pub fn is_complete(&self) -> bool {
193        self.shared.0.lock().unwrap_or_else(|e| e.into_inner()).eof
194    }
195
196    /// Take the native read error, if any. Cancellation is not an error.
197    pub fn take_error(&self) -> Option<io::Error> {
198        self.shared
199            .0
200            .lock()
201            .unwrap_or_else(|e| e.into_inner())
202            .error
203            .take()
204    }
205
206    /// Wait up to `timeout` for EOF/error, then cancel & join the reader. Always closes the
207    /// read end before returning; a writer retained by a descendant cannot extend the wait.
208    /// Returned bytes may be partial: consult [`is_complete`](Self::is_complete).
209    pub fn finish_within(&mut self, timeout: Duration) -> CapturedOutput {
210        {
211            let state = self.shared.0.lock().unwrap_or_else(|e| e.into_inner());
212            drop(
213                self.shared
214                    .1
215                    .wait_timeout_while(state, timeout, |state| !state.finished)
216                    .unwrap_or_else(|e| e.into_inner()),
217            );
218        }
219        self.stop_and_join();
220        self.snapshot()
221    }
222
223    /// Cancel immediately, join the reader & return retained bytes. Idempotent.
224    pub fn cancel(&mut self) -> CapturedOutput {
225        self.stop_and_join();
226        self.snapshot()
227    }
228
229    fn stop_and_join(&mut self) {
230        self.stop.store(true, Ordering::Release);
231        if let Some(handle) = self.handle.take() {
232            handle.thread().unpark();
233            if handle.join().is_err() {
234                self.shared
235                    .0
236                    .lock()
237                    .unwrap_or_else(|e| e.into_inner())
238                    .error = Some(io::Error::other("native pipe reader panicked"));
239            }
240        }
241    }
242}
243
244impl Drop for NativeOutputCapture {
245    fn drop(&mut self) {
246        self.stop_and_join();
247    }
248}
249
250#[cfg(test)]
251pub(crate) mod test_fault {
252    use super::Stream;
253    use std::{cell::Cell, io};
254
255    thread_local! {
256        // Thread-local, not global: `start` runs on the thread that called `run`, so parallel
257        // tests cannot take each other's fault.
258        static ARMED: Cell<[Option<io::ErrorKind>; 2]> = const { Cell::new([None, None]) };
259    }
260
261    pub(crate) fn arm(stream: Stream, kind: io::ErrorKind) {
262        ARMED.with(|armed| {
263            let mut slots = armed.get();
264            slots[stream as usize] = Some(kind);
265            armed.set(slots);
266        });
267    }
268
269    pub(super) fn take(stream: Stream) -> Option<io::ErrorKind> {
270        ARMED.with(|armed| {
271            let mut slots = armed.get();
272            let kind = slots[stream as usize].take();
273            armed.set(slots);
274            kind
275        })
276    }
277}
rightkit_child_process/native_capture.rs

rightkit_child_process/
native_capture.rs
