artifactize_core/process/
capture.rs

1//! Draining a child's output pipes into its bounded capture, its evidence files and the lines a
2//! Tool log shows. How each pipe ended and whether its evidence was saved are kept apart, so a
3//! slow disk is never taken for a process that still holds a pipe.
4
5use std::{io, path::PathBuf, time::Duration};
6
7use tokio::{
8    io::{AsyncRead, AsyncReadExt, AsyncWriteExt},
9    time::{Instant, timeout},
10};
11use tokio_util::sync::CancellationToken;
12
13use crate::{
14    platform,
15    tools::human::log::{Lines, Stream},
16};
17
18/// Drain process pipes in small bounded chunks without one read per byte.
19const CAPTURE_BUFFER_BYTES: usize = 8192;
20
21/// How long a pipe may stay open once the process group is gone. Only a process outside the
22/// group can hold it then, such as a descendant that left the group. The time counts only
23/// while the pipe is waited for, never while its evidence is written, and it is generous
24/// because waiting costs nothing unless such a process exists, while a late wakeup on a busy
25/// runtime must not pass for one.
26pub(crate) const OUTPUT_CLOSE_GRACE: Duration = Duration::from_secs(5);
27
28/// How reading one pipe ended.
29#[derive(Debug)]
30pub(super) enum PipeEnd {
31    /// Every writer closed it.
32    Closed,
33    /// Still open [`OUTPUT_CLOSE_GRACE`] after the process group was gone.
34    Open,
35    Failed(io::Error),
36}
37
38/// One stream's capture: what was read of it, how its pipe ended and whether its evidence was
39/// saved.
40#[derive(Debug)]
41pub(super) struct Captured {
42    pub output: Vec<u8>,
43    pub truncated: bool,
44    pub pipe: PipeEnd,
45    pub evidence: io::Result<()>,
46}
47
48/// Read `pipe` to its end, keeping its first `limit` bytes and writing all of it to `path` and
49/// `lines`. Once `gone` is cancelled, waiting for the pipe is limited to
50/// [`OUTPUT_CLOSE_GRACE`] in all.
51pub(super) async fn capture(
52    mut pipe: impl AsyncRead + Unpin,
53    limit: usize,
54    path: Option<PathBuf>,
55    mut lines: Option<Lines>,
56    gone: &CancellationToken,
57) -> Captured {
58    let mut evidence = Evidence::open(path);
59    let mut output = Vec::new();
60    let mut buffer = [0; CAPTURE_BUFFER_BYTES];
61    let mut truncated = false;
62    let mut grace = OUTPUT_CLOSE_GRACE;
63    let pipe_end = loop {
64        let count = match read(&mut pipe, &mut buffer, gone, &mut grace).await {
65            Some(Ok(0)) => break PipeEnd::Closed,
66            Some(Ok(count)) => count,
67            Some(Err(error)) => break PipeEnd::Failed(error),
68            None => break PipeEnd::Open,
69        };
70        let chunk = &buffer[..count];
71        evidence.write(chunk).await;
72        if let Some(lines) = &mut lines {
73            lines.push(chunk);
74        }
75        let retained = count.min(limit - output.len());
76        output.extend_from_slice(&chunk[..retained]);
77        truncated |= retained < count;
78    };
79    if let Some(lines) = &mut lines {
80        lines.finish();
81    }
82    let evidence = evidence.finish(matches!(pipe_end, PipeEnd::Closed)).await;
83    Captured {
84        output,
85        truncated,
86        pipe: pipe_end,
87        evidence,
88    }
89}
90
91/// One read of `pipe`: unbounded while the process group runs, then within what is left of
92/// `grace`, which the time spent waiting here uses up. `None` when it ran out.
93async fn read(
94    pipe: &mut (impl AsyncRead + Unpin),
95    buffer: &mut [u8],
96    gone: &CancellationToken,
97    grace: &mut Duration,
98) -> Option<io::Result<usize>> {
99    if !gone.is_cancelled() {
100        // A pipe read is cancel safe: a read dropped here has taken no bytes.
101        tokio::select! {
102            biased;
103            read = pipe.read(buffer) => return Some(read),
104            () = gone.cancelled() => {}
105        }
106    }
107    let started = Instant::now();
108    let read = timeout(*grace, pipe.read(buffer)).await;
109    *grace = grace.saturating_sub(started.elapsed());
110    read.ok()
111}
112
113/// A stream's evidence file. A failed write ends the evidence, not the draining, so the pipe is
114/// still read to its end and the failure is reported once it is.
115struct Evidence {
116    file: Option<(PathBuf, tokio::fs::File)>,
117    failure: Option<io::Error>,
118}
119
120impl Evidence {
121    fn open(path: Option<PathBuf>) -> Self {
122        let Some(path) = path else {
123            return Self {
124                file: None,
125                failure: None,
126            };
127        };
128        let mut options = platform::private_options();
129        options.write(true).create_new(true);
130        match options.open(&path) {
131            Ok(file) => Self {
132                file: Some((path, tokio::fs::File::from_std(file))),
133                failure: None,
134            },
135            Err(error) => Self {
136                file: None,
137                failure: Some(error),
138            },
139        }
140    }
141
142    async fn write(&mut self, chunk: &[u8]) {
143        let Some((_, file)) = &mut self.file else {
144            return;
145        };
146        #[cfg(test)]
147        slow_disk::wait().await;
148        // Flushed at once, so captured evidence is available while the process still runs.
149        let written = async {
150            file.write_all(chunk).await?;
151            file.flush().await
152        };
153        if let Err(error) = written.await {
154            self.file = None;
155            self.failure = Some(error);
156        }
157    }
158
159    /// Flush the file and, when the pipe was `closed`, mark the stream complete beside it.
160    async fn finish(self, closed: bool) -> io::Result<()> {
161        if let Some(failure) = self.failure {
162            return Err(failure);
163        }
164        let Some((path, mut file)) = self.file else {
165            return Ok(());
166        };
167        #[cfg(test)]
168        slow_disk::wait().await;
169        file.flush().await?;
170        if closed {
171            let mut options = platform::private_options();
172            options.write(true).create_new(true);
173            let extension = path.extension().expect("stream extension");
174            options
175                .open(path.with_extension(format!("{}.complete", extension.to_string_lossy())))?;
176        }
177        Ok(())
178    }
179}
180
181/// What failed in capturing one stream.
182#[derive(Debug)]
183pub enum StreamFailure {
184    /// Its pipe was still open [`OUTPUT_CLOSE_GRACE`] after the process group was gone.
185    Open,
186    /// Its pipe could not be read.
187    Read(io::Error),
188    /// Its evidence file could not be written.
189    Evidence(io::Error),
190}
191
192/// A capture failure with all the run had come to: how the process ended and the output read.
193#[derive(Debug, thiserror::Error)]
194#[error("process output could not be captured: {}; {}", failure_text(&self.failures), end_text(&self.output.status))]
195pub struct CaptureError {
196    pub output: super::Output,
197    /// Each failure with its stream, never empty.
198    pub failures: Vec<(Stream, StreamFailure)>,
199}
200
201fn failure_text(failures: &[(Stream, StreamFailure)]) -> String {
202    failures
203        .iter()
204        .map(|(stream, failure)| {
205            let stream = match stream {
206                Stream::Stdout => "stdout",
207                Stream::Stderr => "stderr",
208            };
209            match failure {
210                StreamFailure::Open => format!(
211                    "{stream} was still open {}s after the process group was gone, so a process \
212                     outside the group holds it",
213                    OUTPUT_CLOSE_GRACE.as_secs()
214                ),
215                StreamFailure::Read(error) => format!("{stream} could not be read: {error}"),
216                StreamFailure::Evidence(error) => {
217                    format!("{stream} could not be saved: {error}")
218                }
219            }
220        })
221        .collect::<Vec<_>>()
222        .join("; ")
223}
224
225fn end_text(status: &std::process::ExitStatus) -> String {
226    format!(
227        "the process had ended ({})",
228        platform::exit_description(status)
229    )
230}
231
232impl Captured {
233    /// The output read, and what failed.
234    fn split(self) -> ((Vec<u8>, bool), impl Iterator<Item = StreamFailure>) {
235        let pipe = match self.pipe {
236            PipeEnd::Closed => None,
237            PipeEnd::Open => Some(StreamFailure::Open),
238            PipeEnd::Failed(error) => Some(StreamFailure::Read(error)),
239        };
240        let evidence = self.evidence.err().map(StreamFailure::Evidence);
241        (
242            (self.output, self.truncated),
243            pipe.into_iter().chain(evidence),
244        )
245    }
246}
247
248/// Combine how the process ended with its two streams' captures.
249pub(super) fn finish(
250    status: std::process::ExitStatus,
251    stdout: Captured,
252    stderr: Captured,
253    duration: Duration,
254) -> Result<super::Output, super::Error> {
255    let ((stdout, stdout_truncated), stdout_failures) = stdout.split();
256    let ((stderr, stderr_truncated), stderr_failures) = stderr.split();
257    let failures: Vec<_> = stdout_failures
258        .map(|failure| (Stream::Stdout, failure))
259        .chain(stderr_failures.map(|failure| (Stream::Stderr, failure)))
260        .collect();
261    let output = super::Output {
262        status,
263        stdout,
264        stderr,
265        truncated: stdout_truncated || stderr_truncated,
266        duration,
267    };
268    if failures.is_empty() {
269        Ok(output)
270    } else {
271        Err(super::Error::Capture(Box::new(CaptureError {
272            output,
273            failures,
274        })))
275    }
276}
277
278/// Tests stand in for a slow disk by holding every evidence write on their thread until
279/// released.
280#[cfg(test)]
281pub(crate) mod slow_disk {
282    use std::{cell::RefCell, sync::Arc};
283
284    use tokio::sync::watch;
285
286    /// The release switch, and how many evidence writes wait for it.
287    type Hold = (watch::Receiver<bool>, Arc<()>);
288
289    thread_local! {
290        static HOLD: RefCell<Option<Hold>> = const { RefCell::new(None) };
291    }
292
293    /// A disk that holds every evidence write on this thread until it is released.
294    pub(crate) struct SlowDisk {
295        release: watch::Sender<bool>,
296        waiting: Arc<()>,
297    }
298
299    impl SlowDisk {
300        /// How many evidence writes wait for the release now.
301        pub(crate) fn waiting(&self) -> usize {
302            // Each waiting write holds one clone; the disk and the thread's hold hold two.
303            Arc::strong_count(&self.waiting) - 2
304        }
305
306        pub(crate) fn release(&self) {
307            self.release.send_replace(true);
308        }
309    }
310
311    pub(crate) fn hold() -> SlowDisk {
312        let (release, held) = watch::channel(false);
313        let waiting = Arc::new(());
314        HOLD.with(|hold| *hold.borrow_mut() = Some((held, waiting.clone())));
315        SlowDisk { release, waiting }
316    }
317
318    pub(super) async fn wait() {
319        let held = HOLD.with(|hold| hold.borrow().clone());
320        if let Some((mut held, waiting)) = held {
321            let _ = held.wait_for(|released| *released).await;
322            drop(waiting);
323        }
324    }
325}
artifactize_core/process/capture.rs

artifactize_core/process/
capture.rs
