Skip to main content

stuffr_cli/bench/
memory.rs

1//! Peak memory for `stuffr bench`, measured in untimed child processes.
2//!
3//! After a row's timed runs, the parent runs the row once more, once for
4//! compression and once for decompression, each in a child: bench's own
5//! executable re-run as `stuffr __bench-memory-child <job>` (see
6//! [`ChildJob::to_args`]). `getrusage`'s `ru_maxrss` is a per-process
7//! high-water mark, which is why each measurement needs a fresh process,
8//! and why one per file in `--per-file` mode (the row reports the largest).
9//!
10//! **What a child measures.** Its 64 KiB read buffer, its 64 KiB output
11//! buffer and its open files (the compress child's output file included)
12//! are made first; then it samples its peak resident set as the baseline (the
13//! high-water so far, which is its current size, since nothing has run
14//! yet), runs the codec to the end, and reports the peak minus that
15//! baseline, saturating at 0. So the number is the codec's resident working
16//! memory plus a few small fixed buffers; memory reserved but never touched
17//! does not count. The compress child streams the input from the files
18//! (each read stops at the length snapshotted at load, as in the timed
19//! runs), so an encoder that buffers its whole input (the pure-tier zstd
20//! encoder) shows memory growing with the input: that is what using it
21//! costs. The compress child writes its compressed output to a guarded
22//! temp file the parent created and armed, `stuffr-bench-<pid>-<n>`; the
23//! decompress child then reads that file through the same 64 KiB buffered
24//! source the timed runs use. So the row is compressed once more, not
25//! twice (ruling F8), and a compress child that fails leaves nothing for
26//! the decompress side, which then shows `-` too.
27//!
28//! **Same conditions as the timed runs**: level, `--threads` (through
29//! [`governor_for`]) and `--memory-limit`. The backend is not sent: the
30//! child is the same executable, so it has the same one.
31//!
32//! **The wire.** The child prints one line on stdout, `ok <bytes>` or
33//! `err <text>`, and exits 0 or 1 to match. Anything else (a crash, a
34//! signal, no line) is a failed measurement. A failure never fails the row:
35//! its cell is `-` and the text goes to [`Memory::error`]. A refusal by
36//! `--memory-limit` (an `Error::ResourceLimit`) reads `refused by
37//! --memory-limit: …`, so it is told apart from a crash.
38//!
39//! A debug build's child refuses, as `bench` does, so this path is not
40//! reachable from the CLI tests; [`Runner`] lets the parent's logic be
41//! tested with a fake, and [`child_main`] is tested in-process.
42
43use std::ffi::{OsStr, OsString};
44use std::fs::File;
45use std::io::{self, BufWriter, Write};
46use std::path::{Path, PathBuf};
47use std::process::{Command, ExitCode, Stdio};
48use std::sync::atomic::{AtomicI32, Ordering};
49
50use stuffr::source::{GuardedSeek, Source};
51use stuffr::{Codec, DecodeOpts, EncodeOpts, Error, FormatKind};
52
53use super::measure::{
54    CHUNK, Mode, Plan, Snapshot, TempOut, fill, governor_for, read_some, shrank, temp_source,
55};
56use super::score::Memory;
57
58/// The hidden first argument that makes the binary a memory child. Not a
59/// clap subcommand: `main` checks for it before clap parses, so it is in no
60/// help text and the examples page's guard (which walks clap) never sees it.
61pub const CHILD_ARG: &str = "__bench-memory-child";
62
63/// What a debug build's child answers: its numbers would be meaningless.
64const DEBUG_CHILD_REFUSAL: &str = "memory runs from a debug build are meaningless";
65
66/// The pid of the memory child the parent is waiting for, 0 when none. Set
67/// once the child is spawned, cleared once it has been waited for; the
68/// signal cleanup reads it to kill the child before bench dies.
69pub static RUNNING_CHILD: AtomicI32 = AtomicI32::new(0);
70
71/// Which half of the round trip a child measures.
72#[derive(Clone, Copy, Debug, PartialEq, Eq)]
73pub enum Side {
74    Compress,
75    Decompress,
76}
77
78impl Side {
79    fn word(self) -> &'static str {
80        match self {
81            Side::Compress => "compress",
82            Side::Decompress => "decompress",
83        }
84    }
85}
86
87/// What one child runs. Compress: `files` concatenated as one stream,
88/// written to `out` (discarded when `None`). Decompress: exactly one file,
89/// the compressed copy; `out` is `None`.
90#[derive(Clone, Debug, PartialEq, Eq)]
91pub struct ChildJob {
92    pub side: Side,
93    pub codec: &'static str,
94    pub level: Option<i32>,
95    pub threads: Option<usize>,
96    pub memory_limit: u64,
97    pub out: Option<PathBuf>,
98    pub files: Vec<Snapshot>,
99}
100
101/// `-` for an absent level or thread count on the wire.
102const NONE: &str = "-";
103
104impl ChildJob {
105    /// The arguments after [`CHILD_ARG`]:
106    /// `<side> <codec> <level|-> <threads|-> <memory-limit> <out|-> [<len> <path>]…`,
107    /// `side` being `compress` or `decompress`, every number in decimal.
108    /// The decompress side ignores the level (the parent sends `-`). `out`
109    /// is `-` for none; the parent's is always an absolute temp path.
110    /// Paths go as they are (not through UTF-8), each its own argument.
111    pub fn to_args(&self) -> Vec<OsString> {
112        let opt = |v: Option<String>| OsString::from(v.unwrap_or_else(|| NONE.to_string()));
113        let mut args = vec![
114            OsString::from(self.side.word()),
115            OsString::from(self.codec),
116            opt(self.level.map(|l| l.to_string())),
117            opt(self.threads.map(|t| t.to_string())),
118            OsString::from(self.memory_limit.to_string()),
119            self.out
120                .clone()
121                .map_or_else(|| OsString::from(NONE), PathBuf::into_os_string),
122        ];
123        for s in &self.files {
124            args.push(OsString::from(s.len.to_string()));
125            args.push(s.path.clone().into_os_string());
126        }
127        args
128    }
129
130    /// Parses [`ChildJob::to_args`]'s form back. The codec must be one of
131    /// this build's.
132    pub fn from_args(args: &[OsString]) -> Result<Self, String> {
133        fn text<'a>(a: Option<&'a OsString>, what: &str) -> Result<&'a str, String> {
134            a.ok_or_else(|| format!("missing {what}"))?
135                .to_str()
136                .ok_or_else(|| format!("{what} is not UTF-8"))
137        }
138        fn num<T: std::str::FromStr>(s: &str, what: &str) -> Result<T, String> {
139            s.parse().map_err(|_| format!("bad {what} `{s}`"))
140        }
141        fn opt<T: std::str::FromStr>(s: &str, what: &str) -> Result<Option<T>, String> {
142            if s == NONE {
143                Ok(None)
144            } else {
145                num(s, what).map(Some)
146            }
147        }
148        let mut it = args.iter();
149        let side = match text(it.next(), "side")? {
150            "compress" => Side::Compress,
151            "decompress" => Side::Decompress,
152            other => return Err(format!("bad side `{other}`")),
153        };
154        let codec = codec_id(text(it.next(), "codec")?)?.as_str();
155        let level = opt(text(it.next(), "level")?, "level")?;
156        let threads = opt(text(it.next(), "threads")?, "threads")?;
157        let memory_limit = num(text(it.next(), "memory limit")?, "memory limit")?;
158        let out = it.next().ok_or("missing output")?;
159        let out = (out != NONE).then(|| PathBuf::from(out));
160        let mut files = Vec::new();
161        while let Some(len) = it.next() {
162            let len = num(text(Some(len), "length")?, "length")?;
163            let path = it.next().ok_or("a length without its path")?;
164            files.push(Snapshot {
165                path: path.into(),
166                len,
167            });
168        }
169        Ok(Self {
170            side,
171            codec,
172            level,
173            threads,
174            memory_limit,
175            out,
176            files,
177        })
178    }
179
180    /// The plan the timed runs used, as far as a child needs it.
181    fn plan(&self) -> Plan {
182        Plan {
183            mode: Mode::OneStream,
184            repeat: 1,
185            threads: self.threads,
186            memory_limit: self.memory_limit,
187        }
188    }
189}
190
191/// `name` as this build's codec id.
192fn codec_id(name: &str) -> Result<stuffr::FormatId, String> {
193    stuffr::registry()
194        .matrix()
195        .into_iter()
196        .find(|r| r.kind == FormatKind::Codec && r.id.as_str() == name)
197        .map(|r| r.id)
198        .ok_or_else(|| format!("`{name}` is not a codec in this build"))
199}
200
201/// A child's stdout line: `Some(Ok(n))` for `ok <n>`, `Some(Err(text))` for
202/// `err <text>` (text not empty), `None` for anything else.
203pub fn parse_reply(stdout: &[u8]) -> Option<Result<u64, String>> {
204    let line = std::str::from_utf8(stdout).ok()?.lines().next()?;
205    if let Some(n) = line.strip_prefix("ok ") {
206        if !n.bytes().all(|b| b.is_ascii_digit()) {
207            return None;
208        }
209        return n.parse().ok().map(Ok);
210    }
211    let text = line.strip_prefix("err ")?.trim();
212    (!text.is_empty()).then(|| Err(text.to_string()))
213}
214
215/// A child's result as its one stdout line (no newline inside).
216fn reply_line(r: &Result<u64, String>) -> String {
217    match r {
218        Ok(n) => format!("ok {n}"),
219        Err(e) => format!("err {}", e.replace(['\n', '\r'], " ")),
220    }
221}
222
223/// How a measurement failure reads in `mem_error`: a refusal by
224/// `--memory-limit` (`ResourceLimit`, which is also what an allocation
225/// failure becomes) says so; everything else is its own message.
226pub fn describe(e: &Error) -> String {
227    match e {
228        Error::ResourceLimit(s) => format!("refused by --memory-limit: {s}"),
229        other => other.to_string(),
230    }
231}
232
233/// A finished child: its reply when it gave a well-formed one and its exit
234/// status agrees (0 for `ok`, failure for `err`); otherwise a failure
235/// naming `status` and the first line of its stderr, if any.
236pub fn interpret(success: bool, status: &str, stdout: &[u8], stderr: &[u8]) -> Result<u64, String> {
237    match (parse_reply(stdout), success) {
238        (Some(Ok(n)), true) => Ok(n),
239        (Some(Err(e)), false) => Err(e),
240        _ => {
241            let why = String::from_utf8_lossy(stderr);
242            match why.lines().map(str::trim).find(|l| !l.is_empty()) {
243                Some(line) => Err(format!("the memory run failed ({status}): {line}")),
244                None => Err(format!("the memory run failed ({status})")),
245            }
246        }
247    }
248}
249
250/// The hidden entry `main` calls first: `None` unless `args[1]` is
251/// [`CHILD_ARG`]; otherwise runs the child, prints its reply line and
252/// returns its exit code. A debug build refuses.
253pub fn child_entry(args: &[OsString]) -> Option<ExitCode> {
254    if args.get(1).map(OsString::as_os_str) != Some(OsStr::new(CHILD_ARG)) {
255        return None;
256    }
257    // A malformed job is a usage error (exit 2): exit 1 is reserved for
258    // stuffr itself failing. The parent takes any non-zero exit with `err`.
259    let (result, code) = if cfg!(debug_assertions) {
260        (Err(DEBUG_CHILD_REFUSAL.to_string()), ExitCode::FAILURE)
261    } else {
262        match ChildJob::from_args(&args[2..]) {
263            Ok(job) => (child_main(job), ExitCode::FAILURE),
264            Err(e) => (Err(e), ExitCode::from(2)),
265        }
266    };
267    let mut out = io::stdout().lock();
268    let printed = writeln!(out, "{}", reply_line(&result)).and_then(|()| out.flush());
269    Some(if result.is_ok() && printed.is_ok() {
270        ExitCode::SUCCESS
271    } else if result.is_err() {
272        code
273    } else {
274        ExitCode::FAILURE
275    })
276}
277
278/// The child: runs `job` and returns its peak resident memory above the
279/// baseline (see the module doc), in bytes.
280pub fn child_main(job: ChildJob) -> Result<u64, String> {
281    let reg = stuffr::registry();
282    let id = codec_id(job.codec)?;
283    match job.side {
284        Side::Compress => {
285            let codec = reg.require_encoder(id).map_err(|e| describe(&e))?;
286            let encode = EncodeOpts {
287                level: job.level,
288                governor: governor_for(&job.plan()),
289                ..Default::default()
290            };
291            codec.check_encode_opts(&encode).map_err(|e| describe(&e))?;
292            compress_child(codec.as_ref(), &encode, &job.files, job.out.as_deref())
293                .map_err(|e| describe(&e))
294        }
295        Side::Decompress => {
296            let codec = reg.require_decoder(id).map_err(|e| describe(&e))?;
297            let decode = DecodeOpts {
298                memory_limit: Some(job.memory_limit),
299                ..Default::default()
300            };
301            let [copy] = &job.files[..] else {
302                return Err(format!(
303                    "a decompress run takes one file, not {}",
304                    job.files.len()
305                ));
306            };
307            decompress_child(codec.as_ref(), &decode, copy).map_err(|e| describe(&e))
308        }
309    }
310}
311
312/// Streams `files` through the encoder into `out` (the parent's temp file,
313/// which it already created: opened here, truncated), or into a discarding
314/// sink when there is none.
315fn compress_child(
316    codec: &dyn Codec,
317    encode: &EncodeOpts,
318    files: &[Snapshot],
319    out: Option<&Path>,
320) -> stuffr::Result<u64> {
321    let mut feeds = files
322        .iter()
323        .map(|s| s.open().map(|f| (s, f)))
324        .collect::<io::Result<Vec<_>>>()?;
325    let mut buf = vec![0u8; CHUNK];
326    let dst: Box<dyn Write + Send> = match out {
327        Some(path) => Box::new(File::options().write(true).truncate(true).open(path)?),
328        None => Box::new(io::sink()),
329    };
330    let dst = BufWriter::with_capacity(CHUNK, dst);
331    let baseline = baseline_rss()?;
332    let mut sink = codec.encoder(Box::new(dst), encode)?;
333    for (snap, f) in &mut feeds {
334        feed(snap, f, &mut buf, &mut *sink)?;
335    }
336    sink.finish()?;
337    Ok(peak_rss()?.saturating_sub(baseline))
338}
339
340/// Reads `copy` through the decoder to its end, discarding the output.
341fn decompress_child(
342    codec: &dyn Codec,
343    decode: &DecodeOpts,
344    copy: &Snapshot,
345) -> stuffr::Result<u64> {
346    let src: Box<dyn Source> = Box::new(temp_source(
347        GuardedSeek::new(File::open(&copy.path)?),
348        copy.len,
349    ));
350    let mut buf = vec![0u8; CHUNK];
351    let baseline = baseline_rss()?;
352    let mut decoder = codec.decoder(src, decode)?;
353    while fill(&mut decoder, &mut buf).map_err(Error::from_decode_io)? == buf.len() {}
354    drop(decoder);
355    Ok(peak_rss()?.saturating_sub(baseline))
356}
357
358/// Writes all of `f` (`snap` opened) to `sink` through `buf`; a file that
359/// ends short of its loaded length is an error, as in the timed runs.
360fn feed(
361    snap: &Snapshot,
362    f: &mut io::Take<File>,
363    buf: &mut [u8],
364    sink: &mut dyn Write,
365) -> stuffr::Result<()> {
366    let mut read = 0u64;
367    loop {
368        let n = read_some(f, buf)?;
369        if n == 0 {
370            break;
371        }
372        read += n as u64;
373        sink.write_all(&buf[..n])?;
374    }
375    if read < snap.len {
376        return Err(shrank(snap, read));
377    }
378    Ok(())
379}
380
381/// Whether this platform's memory figures have been verified: macOS
382/// (`ru_maxrss`) and Linux (`/proc/self/status`). Elsewhere (the BSDs, and
383/// off unix) bench reports memory as not measured (ruling F9).
384pub const MEASURES: bool = cfg!(any(target_vendor = "apple", target_os = "linux"));
385
386/// The resident size just before the codec starts, in bytes.
387///
388/// macOS: `ru_maxrss`, the high-water so far, which is the current size
389/// since nothing has run yet. Linux: `VmRSS` from `/proc/self/status`, not
390/// `ru_maxrss`: exec copies the spawning process's high-water mark into the
391/// new image's `ru_maxrss` (fs/exec.c, `setmax_mm_hiwater_rss`), so there it
392/// starts at bench's own lifetime peak and would swallow the codec's.
393fn baseline_rss() -> io::Result<u64> {
394    #[cfg(target_os = "linux")]
395    {
396        proc_status_bytes(b"VmRSS:")
397    }
398    #[cfg(not(target_os = "linux"))]
399    {
400        peak_rss()
401    }
402}
403
404/// The address space's peak resident size so far, in bytes. macOS:
405/// `ru_maxrss` (bytes there). Linux: `VmHWM` from `/proc/self/status`, the
406/// high-water of this address space only, fresh at exec.
407#[cfg(target_vendor = "apple")]
408fn peak_rss() -> io::Result<u64> {
409    // SAFETY: `getrusage` only writes the `rusage` it is given, which is
410    // plain data and valid zeroed.
411    let mut usage: libc::rusage = unsafe { std::mem::zeroed() };
412    // SAFETY: as above; `usage` is a valid, writable `rusage`.
413    if unsafe { libc::getrusage(libc::RUSAGE_SELF, &mut usage) } != 0 {
414        return Err(io::Error::last_os_error());
415    }
416    Ok(u64::try_from(usage.ru_maxrss).unwrap_or(0))
417}
418
419#[cfg(target_os = "linux")]
420fn peak_rss() -> io::Result<u64> {
421    proc_status_bytes(b"VmHWM:")
422}
423
424/// Elsewhere memory is not measured; bench never starts a child there.
425#[cfg(not(any(target_vendor = "apple", target_os = "linux")))]
426fn peak_rss() -> io::Result<u64> {
427    Err(io::Error::new(
428        io::ErrorKind::Unsupported,
429        "memory is not measured on this platform",
430    ))
431}
432
433/// `field` (with its colon, e.g. `VmHWM:`) of `/proc/self/status`, given
434/// there in kB, as bytes. Read into a buffer on the stack, so sampling
435/// allocates nothing on the heap that the codec's figure could include.
436#[cfg(target_os = "linux")]
437fn proc_status_bytes(field: &[u8]) -> io::Result<u64> {
438    use std::io::Read;
439    let mut buf = [0u8; 8192];
440    let mut f = File::open("/proc/self/status")?;
441    let mut len = 0;
442    while len < buf.len() {
443        match f.read(&mut buf[len..]) {
444            Ok(0) => break,
445            Ok(n) => len += n,
446            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
447            Err(e) => return Err(e),
448        }
449    }
450    status_field_kib(&buf[..len], field)
451        .map(|kib| kib.saturating_mul(1024))
452        .ok_or_else(|| io::Error::other("no memory figure in /proc/self/status"))
453}
454
455/// The number on the line starting with `field` in a `/proc/<pid>/status`
456/// text (`VmHWM:\t   1234 kB`), in kB.
457#[cfg(any(target_os = "linux", test))]
458fn status_field_kib(text: &[u8], field: &[u8]) -> Option<u64> {
459    let line = text.split(|&b| b == b'\n').find(|l| l.starts_with(field))?;
460    let rest = line[field.len()..].trim_ascii_start();
461    let digits = rest.iter().take_while(|b| b.is_ascii_digit()).count();
462    if digits == 0 {
463        return None;
464    }
465    rest[..digits].iter().try_fold(0u64, |n, &b| {
466        n.checked_mul(10)?.checked_add(u64::from(b - b'0'))
467    })
468}
469
470/// Runs one child job and returns its reply. The real one spawns a
471/// process ([`ProcessRunner`]); tests inject a fake.
472pub trait Runner {
473    fn run(&mut self, job: &ChildJob) -> Result<u64, String>;
474}
475
476/// Runs each job as `<this executable> __bench-memory-child …`, with stdin
477/// closed and stdout and stderr captured, keeping its pid in
478/// [`RUNNING_CHILD`] while it runs. The executable is resolved once, when
479/// the runner is made at the start of the run, and reused for every child.
480pub struct ProcessRunner {
481    exe: Result<PathBuf, String>,
482}
483
484impl Default for ProcessRunner {
485    fn default() -> Self {
486        Self::new()
487    }
488}
489
490impl ProcessRunner {
491    pub fn new() -> Self {
492        Self {
493            exe: std::env::current_exe()
494                .map_err(|e| format!("cannot find bench's own executable: {e}")),
495        }
496    }
497}
498
499impl Runner for ProcessRunner {
500    fn run(&mut self, job: &ChildJob) -> Result<u64, String> {
501        let exe = self.exe.clone()?;
502        // Bench's signals are held back from the spawn until the pid is
503        // stored, so a signal sent to bench alone in between is delivered
504        // with the pid known, and the cleanup kills the child (no orphan).
505        // The child does not inherit the mask: std resets it before exec.
506        let child = crate::compat::with_run_signals_held(|| {
507            let child = Command::new(exe)
508                .arg(CHILD_ARG)
509                .args(job.to_args())
510                .stdin(Stdio::null())
511                .stdout(Stdio::piped())
512                .stderr(Stdio::piped())
513                .spawn()?;
514            RUNNING_CHILD.store(i32::try_from(child.id()).unwrap_or(0), Ordering::SeqCst);
515            Ok::<_, io::Error>(child)
516        })
517        .map_err(|e| format!("cannot start the memory run: {e}"))?;
518        // Known window, accepted (ruling F7): `wait_with_output` reaps the
519        // child before the clear below, so a signal arriving in between
520        // makes the cleanup `kill()` a pid that is no longer ours. It lasts
521        // microseconds and pids are handed out sequentially, so reuse
522        // inside it is theoretical. Clearing before the wait instead would
523        // leave the child running (orphaned) on a bare SIGTERM to bench.
524        let out = child.wait_with_output();
525        RUNNING_CHILD.store(0, Ordering::SeqCst);
526        let out = out.map_err(|e| format!("the memory run failed: {e}"))?;
527        interpret(
528            out.status.success(),
529            &out.status.to_string(),
530            &out.stdout,
531            &out.stderr,
532        )
533    }
534}
535
536/// [`measure_memory_with`] with real child processes (`runner`). Only
537/// where [`MEASURES`]: elsewhere the caller reports memory as not measured.
538pub fn measure_memory(
539    runner: &mut ProcessRunner,
540    codec: &'static str,
541    level: Option<i32>,
542    plan: &Plan,
543    input_files: &[Snapshot],
544) -> Memory {
545    measure_memory_with(runner, codec, level, plan, input_files)
546}
547
548/// One row's memory under `plan`: a compress child and a decompress child
549/// per stream (`plan.mode`: one stream over every file, or one per file),
550/// reporting each side's largest peak. Per stream the parent creates an
551/// armed temp file, the compress child writes the compressed stream into
552/// it, and the decompress child reads it back; it is removed before the
553/// next stream begins. Strictly after the row's timed runs.
554///
555/// A side whose child fails for any stream is `None`, and the first
556/// failure's text, prefixed by its side, is `error`. A failed decompress
557/// child leaves the compress side standing; a failed compress child (or a
558/// temp file that cannot be made) leaves no copy, so both sides are `None`.
559/// Never an error of its own.
560pub fn measure_memory_with(
561    runner: &mut dyn Runner,
562    codec: &'static str,
563    level: Option<i32>,
564    plan: &Plan,
565    input_files: &[Snapshot],
566) -> Memory {
567    // bench always has at least one input; with none there is nothing to
568    // measure, and a zero would wrongly win "least memory".
569    if input_files.is_empty() {
570        return Memory {
571            error: Some("no input to measure".into()),
572            ..Memory::default()
573        };
574    }
575    let streams: Vec<&[Snapshot]> = match plan.mode {
576        Mode::OneStream => vec![input_files],
577        Mode::PerFile => input_files.chunks(1).collect(),
578    };
579    // The decompress side has no level: it is sent as `-`.
580    let job = |side, out, files| ChildJob {
581        side,
582        codec,
583        level: if side == Side::Compress { level } else { None },
584        threads: plan.threads,
585        memory_limit: plan.memory_limit,
586        out,
587        files,
588    };
589    let mut compress = Some(0u64);
590    let mut decompress = Some(0u64);
591    let mut error = None;
592    let mut fail = |side: Side, e: String, cell: &mut Option<u64>| {
593        *cell = None;
594        error.get_or_insert_with(|| format!("{}: {e}", side.word()));
595    };
596    for files in streams {
597        let Some(peak) = compress else { break };
598        // The copy lives until this stream's decompress child is done; the
599        // compress child opens it by name, so the parent's handle closes.
600        let copy = match TempOut::create() {
601            Ok((copy, _file)) => copy,
602            Err(e) => {
603                fail(
604                    Side::Compress,
605                    format!("cannot make the temp file: {e}"),
606                    &mut compress,
607                );
608                decompress = None;
609                break;
610            }
611        };
612        let out = Some(copy.path().to_path_buf());
613        match runner.run(&job(Side::Compress, out, files.to_vec())) {
614            Ok(n) => compress = Some(peak.max(n)),
615            Err(e) => {
616                fail(
617                    Side::Compress,
618                    format!("{e} (so no decompress run)"),
619                    &mut compress,
620                );
621                decompress = None;
622                break;
623            }
624        }
625        if let Some(peak) = decompress {
626            let r = std::fs::metadata(copy.path())
627                .map_err(|e| format!("cannot read the compressed copy: {e}"))
628                .and_then(|m| {
629                    let snap = Snapshot {
630                        path: copy.path().to_path_buf(),
631                        len: m.len(),
632                    };
633                    runner.run(&job(Side::Decompress, None, vec![snap]))
634                });
635            match r {
636                Ok(n) => decompress = Some(peak.max(n)),
637                Err(e) => fail(Side::Decompress, e, &mut decompress),
638            }
639        }
640    }
641    Memory {
642        compress,
643        decompress,
644        error,
645    }
646}
647
648#[cfg(test)]
649mod tests {
650    use super::*;
651    use std::fs;
652    use std::path::PathBuf;
653    use std::sync::atomic::{AtomicU32, Ordering};
654
655    /// A scratch directory `stuffr-benchmem-<pid>-<n>`, removed on drop.
656    struct Scratch(PathBuf);
657
658    impl Scratch {
659        fn new() -> Self {
660            static N: AtomicU32 = AtomicU32::new(0);
661            let n = N.fetch_add(1, Ordering::Relaxed);
662            let dir =
663                std::env::temp_dir().join(format!("stuffr-benchmem-{}-{n}", std::process::id()));
664            fs::create_dir_all(&dir).unwrap();
665            Scratch(dir)
666        }
667
668        fn file(&self, name: &str, bytes: &[u8]) -> Snapshot {
669            let path = self.0.join(name);
670            fs::write(&path, bytes).unwrap();
671            Snapshot {
672                path,
673                len: bytes.len() as u64,
674            }
675        }
676    }
677
678    impl Drop for Scratch {
679        fn drop(&mut self) {
680            let _ = fs::remove_dir_all(&self.0);
681        }
682    }
683
684    /// Serialises the tests that arm the process-wide signal cleanup (the
685    /// parent's temp copies are armed).
686    fn armed_lock() -> std::sync::MutexGuard<'static, ()> {
687        crate::compat::ARM_TEST_LOCK
688            .lock()
689            .unwrap_or_else(|e| e.into_inner())
690    }
691
692    fn text(len: usize, seed: u8) -> Vec<u8> {
693        b"the quick brown fox jumps over the lazy dog, "
694            .iter()
695            .cycle()
696            .enumerate()
697            .map(|(i, &b)| b.wrapping_add(seed).wrapping_add((i / 997) as u8))
698            .take(len)
699            .collect()
700    }
701
702    fn plan(mode: Mode) -> Plan {
703        Plan {
704            mode,
705            repeat: 1,
706            threads: None,
707            memory_limit: 1 << 30,
708        }
709    }
710
711    fn bench_temps() -> Vec<String> {
712        let prefix = format!("stuffr-bench-{}-", std::process::id());
713        fs::read_dir(std::env::temp_dir())
714            .unwrap()
715            .filter_map(|e| e.ok()?.file_name().into_string().ok())
716            .filter(|n| n.starts_with(&prefix))
717            .collect()
718    }
719
720    /// A runner that answers from `answer` and records every job it saw,
721    /// with whether each job's files (and output) existed while it ran. A
722    /// compress job that succeeds writes a stand-in copy to its output, as
723    /// the real child writes the compressed stream there.
724    struct Fake<F: FnMut(&ChildJob) -> Result<u64, String>> {
725        answer: F,
726        seen: Vec<(ChildJob, bool)>,
727    }
728
729    impl<F: FnMut(&ChildJob) -> Result<u64, String>> Runner for Fake<F> {
730        fn run(&mut self, job: &ChildJob) -> Result<u64, String> {
731            let present = job.files.iter().all(|s| s.path.exists())
732                && job.out.as_ref().is_none_or(|o| o.exists());
733            self.seen.push((job.clone(), present));
734            let r = (self.answer)(job);
735            if let (Ok(_), Some(out)) = (&r, &job.out) {
736                fs::write(out, b"copy").unwrap();
737            }
738            r
739        }
740    }
741
742    fn fake<F: FnMut(&ChildJob) -> Result<u64, String>>(answer: F) -> Fake<F> {
743        Fake {
744            answer,
745            seen: Vec::new(),
746        }
747    }
748
749    #[test]
750    fn child_args_round_trip() {
751        #[cfg(unix)]
752        let odd: PathBuf = {
753            use std::os::unix::ffi::OsStringExt;
754            OsString::from_vec(b"/tmp/a b\xff/-c".to_vec()).into()
755        };
756        #[cfg(not(unix))]
757        let odd = PathBuf::from("C:\\a b\\-c");
758        let jobs = [
759            ChildJob {
760                side: Side::Compress,
761                codec: "xz",
762                level: Some(9),
763                threads: Some(4),
764                memory_limit: 123_456_789,
765                out: Some(odd.join("out")),
766                files: vec![
767                    Snapshot {
768                        path: odd.clone(),
769                        len: 7,
770                    },
771                    Snapshot {
772                        path: PathBuf::from("-"),
773                        len: 0,
774                    },
775                ],
776            },
777            ChildJob {
778                side: Side::Decompress,
779                codec: "zstd",
780                level: Some(-5),
781                threads: None,
782                memory_limit: u64::MAX,
783                out: None,
784                files: vec![Snapshot {
785                    path: odd,
786                    len: u64::MAX,
787                }],
788            },
789            ChildJob {
790                side: Side::Compress,
791                codec: "lz4",
792                level: None,
793                threads: None,
794                memory_limit: 0,
795                out: None,
796                files: Vec::new(),
797            },
798        ];
799        for job in jobs {
800            let args = job.to_args();
801            assert_eq!(ChildJob::from_args(&args), Ok(job.clone()), "{args:?}");
802        }
803        // Not a codec of this build, a bad side, a missing output, a
804        // missing length, a bad level.
805        for bad in [
806            &["compress", "nope", "-", "-", "1", "-"][..],
807            &["sideways", "lz4", "-", "-", "1", "-"],
808            &["compress", "lz4", "-", "-", "1"],
809            &["compress", "lz4", "-", "-", "1", "/out", "/x"],
810            &["compress", "lz4", "x", "-", "1", "-"],
811        ] {
812            let args: Vec<OsString> = bad.iter().map(OsString::from).collect();
813            assert!(ChildJob::from_args(&args).is_err(), "{bad:?}");
814        }
815    }
816
817    #[test]
818    fn proc_status_fields_parse_in_kib() {
819        // The kernel pads the numbers with spaces; built, not spelled, so
820        // no literal carries a run of them.
821        let pad = " ".repeat(4);
822        let status = format!(
823            "Name:\tstuffr\nVmPeak:\t{pad}123456 kB\nVmHWM:\t{pad}6372 kB\nVmRSS:\t{pad}6100 kB\n"
824        );
825        let status = status.as_bytes();
826        assert_eq!(status_field_kib(status, b"VmHWM:"), Some(6372));
827        assert_eq!(status_field_kib(status, b"VmRSS:"), Some(6100));
828        assert_eq!(status_field_kib(status, b"VmSwap:"), None);
829        assert_eq!(status_field_kib(b"VmHWM:\t kB\n", b"VmHWM:"), None);
830    }
831
832    #[test]
833    fn child_reply_parses_ok_and_err() {
834        assert_eq!(parse_reply(b"ok 12345\n"), Some(Ok(12345)));
835        assert_eq!(parse_reply(b"ok 0\n"), Some(Ok(0)));
836        assert_eq!(
837            parse_reply(b"err refused by --memory-limit: too big\n"),
838            Some(Err("refused by --memory-limit: too big".into()))
839        );
840        for junk in [
841            &b""[..],
842            b"ok\n",
843            b"ok -3\n",
844            b"ok 12x\n",
845            b"hello\n",
846            b"err \n",
847        ] {
848            assert_eq!(
849                parse_reply(junk),
850                None,
851                "{:?}",
852                String::from_utf8_lossy(junk)
853            );
854        }
855    }
856
857    /// The env vars that make [`memory_probe`] run, and the prefix its
858    /// reply line carries in libtest's captured output.
859    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
860    const PROBE_CODEC: &str = "STUFFR_BENCH_MEMORY_PROBE_CODEC";
861    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
862    const PROBE_LEVEL: &str = "STUFFR_BENCH_MEMORY_PROBE_LEVEL";
863    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
864    const PROBE_FILE: &str = "STUFFR_BENCH_MEMORY_PROBE_FILE";
865    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
866    const PROBE_MARK: &str = "memory-probe-reply: ";
867
868    /// Not a test of its own: [`fresh_processes_measure_xz_9_above_lz4`]
869    /// re-runs this test binary with only this test, so the compress child
870    /// it runs has a process (and an `ru_maxrss` high-water mark) to
871    /// itself. Prints the child's reply line after [`PROBE_MARK`]. Does
872    /// nothing without the env vars.
873    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
874    #[test]
875    #[ignore = "a probe that fresh_processes_measure_xz_9_above_lz4 runs in its own process"]
876    fn memory_probe() {
877        let (Ok(codec), Ok(level), Some(file)) = (
878            std::env::var(PROBE_CODEC),
879            std::env::var(PROBE_LEVEL),
880            std::env::var_os(PROBE_FILE),
881        ) else {
882            return;
883        };
884        let path = PathBuf::from(file);
885        let len = fs::metadata(&path).unwrap().len();
886        let args: Vec<OsString> = ["compress", &codec, &level, "-", "1073741824", "-"]
887            .iter()
888            .map(OsString::from)
889            .chain([OsString::from(len.to_string()), path.into_os_string()])
890            .collect();
891        let job = ChildJob::from_args(&args).unwrap();
892        println!("{PROBE_MARK}{}", reply_line(&child_main(job)));
893    }
894
895    /// The compress child's figures for xz -9 and lz4 over 1 MiB, each
896    /// from a fresh process running [`memory_probe`]: xz -9's encoder (a
897    /// 64 MiB dictionary's hash table) touches more than lz4's. Each figure
898    /// is that process's own peak above its own baseline, so no other test
899    /// can move it.
900    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
901    #[test]
902    fn fresh_processes_measure_xz_9_above_lz4() {
903        let dir = Scratch::new();
904        let input = dir.file("in", &text(1 << 20, 3));
905        let probe = |codec: &str, level: &str| -> u64 {
906            let out = std::process::Command::new(std::env::current_exe().unwrap())
907                .args([
908                    "--exact",
909                    "bench::memory::tests::memory_probe",
910                    "--ignored",
911                    "--nocapture",
912                    "--test-threads=1",
913                ])
914                .env(PROBE_CODEC, codec)
915                .env(PROBE_LEVEL, level)
916                .env(PROBE_FILE, &input.path)
917                .stdin(std::process::Stdio::null())
918                .output()
919                .unwrap();
920            let stdout = String::from_utf8_lossy(&out.stdout);
921            let reply = stdout
922                .lines()
923                .find_map(|l| l.split_once(PROBE_MARK).map(|(_, r)| r));
924            match reply.and_then(|r| parse_reply(r.as_bytes())) {
925                Some(Ok(n)) if out.status.success() => n,
926                _ => panic!(
927                    "{codec} probe gave no `ok` reply ({}):\n--- stdout\n{stdout}\n--- stderr\n{}",
928                    out.status,
929                    String::from_utf8_lossy(&out.stderr)
930                ),
931            }
932        };
933        let lz4 = probe("lz4", "-");
934        let xz = probe("xz", "9");
935        assert!(xz > lz4, "xz -9 {xz} vs lz4 {lz4}");
936    }
937
938    /// `input` compressed by the compress child (in-process) into the
939    /// existing file `name` in `dir`, as the parent's temp file is filled.
940    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
941    fn compressed_by_the_child(
942        dir: &Scratch,
943        name: &str,
944        codec: &'static str,
945        level: Option<i32>,
946        input: &Snapshot,
947    ) -> Snapshot {
948        let out = dir.file(name, b"stale bytes the child must truncate");
949        let c = child_main(ChildJob {
950            side: Side::Compress,
951            codec,
952            level,
953            threads: None,
954            memory_limit: 1 << 30,
955            out: Some(out.path.clone()),
956            files: vec![input.clone()],
957        });
958        assert!(c.is_ok(), "{codec}: {c:?}");
959        let len = fs::metadata(&out.path).unwrap().len();
960        Snapshot {
961            path: out.path,
962            len,
963        }
964    }
965
966    /// The compress child writes the compressed stream to its output file
967    /// (truncating what was there), and the decompress child runs to the
968    /// end over it and gets the input back. Only that each succeeds is
969    /// checked of the figures: in-process, they are this whole test
970    /// binary's high-water mark, shared with every test beside it.
971    #[cfg(any(target_vendor = "apple", target_os = "linux"))]
972    #[test]
973    fn the_compress_child_writes_the_copy_the_decompress_child_reads() {
974        let dir = Scratch::new();
975        let original = text(1 << 20, 3);
976        let input = dir.file("in", &original);
977        for (codec, level) in [("lz4", None), ("xz", Some(9))] {
978            let copy = compressed_by_the_child(&dir, codec, codec, level, &input);
979            let id = codec_id(codec).unwrap();
980            let decoder = stuffr::registry().require_decoder(id).unwrap();
981            let mut back = Vec::new();
982            let src: Box<dyn Source> = Box::new(temp_source(
983                GuardedSeek::new(File::open(&copy.path).unwrap()),
984                copy.len,
985            ));
986            io::Read::read_to_end(
987                &mut decoder.decoder(src, &DecodeOpts::default()).unwrap(),
988                &mut back,
989            )
990            .unwrap();
991            assert!(
992                back == original,
993                "{codec}: the copy does not decode to the input"
994            );
995            let d = child_main(ChildJob {
996                side: Side::Decompress,
997                codec,
998                level: None,
999                threads: None,
1000                memory_limit: 1 << 30,
1001                out: None,
1002                files: vec![copy],
1003            });
1004            assert!(d.is_ok(), "{codec}: {d:?}");
1005        }
1006    }
1007
1008    #[test]
1009    fn per_file_reports_the_maximum() {
1010        let _g = armed_lock();
1011        let dir = Scratch::new();
1012        let files = vec![
1013            dir.file("a", &text(10_000, 0)),
1014            dir.file("b", &text(30_000, 1)),
1015            dir.file("c", &text(20_000, 2)),
1016        ];
1017        // Compress answers 100 x the file's length; decompress answers by
1018        // order: 5, 9, 2.
1019        let mut d = [5u64, 9, 2].into_iter();
1020        let mut r = fake(|job: &ChildJob| match job.side {
1021            Side::Compress => Ok(job.files[0].len * 100),
1022            Side::Decompress => Ok(d.next().unwrap()),
1023        });
1024        let m = measure_memory_with(&mut r, "gzip", Some(6), &plan(Mode::PerFile), &files);
1025        assert_eq!(
1026            m,
1027            Memory {
1028                compress: Some(3_000_000),
1029                decompress: Some(9),
1030                error: None
1031            }
1032        );
1033        // One child pair per file, each child given exactly its one file
1034        // (compress: the input, writing to the temp copy; decompress: that
1035        // copy, present, at the length the compress child left it).
1036        let sides: Vec<Side> = r.seen.iter().map(|(j, _)| j.side).collect();
1037        use Side::{Compress as C, Decompress as D};
1038        assert_eq!(sides, [C, D, C, D, C, D]);
1039        for (i, (job, present)) in r.seen.iter().enumerate() {
1040            assert_eq!(job.files.len(), 1);
1041            assert!(present, "job {i}'s file and output should exist");
1042            let level = (job.side == Side::Compress).then_some(6);
1043            assert_eq!((job.codec, job.level), ("gzip", level));
1044            if job.side == Side::Compress {
1045                assert_eq!(job.files[0], files[i / 2]);
1046                assert!(job.out.is_some(), "a compress child writes the copy");
1047            } else {
1048                assert_eq!(job.out, None);
1049                assert_eq!(job.files[0].len, 4, "the length the compress child wrote");
1050                assert_eq!(Some(&job.files[0].path), r.seen[i - 1].0.out.as_ref());
1051            }
1052        }
1053
1054        // One stream: one pair, the compress child given every file.
1055        let mut r = fake(|job: &ChildJob| Ok(job.files.iter().map(|s| s.len).sum()));
1056        let m = measure_memory_with(&mut r, "gzip", None, &plan(Mode::OneStream), &files);
1057        assert_eq!(m.compress, Some(60_000));
1058        assert!(m.decompress.is_some_and(|n| n > 0));
1059        assert_eq!(r.seen.len(), 2);
1060        assert_eq!(r.seen[0].0.files, files);
1061        assert_eq!(r.seen[1].0.files.len(), 1);
1062    }
1063
1064    #[test]
1065    fn a_failed_child_gives_dashes_and_an_error() {
1066        let _g = armed_lock();
1067        let dir = Scratch::new();
1068        let files = vec![
1069            dir.file("a", &text(10_000, 0)),
1070            dir.file("b", &text(10_000, 1)),
1071        ];
1072        // The second file's compress child fails: it left no copy, so
1073        // there is no decompress run for that file and both cells are
1074        // dashes (ruling F8); the error says why.
1075        let mut n = 0;
1076        let mut r = fake(|job: &ChildJob| match job.side {
1077            Side::Compress => {
1078                n += 1;
1079                if n == 2 {
1080                    Err("the memory run failed (signal: 9 (SIGKILL))".into())
1081                } else {
1082                    Ok(10)
1083                }
1084            }
1085            Side::Decompress => Ok(20),
1086        });
1087        let m = measure_memory_with(&mut r, "gzip", None, &plan(Mode::PerFile), &files);
1088        assert_eq!(m.compress, None);
1089        assert_eq!(m.decompress, None);
1090        let e = m.error.unwrap();
1091        assert!(e.starts_with("compress: ") && e.contains("SIGKILL"), "{e}");
1092        assert!(e.contains("no decompress run"), "{e}");
1093        // No child ran after the failure: two compress, one decompress.
1094        assert_eq!(r.seen.len(), 3);
1095
1096        // A failed decompress child leaves the compress side standing
1097        // (spec call 16).
1098        let mut r = fake(|job: &ChildJob| match job.side {
1099            Side::Compress => Ok(10),
1100            Side::Decompress => Err("the memory run failed (exit status: 1)".into()),
1101        });
1102        let m = measure_memory_with(&mut r, "gzip", None, &plan(Mode::PerFile), &files);
1103        assert_eq!((m.compress, m.decompress), (Some(10), None));
1104        assert!(m.error.unwrap().starts_with("decompress: "));
1105
1106        // Both sides failing: two dashes.
1107        let mut r = fake(|_: &ChildJob| Err("boom".to_string()));
1108        let m = measure_memory_with(&mut r, "gzip", None, &plan(Mode::OneStream), &files);
1109        assert_eq!((m.compress, m.decompress), (None, None));
1110        assert!(m.error.unwrap().contains("boom"));
1111
1112        // A child that printed no reply, or exited badly, is a failure.
1113        assert!(interpret(false, "exit status: 101", b"", b"thread 'main' panicked\n").is_err());
1114        assert!(interpret(true, "exit status: 0", b"garbage\n", b"").is_err());
1115        assert!(interpret(false, "signal: 9 (SIGKILL)", b"ok 5\n", b"").is_err());
1116        assert_eq!(interpret(true, "exit status: 0", b"ok 5\n", b""), Ok(5));
1117        assert_eq!(
1118            interpret(false, "exit status: 1", b"err nope\n", b""),
1119            Err("nope".into())
1120        );
1121    }
1122
1123    #[test]
1124    fn a_refusal_names_the_memory_limit() {
1125        let refusal = describe(&stuffr::Error::ResourceLimit(
1126            "xz: header declares a dictionary needing 64 MiB".into(),
1127        ));
1128        assert!(
1129            refusal.starts_with("refused by --memory-limit"),
1130            "{refusal}"
1131        );
1132        assert!(refusal.contains("64 MiB"), "{refusal}");
1133        // Anything else keeps its own text.
1134        let corrupt = describe(&stuffr::Error::Corrupt("bad".into()));
1135        assert!(!corrupt.contains("--memory-limit"), "{corrupt}");
1136
1137        // The parent carries the child's refusal into the row's error.
1138        let _g = armed_lock();
1139        let dir = Scratch::new();
1140        let files = vec![dir.file("a", &text(10_000, 0))];
1141        let mut r = fake(|job: &ChildJob| match job.side {
1142            Side::Compress => Err(refusal.clone()),
1143            Side::Decompress => Ok(1),
1144        });
1145        let m = measure_memory_with(&mut r, "gzip", None, &plan(Mode::OneStream), &files);
1146        assert_eq!(m.compress, None);
1147        assert!(
1148            m.error.unwrap().contains("refused by --memory-limit"),
1149            "the refusal is distinguishable from a crash"
1150        );
1151    }
1152
1153    /// The pure xz decoder refuses a header whose dictionary exceeds the
1154    /// limit; the child reports that as a refusal. (The C tier's decoder
1155    /// takes no limit, so this is pure-tier only.)
1156    #[cfg(all(
1157        any(target_vendor = "apple", target_os = "linux"),
1158        not(feature = "xz-c")
1159    ))]
1160    #[test]
1161    fn a_child_refusal_names_the_memory_limit() {
1162        let dir = Scratch::new();
1163        let input = dir.file("in", &text(100_000, 0));
1164        let copy = compressed_by_the_child(&dir, "in.xz", "xz", Some(9), &input);
1165        let r = child_main(ChildJob {
1166            side: Side::Decompress,
1167            codec: "xz",
1168            level: None,
1169            threads: None,
1170            memory_limit: 1 << 20,
1171            out: None,
1172            files: vec![copy],
1173        });
1174        let e = r.unwrap_err();
1175        assert!(e.starts_with("refused by --memory-limit"), "{e}");
1176    }
1177
1178    #[test]
1179    fn temp_copies_are_removed() {
1180        let _g = armed_lock();
1181        let dir = Scratch::new();
1182        let files = vec![
1183            dir.file("a", &text(10_000, 0)),
1184            dir.file("b", &text(10_000, 1)),
1185        ];
1186        assert_eq!(bench_temps(), Vec::<String>::new());
1187        for mode in [Mode::OneStream, Mode::PerFile] {
1188            // While a decompress child runs, its copy exists and is the
1189            // only bench temp file: per-file copies do not pile up.
1190            let mut r = fake(|job: &ChildJob| {
1191                if job.side == Side::Decompress {
1192                    assert_eq!(bench_temps().len(), 1);
1193                }
1194                Ok(1)
1195            });
1196            let m = measure_memory_with(&mut r, "gzip", None, &plan(mode), &files);
1197            assert_eq!(m.error, None);
1198            assert!(r.seen.iter().all(|(_, present)| *present));
1199            assert_eq!(bench_temps(), Vec::<String>::new(), "{mode:?}");
1200        }
1201        // A failed child removes its copy too.
1202        let mut r = fake(|_: &ChildJob| Err("boom".to_string()));
1203        let _ = measure_memory_with(&mut r, "gzip", None, &plan(Mode::PerFile), &files);
1204        assert_eq!(bench_temps(), Vec::<String>::new());
1205    }
1206}