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

rightkit_child_process/
native_capture.rs
