Skip to main content

stuffr_cli/bench/
measure.rs

1//! The timing harness for `stuffr bench`: one format at one level, run over
2//! the user's input, compressed and decompressed, with only codec work
3//! inside a timed region, except the temp file's writes when the output is
4//! too large to keep in memory (below).
5//!
6//! Compression's timed regions hold encoder creation, `write_all` of the
7//! input and `Sink::finish` (which consumes the sink, so its teardown is
8//! inside too). Decompression's hold decoder creation and reading the
9//! decoder to its end into a buffer whose bytes are only counted. Building
10//! the destination, opening files, reading streamed input, the temp file's
11//! creation and removal, and dropping the decoder are all outside.
12//!
13//! Each phase is the sum of several timed pieces, each one pair of clock
14//! reads. Compression: encoder creation, then one piece per 64 KiB
15//! `write_all` of the input, then `finish`. Decompression: one piece per
16//! 64 KiB the decoder fills, the decoder's creation in the first, its own
17//! reads of the compressed bytes being part of decoding. In memory and
18//! streamed alike; streamed, the input's reads fall between the pieces.
19//! Otherwise nothing runs between them but the progress report, so the
20//! pieces cost every format the same: one pair of clock reads per 64 KiB.
21//!
22//! Each codec gets the destination its own worst case allows, decided
23//! before any timing: memory whenever that worst case (input + 1/8 +
24//! 1 KiB, or 2 x input + 1 KiB for `compress`, LZW) fits under the memory
25//! left for it, half of `--memory-limit` when the input is held in memory,
26//! all of it when the input is streamed. Only past that does its output go
27//! to one temp file, and then its writes (to the page cache, through the
28//! 64 KiB buffer) fall inside its compression regions, so a row that went
29//! to disk is marked ([`Measured::on_disk`]; the table's `*`, the `output`
30//! field of JSON and CSV). With `--per-file` each file's stream is decided
31//! by its own size. Both destinations have the same shape: a 64 KiB
32//! `BufWriter` over a pre-reserved `Vec`, or over the temp file, so a codec
33//! that writes in small pieces pays the same either way. Decoding reads the
34//! same shape back: the `Vec` as a seekable, sized
35//! source, or the temp file through a 64 KiB `BufReader`, so a decoder that
36//! reads a byte at a time (LZW's) costs a memory copy per byte, not a
37//! system call.
38//!
39//! Each file's length is taken once, when the input is loaded; every later
40//! read stops there, so a file still being written is benched as it was.
41//!
42//! bench ignores STUFFR_THREADS; threads come only from --threads.
43//!
44//! With `--threads` above 1, a multi-threaded encoder's workers keep running
45//! in the untimed gaps between pieces, so its compression time can be
46//! under-reported against a single-threaded one. Streamed, the gaps hold
47//! the input's reads; in memory, only the progress report (a clock read,
48//! and a redraw at most every 100 ms).
49
50use std::fs::{self, File};
51use std::io::{self, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write};
52use std::path::{Path, PathBuf};
53use std::sync::atomic::{AtomicU32, Ordering};
54use std::sync::{Arc, Mutex};
55use std::time::Instant;
56
57use stuffr::ops::{CompressOpts, resolved_budget};
58use stuffr::source::{GuardedSeek, SeekRead, Source, SourceCaps};
59use stuffr::{Codec, DecodeOpts, EncodeOpts, Error, FormatKind, Governor};
60
61use super::score::{Measured, Outcome};
62use crate::compat::{CleanupGuard, create_armed};
63
64/// The copy buffer, as `ops` uses.
65pub(super) const CHUNK: usize = 64 * 1024;
66
67/// A monotonic clock in nanoseconds; injected so tests can count its reads.
68pub trait Clock {
69    fn now(&mut self) -> u64;
70}
71
72/// The real clock: nanoseconds since it was made.
73pub struct RealClock(Instant);
74
75impl RealClock {
76    pub fn new() -> Self {
77        Self(Instant::now())
78    }
79}
80
81impl Default for RealClock {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl Clock for RealClock {
88    fn now(&mut self) -> u64 {
89        u64::try_from(self.0.elapsed().as_nanos()).unwrap_or(u64::MAX)
90    }
91}
92
93/// The input, in operand order: held in memory (one buffer per file) or
94/// streamed from the files.
95#[derive(Debug)]
96pub enum Input {
97    Memory(Vec<Vec<u8>>),
98    Files(Vec<Snapshot>),
99}
100
101/// A streamed file and its length when the input was loaded: every read of
102/// it stops at `len`, so bytes appended later are never benched.
103#[derive(Clone, Debug, PartialEq, Eq)]
104pub struct Snapshot {
105    pub path: PathBuf,
106    pub len: u64,
107}
108
109impl Snapshot {
110    /// The file, opened and limited to its loaded length.
111    pub(super) fn open(&self) -> io::Result<io::Take<File>> {
112        Ok(File::open(&self.path)?.take(self.len))
113    }
114
115    /// [`Snapshot::open`] while the input is loaded: a file that cannot be
116    /// opened is the user's to fix, so a usage error naming it.
117    fn open_input(&self) -> stuffr::Result<io::Take<File>> {
118        self.open().map_err(|e| super::unreadable(&self.path, &e))
119    }
120}
121
122/// One stream of the files concatenated, or one stream per file.
123#[derive(Clone, Copy, Debug, PartialEq, Eq)]
124pub enum Mode {
125    OneStream,
126    PerFile,
127}
128
129/// The conditions every row runs under.
130#[derive(Clone, Copy, Debug, PartialEq, Eq)]
131pub struct Plan {
132    pub mode: Mode,
133    /// Runs per row; the reported times are the medians. 0 counts as 1.
134    pub repeat: u32,
135    /// As `pack --threads`.
136    pub threads: Option<usize>,
137    pub memory_limit: u64,
138}
139
140/// Loads `paths` into memory when their total size is at most half of
141/// `memory_limit`, leaving room for the compressed copy; otherwise the input
142/// is streamed from the files.
143///
144/// Also returns the input's size, measured here once: the bytes read into
145/// memory, or the files' sizes when streamed. It is the ratio's denominator,
146/// and each file's length is snapshotted with it (reads stop there), so a
147/// file still growing while the rows run changes neither.
148///
149/// A file whose size cannot be read, or that cannot be opened, is a usage
150/// error naming it (exit 2); only a failing read is an i/o error.
151pub fn load_input(paths: &[PathBuf], memory_limit: u64) -> stuffr::Result<(Input, u64)> {
152    let mut total = 0u64;
153    let mut snaps = Vec::with_capacity(paths.len());
154    for p in paths {
155        let len = fs::metadata(p).map_err(|e| super::unreadable(p, &e))?.len();
156        total = total.saturating_add(len);
157        snaps.push(Snapshot {
158            path: p.clone(),
159            len,
160        });
161    }
162    if total > memory_limit / 2 {
163        return Ok((Input::Files(snaps), total));
164    }
165    let mut parts = Vec::with_capacity(snaps.len());
166    for s in &snaps {
167        let mut buf = Vec::with_capacity(usize::try_from(s.len).unwrap_or(0));
168        s.open_input()?.read_to_end(&mut buf)?;
169        parts.push(buf);
170    }
171    let read = parts.iter().map(|b| b.len() as u64).sum();
172    Ok((Input::Memory(parts), read))
173}
174
175/// One untimed read of every file, so the first timed run does not pay for
176/// a cold cache. Nothing to do for input already in memory. A file that
177/// cannot be opened is a usage error naming it, as in [`load_input`].
178pub fn warm(input: &Input) -> stuffr::Result<()> {
179    if let Input::Files(files) = input {
180        let mut buf = vec![0u8; CHUNK];
181        for s in files {
182            let mut f = s.open_input()?;
183            while read_some(&mut f, &mut buf)? > 0 {}
184        }
185    }
186    Ok(())
187}
188
189/// Which half of a row's round trip a progress report is about.
190#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191pub enum Phase {
192    Compress,
193    Decompress,
194}
195
196/// Compresses and decompresses `input` with `codec` at `level` (`None`: the
197/// encoder's own default) under `plan`. Every error, a rejected level
198/// included, is a failed row; nothing panics.
199///
200/// `progress` is called only between timed pieces, never inside one, with
201/// `(phase, done, total)`; `total` is the input's size. Each repeat opens
202/// with exactly one `(Compress, 0, total)`; after it, `Compress` reports
203/// the input bytes the encoder has taken, rising strictly, and `Decompress`
204/// the bytes the decoder has given back, never falling; both end at
205/// `total` and never pass it. With `Mode::PerFile` the two phases take
206/// turns, file by file, each counting on across the files.
207pub fn measure(
208    codec: &str,
209    level: Option<i32>,
210    input: &Input,
211    plan: &Plan,
212    clock: &mut dyn Clock,
213    progress: &mut dyn FnMut(Phase, u64, u64),
214) -> Outcome {
215    match run(codec, level, input, plan, clock, progress) {
216        Ok(m) => Outcome::Ok(m),
217        Err(e) => Outcome::Failed(e.to_string()),
218    }
219}
220
221/// One file's input as the harness reads it.
222#[derive(Clone, Copy)]
223enum Part<'a> {
224    Memory(&'a [u8]),
225    File(&'a Snapshot),
226}
227
228impl Part<'_> {
229    /// The length the input was loaded at, never re-read from the file.
230    fn len(&self) -> u64 {
231        match self {
232            Part::Memory(b) => b.len() as u64,
233            Part::File(s) => s.len,
234        }
235    }
236}
237
238/// A part as compression takes it: its bytes, or its file opened.
239enum Feed<'a> {
240    Memory(&'a [u8]),
241    File(&'a Snapshot, io::Take<File>),
242}
243
244/// Where one stream's compressed output goes. Decided per codec, from its
245/// own worst case, before any timing (see [`destination`]), and the same
246/// for all of that codec's repeats. A row with any stream on the temp
247/// file is marked as such ([`Measured::on_disk`]).
248#[derive(Clone, Copy, Debug, PartialEq, Eq)]
249enum Dest {
250    Memory,
251    TempFile,
252}
253
254/// Memory whenever `codec`'s worst-case output for `input_len` bytes (see
255/// [`output_room`]) fits in the memory left for it: half of `memory_limit`
256/// beside input held in memory, all of it when the input is `streamed` (it
257/// holds no copy of the input). The temp file only past that, so the
258/// reservation never exceeds the budget. Each codec gets the destination
259/// its own worst case allows; the row reports which it got.
260fn destination(codec: &str, input_len: u64, memory_limit: u64, streamed: bool) -> Dest {
261    // Small test inputs always fit; this is how the tests reach the temp
262    // file with a memory limit the decoders can still work under.
263    #[cfg(test)]
264    if tests::FORCE_DISK.with(std::cell::Cell::get) {
265        return Dest::TempFile;
266    }
267    let budget = if streamed {
268        memory_limit
269    } else {
270        memory_limit / 2
271    };
272    match output_room(codec, input_len) {
273        Some(w) if w <= budget => Dest::Memory,
274        _ => Dest::TempFile,
275    }
276}
277
278/// The capacity reserved for `codec`'s output in memory, large enough that
279/// the `Vec` never reallocates inside a timed region. LZW (`compress`)
280/// writes no code wider than 16 bits and every code stands for at least one
281/// input byte, so it never writes more than two bytes per input byte (plus
282/// its 3-byte header and the padding at each code-width step); every other
283/// codec stays within 1/8 + 1 KiB, which
284/// `the_reserved_room_holds_incompressible_output` measures. Reserved pages
285/// a codec never writes are never touched. `None`: past `u64`, so it fits
286/// nowhere.
287fn output_room(codec: &str, input_len: u64) -> Option<u64> {
288    match codec {
289        "compress" => input_len.checked_mul(2),
290        _ => input_len.checked_add(input_len / 8),
291    }
292    .and_then(|n| n.checked_add(1024))
293}
294
295fn run(
296    name: &str,
297    level: Option<i32>,
298    input: &Input,
299    plan: &Plan,
300    clock: &mut dyn Clock,
301    report: &mut dyn FnMut(Phase, u64, u64),
302) -> stuffr::Result<Measured> {
303    let reg = stuffr::registry();
304    let id = reg
305        .matrix()
306        .into_iter()
307        .find(|r| r.kind == FormatKind::Codec && r.id.as_str() == name)
308        .map(|r| r.id)
309        .ok_or_else(|| Error::Usage(format!("`{name}` is not a codec in this build")))?;
310    let codec = reg.require_encoder(id)?;
311    reg.require_decoder(id)?;
312
313    let encode = EncodeOpts {
314        level,
315        governor: governor_for(plan),
316        ..Default::default()
317    };
318    codec.check_encode_opts(&encode)?;
319    let decode = DecodeOpts {
320        memory_limit: Some(plan.memory_limit),
321        ..Default::default()
322    };
323
324    let parts: Vec<Part<'_>> = match input {
325        Input::Memory(bufs) => bufs.iter().map(|b| Part::Memory(b)).collect(),
326        Input::Files(files) => files.iter().map(Part::File).collect(),
327    };
328    let groups: Vec<&[Part<'_>]> = match plan.mode {
329        Mode::OneStream => vec![&parts[..]],
330        Mode::PerFile => parts.chunks(1).collect(),
331    };
332    let group_len = |g: &[Part<'_>]| g.iter().map(Part::len).sum::<u64>();
333    let streamed = matches!(input, Input::Files(_));
334    let dest = |g: &[Part<'_>]| destination(name, group_len(g), plan.memory_limit, streamed);
335    let temp = if groups.iter().any(|g| dest(g) == Dest::TempFile) {
336        Some(TempOut::create()?)
337    } else {
338        None
339    };
340
341    let mut buf = vec![0u8; CHUNK];
342    let mut bytes_out = None;
343    let mut on_disk = false;
344    let mut compress = Vec::new();
345    let mut decompress = Vec::new();
346    let mut progress = Progress::new(report, parts.iter().map(Part::len).sum());
347    for _ in 0..plan.repeat.max(1) {
348        progress.start();
349        let (mut out, mut c_ns, mut d_ns) = (0u64, 0u64, 0u64);
350        for group in &groups {
351            let expected = group_len(group);
352            let job = Job {
353                name: id.as_str(),
354                codec: codec.as_ref(),
355                dest: dest(group),
356                encode: &encode,
357                decode: &decode,
358                temp: temp.as_ref(),
359                buf: &mut buf,
360            };
361            let row = job.run(group, expected, clock, &mut progress)?;
362            out += row.bytes_out;
363            c_ns += row.compress_ns;
364            d_ns += row.decompress_ns;
365            on_disk |= row.on_disk;
366        }
367        bytes_out.get_or_insert(out);
368        compress.push(c_ns);
369        decompress.push(d_ns);
370    }
371    Ok(Measured {
372        bytes_out: bytes_out.unwrap_or(0),
373        compress_ns: median(compress),
374        decompress_ns: median(decompress),
375        on_disk,
376    })
377}
378
379/// The encode governor for `plan`: `pack --threads`'s mapping, except that
380/// no `--threads` means single-threaded whatever `STUFFR_THREADS` says
381/// (`resolved_budget` would read it). With a CLI value the variable is
382/// already ignored: `resolve_workers` takes the CLI value first.
383pub fn governor_for(plan: &Plan) -> Option<Arc<Governor>> {
384    let t = plan.threads?;
385    resolved_budget(&CompressOpts {
386        threads: Some(t),
387        memory_limit: Some(plan.memory_limit),
388        ..Default::default()
389    })
390}
391
392/// The middle value; the lower of the two middles for an even count.
393fn median(mut v: Vec<u64>) -> u64 {
394    v.sort_unstable();
395    v.get(v.len().saturating_sub(1) / 2).copied().unwrap_or(0)
396}
397
398/// The progress reports of one row, made between timed pieces only (see
399/// [`measure`] for what they promise).
400struct Progress<'a> {
401    report: &'a mut dyn FnMut(Phase, u64, u64),
402    total: u64,
403    fed: u64,
404    decoded: u64,
405}
406
407impl<'a> Progress<'a> {
408    fn new(report: &'a mut dyn FnMut(Phase, u64, u64), total: u64) -> Self {
409        Self {
410            report,
411            total,
412            fed: 0,
413            decoded: 0,
414        }
415    }
416
417    /// A repeat begins: its one `(Compress, 0, total)`.
418    fn start(&mut self) {
419        self.fed = 0;
420        self.decoded = 0;
421        (self.report)(Phase::Compress, 0, self.total);
422    }
423
424    /// The encoder took `n` more input bytes.
425    fn fed(&mut self, n: usize) {
426        let fed = self.fed.saturating_add(n as u64).min(self.total);
427        // Never a second `(Compress, 0, _)`: that opens a repeat.
428        if fed > self.fed {
429            self.fed = fed;
430            (self.report)(Phase::Compress, fed, self.total);
431        }
432    }
433
434    /// The decoder gave back `n` more bytes.
435    fn decoded(&mut self, n: usize) {
436        self.decoded = self.decoded.saturating_add(n as u64).min(self.total);
437        (self.report)(Phase::Decompress, self.decoded, self.total);
438    }
439}
440
441/// One timed piece: `f` between two clock reads, and nothing else.
442fn timed<T>(clock: &mut dyn Clock, f: impl FnOnce() -> T) -> (T, u64) {
443    let t0 = clock.now();
444    let r = f();
445    let t1 = clock.now();
446    (r, t1.saturating_sub(t0))
447}
448
449/// Everything one stream's round trip needs.
450struct Job<'a> {
451    name: &'a str,
452    codec: &'a dyn Codec,
453    dest: Dest,
454    encode: &'a EncodeOpts,
455    decode: &'a DecodeOpts,
456    temp: Option<&'a TempOut>,
457    buf: &'a mut [u8],
458}
459
460impl Job<'_> {
461    /// Compresses `group` as one stream, then decompresses it, checking the
462    /// round trip gives back `expected` bytes.
463    fn run(
464        mut self,
465        group: &[Part<'_>],
466        expected: u64,
467        clock: &mut dyn Clock,
468        progress: &mut Progress<'_>,
469    ) -> stuffr::Result<Measured> {
470        let on_disk = matches!((self.dest, self.temp), (Dest::TempFile, Some(_)));
471        let (bytes_out, compress_ns, src) = match (self.dest, self.temp) {
472            (Dest::Memory, _) | (Dest::TempFile, None) => {
473                // Room for incompressible output, so the outer Vec does not
474                // reallocate inside the region. (`TempFile` always comes
475                // with a temp file; `run` creates one whenever any group
476                // needs it.)
477                let room = output_room(self.name, expected).unwrap_or(0);
478                let shared = SharedVec::with_capacity(usize::try_from(room).unwrap_or(0));
479                let dst = BufWriter::with_capacity(CHUNK, shared.clone());
480                let ns = self.compress(group, Box::new(dst), clock, progress)?;
481                let bytes = shared.take();
482                let len = bytes.len() as u64;
483                (len, ns, Box::new(MemSource::new(bytes)) as Box<dyn Source>)
484            }
485            (Dest::TempFile, Some(temp)) => {
486                // Through the descriptor `create` made, never the name
487                // again (HOST-9): emptied and rewound for this row.
488                let mut file = temp.handle()?;
489                file.set_len(0)?;
490                file.seek(SeekFrom::Start(0))?;
491                let dst = BufWriter::with_capacity(CHUNK, file);
492                let ns = self.compress(group, Box::new(dst), clock, progress)?;
493                let mut file = temp.handle()?;
494                file.seek(SeekFrom::Start(0))?;
495                let len = file.metadata()?.len();
496                let src = temp_source(GuardedSeek::new(counted(file)), len);
497                (len, ns, Box::new(src) as Box<dyn Source>)
498            }
499        };
500        let (decoded, decompress_ns) =
501            decode(self.codec, self.decode, src, self.buf, clock, progress)?;
502        if decoded != expected {
503            return Err(Error::Corrupt(format!(
504                "the round trip gave back {decoded} bytes of {expected}"
505            )));
506        }
507        Ok(Measured {
508            bytes_out,
509            compress_ns,
510            decompress_ns,
511            on_disk,
512        })
513    }
514
515    /// The compression pieces' total: encoder creation, one `write_all`
516    /// per 64 KiB of input, `finish`. Streamed, each chunk is read between
517    /// two pieces.
518    fn compress(
519        &mut self,
520        group: &[Part<'_>],
521        dst: Box<dyn Write + Send>,
522        clock: &mut dyn Clock,
523        progress: &mut Progress<'_>,
524    ) -> stuffr::Result<u64> {
525        // Files are opened before any timing.
526        let mut feeds = Vec::with_capacity(group.len());
527        for p in group {
528            feeds.push(match p {
529                Part::Memory(b) => Feed::Memory(b),
530                Part::File(snap) => Feed::File(snap, snap.open()?),
531            });
532        }
533        let (sink, mut ns) = timed(clock, || self.codec.encoder(dst, self.encode));
534        let mut sink = sink?;
535        for feed in &mut feeds {
536            match feed {
537                Feed::Memory(b) => {
538                    for chunk in b.chunks(CHUNK) {
539                        let (r, t) = timed(clock, || sink.write_all(chunk));
540                        r?;
541                        ns += t;
542                        progress.fed(chunk.len());
543                    }
544                }
545                Feed::File(snap, f) => {
546                    let mut read = 0u64;
547                    loop {
548                        let n = read_some(f, self.buf)?;
549                        if n == 0 {
550                            if read < snap.len {
551                                return Err(shrank(snap, read));
552                            }
553                            break;
554                        }
555                        read += n as u64;
556                        let chunk = &self.buf[..n];
557                        let (r, t) = timed(clock, || sink.write_all(chunk));
558                        r?;
559                        ns += t;
560                        progress.fed(n);
561                    }
562                }
563            }
564        }
565        let (r, t) = timed(clock, || sink.finish());
566        r?;
567        Ok(ns + t)
568    }
569}
570
571/// Decompression: one piece per 64 KiB the decoder fills (`buf` is
572/// [`CHUNK`] long), the decoder made in the first, until a fill comes back
573/// short at the end; the bytes are only counted. Returns the count and the
574/// pieces' total. The decoder is dropped outside them.
575fn decode(
576    codec: &dyn Codec,
577    opts: &DecodeOpts,
578    src: Box<dyn Source>,
579    buf: &mut [u8],
580    clock: &mut dyn Clock,
581    progress: &mut Progress<'_>,
582) -> stuffr::Result<(u64, u64)> {
583    let (first, mut ns) = timed(clock, || {
584        let mut decoder = codec.decoder(src, opts)?;
585        let n = fill(&mut decoder, buf).map_err(Error::from_decode_io)?;
586        Ok::<_, Error>((decoder, n))
587    });
588    let (mut decoder, mut n) = first?;
589    let mut total = 0u64;
590    loop {
591        total += n as u64;
592        progress.decoded(n);
593        if n < buf.len() {
594            break;
595        }
596        let (r, t) = timed(clock, || fill(&mut decoder, buf));
597        ns += t;
598        n = r.map_err(Error::from_decode_io)?;
599    }
600    drop(decoder);
601    Ok((total, ns))
602}
603
604/// A streamed file ended at `read` bytes, short of the length it was
605/// loaded at.
606pub(super) fn shrank(snap: &Snapshot, read: u64) -> Error {
607    Error::Io(io::Error::new(
608        io::ErrorKind::UnexpectedEof,
609        format!(
610            "{} shrank to {read} bytes from {} while bench ran",
611            snap.path.display(),
612            snap.len
613        ),
614    ))
615}
616
617/// Reads into `buf` until it is full or the reader ends, retrying
618/// `Interrupted`. Short only at the end.
619pub(super) fn fill(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
620    let mut filled = 0;
621    while filled < buf.len() {
622        match r.read(&mut buf[filled..]) {
623            Ok(0) => break,
624            Ok(n) => filled += n,
625            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
626            Err(e) => return Err(e),
627        }
628    }
629    Ok(filled)
630}
631
632/// `read`, retried on `Interrupted`.
633pub(super) fn read_some(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
634    loop {
635        match r.read(buf) {
636            Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
637            other => return other,
638        }
639    }
640}
641
642/// The in-memory destination: the encoder owns one handle, the harness
643/// keeps the other to take the bytes back after `finish`.
644#[derive(Clone)]
645struct SharedVec(Arc<Mutex<Vec<u8>>>);
646
647impl SharedVec {
648    fn with_capacity(n: usize) -> Self {
649        Self(Arc::new(Mutex::new(Vec::with_capacity(n))))
650    }
651
652    #[cfg(test)]
653    fn capacity(&self) -> usize {
654        self.0.lock().unwrap_or_else(|e| e.into_inner()).capacity()
655    }
656
657    fn take(&self) -> Vec<u8> {
658        std::mem::take(&mut *self.0.lock().unwrap_or_else(|e| e.into_inner()))
659    }
660}
661
662impl Write for SharedVec {
663    fn write(&mut self, b: &[u8]) -> io::Result<usize> {
664        self.0
665            .lock()
666            .unwrap_or_else(|e| e.into_inner())
667            .extend_from_slice(b);
668        Ok(b.len())
669    }
670
671    fn flush(&mut self) -> io::Result<()> {
672        Ok(())
673    }
674}
675
676/// Compressed bytes in memory, as seekable and sized as the [`BufSource`]
677/// over the temp file, so a decoder sees the same capabilities either way.
678struct MemSource {
679    cursor: Cursor<Vec<u8>>,
680    len: u64,
681}
682
683impl MemSource {
684    fn new(bytes: Vec<u8>) -> Self {
685        let len = bytes.len() as u64;
686        Self {
687            cursor: Cursor::new(bytes),
688            len,
689        }
690    }
691}
692
693impl Read for MemSource {
694    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
695        self.cursor.read(buf)
696    }
697}
698
699impl Seek for MemSource {
700    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
701        self.cursor.seek(pos)
702    }
703}
704
705impl Source for MemSource {
706    fn caps(&self) -> SourceCaps {
707        SourceCaps {
708            seekable: true,
709            len: Some(self.len),
710        }
711    }
712
713    fn as_seek(&mut self) -> Option<&mut dyn SeekRead> {
714        Some(self)
715    }
716}
717
718/// The temp file's compressed bytes as the decoder reads them: through a
719/// 64 KiB buffer, so a decoder reading a byte at a time pays a memory copy
720/// per byte, as it does from [`MemSource`], not a system call. Seekable and
721/// sized like `MemSource`.
722pub(super) fn temp_source<R: Read + Seek + Send>(inner: R, len: u64) -> BufSource<R> {
723    BufSource {
724        inner: BufReader::with_capacity(CHUNK, inner),
725        len,
726    }
727}
728
729pub(super) struct BufSource<R> {
730    inner: BufReader<R>,
731    len: u64,
732}
733
734impl<R: Read> Read for BufSource<R> {
735    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
736        self.inner.read(buf)
737    }
738}
739
740impl<R: Seek> Seek for BufSource<R> {
741    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
742        self.inner.seek(pos)
743    }
744}
745
746impl<R: Read + Seek + Send> Source for BufSource<R> {
747    fn caps(&self) -> SourceCaps {
748        SourceCaps {
749            seekable: true,
750            len: Some(self.len),
751        }
752    }
753
754    fn as_seek(&mut self) -> Option<&mut dyn SeekRead> {
755        Some(self)
756    }
757}
758
759/// The temp file as `temp_source` reads it. Production: the file itself.
760#[cfg(not(test))]
761fn counted(file: File) -> File {
762    file
763}
764
765/// Under test, a file whose reads are counted per thread, so a test sees
766/// how many reach it through the path `measure` takes.
767#[cfg(test)]
768fn counted(file: File) -> tests::CountedFile {
769    tests::CountedFile(file)
770}
771
772/// The next temp file's `<n>`; also how the tests see whether a run made one.
773static TEMP_SEQ: AtomicU32 = AtomicU32::new(0);
774
775/// The streamed run's compressed output, `stuffr-bench-<pid>-<n>` in the
776/// temp directory: armed for removal on SIGINT, SIGTERM and SIGHUP while it
777/// lives, removed when it drops (a panic included).
778pub(super) struct TempOut {
779    path: PathBuf,
780    /// The descriptor `create` opened (read and write, `create_new`, 0600):
781    /// every use of the file goes through it ([`Self::handle`]), never
782    /// through the name, so nothing put at the name after the create —
783    /// a symlink to a victim file — is ever opened, truncated or written
784    /// (HOST-9). Only the memory child, another process, opens it by name,
785    /// and checks what it opened (`memory::open_copy`).
786    file: File,
787    /// Dropped after `drop` has removed the file, so there is no moment the
788    /// file exists unarmed.
789    _armed: CleanupGuard,
790}
791
792impl TempOut {
793    pub(super) fn create() -> io::Result<Self> {
794        let dir = std::env::temp_dir();
795        let mut tries = 0;
796        loop {
797            let n = TEMP_SEQ.fetch_add(1, Ordering::Relaxed);
798            let path = dir.join(format!("stuffr-bench-{}-{n}", std::process::id()));
799            match create_armed(&path) {
800                Ok((file, armed)) => {
801                    return Ok(Self {
802                        path,
803                        file,
804                        _armed: armed,
805                    });
806                }
807                Err(e) if e.kind() == io::ErrorKind::AlreadyExists && tries < 1000 => tries += 1,
808                Err(e) => return Err(e),
809            }
810        }
811    }
812
813    pub(super) fn path(&self) -> &Path {
814        &self.path
815    }
816
817    /// Another handle on the file `create` made (`try_clone`: the same open
818    /// file, sharing its offset — callers seek before use).
819    pub(super) fn handle(&self) -> io::Result<File> {
820        self.file.try_clone()
821    }
822}
823
824impl Drop for TempOut {
825    fn drop(&mut self) {
826        let _ = fs::remove_file(&self.path);
827    }
828}
829
830#[cfg(test)]
831mod tests {
832    use super::*;
833    use std::fs;
834    use std::path::PathBuf;
835    use std::sync::atomic::{AtomicU32, Ordering};
836
837    /// 0, 10, 20, …: every region it times lasts exactly 10.
838    struct FakeClock(u64);
839
840    impl Clock for FakeClock {
841        fn now(&mut self) -> u64 {
842            let t = self.0;
843            self.0 += 10;
844            t
845        }
846    }
847
848    /// Times region `i` (reads `2i` and `2i + 1`) as `durations[i]`.
849    struct ScriptClock {
850        durations: Vec<u64>,
851        reads: usize,
852        t: u64,
853    }
854
855    impl Clock for ScriptClock {
856        fn now(&mut self) -> u64 {
857            if self.reads % 2 == 1 {
858                self.t += self.durations[self.reads / 2];
859            } else {
860                self.t += 1_000;
861            }
862            self.reads += 1;
863            self.t
864        }
865    }
866
867    /// A scratch directory `stuffr-benchtest-<pid>-<n>`, removed on drop.
868    struct Scratch(PathBuf);
869
870    impl Scratch {
871        fn new() -> Self {
872            static N: AtomicU32 = AtomicU32::new(0);
873            let n = N.fetch_add(1, Ordering::Relaxed);
874            let dir =
875                std::env::temp_dir().join(format!("stuffr-benchtest-{}-{n}", std::process::id()));
876            fs::create_dir_all(&dir).unwrap();
877            Scratch(dir)
878        }
879
880        fn file(&self, name: &str, bytes: &[u8]) -> PathBuf {
881            let p = self.0.join(name);
882            fs::write(&p, bytes).unwrap();
883            p
884        }
885    }
886
887    impl Drop for Scratch {
888        fn drop(&mut self) {
889            let _ = fs::remove_dir_all(&self.0);
890        }
891    }
892
893    /// Serialises the tests that arm the process-wide signal cleanup.
894    fn armed_lock() -> std::sync::MutexGuard<'static, ()> {
895        crate::compat::ARM_TEST_LOCK
896            .lock()
897            .unwrap_or_else(|e| e.into_inner())
898    }
899
900    fn text(len: usize, seed: u8) -> Vec<u8> {
901        b"the quick brown fox jumps over the lazy dog, "
902            .iter()
903            .cycle()
904            .enumerate()
905            .map(|(i, &b)| b.wrapping_add(seed).wrapping_add((i / 997) as u8))
906            .take(len)
907            .collect()
908    }
909
910    fn plan(mode: Mode, repeat: u32) -> Plan {
911        Plan {
912            mode,
913            repeat,
914            threads: None,
915            memory_limit: 1 << 30,
916        }
917    }
918
919    thread_local! {
920        /// Set by `on_disk`: every output goes to the temp file.
921        pub(super) static FORCE_DISK: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
922    }
923
924    /// Runs `f` with every compressed output sent to the temp file.
925    fn on_disk<T>(f: impl FnOnce() -> T) -> T {
926        FORCE_DISK.with(|d| d.set(true));
927        let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
928        FORCE_DISK.with(|d| d.set(false));
929        r.unwrap_or_else(|e| std::panic::resume_unwind(e))
930    }
931
932    thread_local! {
933        /// Reads that reached a temp file on this thread (see `counted`).
934        static FILE_READS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
935    }
936
937    /// The temp file under test: counts each read that reaches it.
938    pub(super) struct CountedFile(pub(super) File);
939
940    impl Read for CountedFile {
941        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
942            FILE_READS.with(|r| r.set(r.get() + 1));
943            self.0.read(buf)
944        }
945    }
946
947    impl Seek for CountedFile {
948        fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
949            self.0.seek(pos)
950        }
951    }
952
953    /// The temp files created while `f` ran. Hold `armed_lock`: every test
954    /// that can create one does.
955    fn temps_made(f: impl FnOnce()) -> u32 {
956        let before = TEMP_SEQ.load(Ordering::Relaxed);
957        f();
958        TEMP_SEQ.load(Ordering::Relaxed) - before
959    }
960
961    /// Deterministic noise no codec can shrink.
962    fn noise(len: usize) -> Vec<u8> {
963        let mut x = 0x9e37_79b9_7f4a_7c15u64;
964        (0..len)
965            .map(|_| {
966                x ^= x << 13;
967                x ^= x >> 7;
968                x ^= x << 17;
969                (x >> 24) as u8
970            })
971            .collect()
972    }
973
974    fn ok(o: Outcome) -> Measured {
975        match o {
976            Outcome::Ok(m) => m,
977            Outcome::Failed(e) => panic!("failed: {e}"),
978        }
979    }
980
981    fn run(codec: &str, input: &Input, plan: &Plan) -> Measured {
982        ok(measure(
983            codec,
984            None,
985            input,
986            plan,
987            &mut FakeClock(0),
988            &mut |_, _, _| {},
989        ))
990    }
991
992    fn bench_temps() -> Vec<String> {
993        let prefix = format!("stuffr-bench-{}-", std::process::id());
994        fs::read_dir(std::env::temp_dir())
995            .unwrap()
996            .filter_map(|e| e.ok()?.file_name().into_string().ok())
997            .filter(|n| n.starts_with(&prefix))
998            .collect()
999    }
1000
1001    #[test]
1002    fn one_stream_compresses_the_concatenation() {
1003        let (a, b) = (text(20_000, 0), text(30_000, 7));
1004        let both = [a.clone(), b.clone()].concat();
1005        let p = plan(Mode::OneStream, 1);
1006        for codec in ["gzip", "zstd", "bzip2"] {
1007            let parts = run(codec, &Input::Memory(vec![a.clone(), b.clone()]), &p);
1008            let whole = run(codec, &Input::Memory(vec![both.clone()]), &p);
1009            assert_eq!(parts.bytes_out, whole.bytes_out, "{codec}");
1010        }
1011    }
1012
1013    #[test]
1014    fn per_file_sums_each_file() {
1015        let (a, b) = (text(20_000, 0), text(30_000, 7));
1016        let p = plan(Mode::PerFile, 1);
1017        let both = run("gzip", &Input::Memory(vec![a.clone(), b.clone()]), &p);
1018        let one = run("gzip", &Input::Memory(vec![a]), &p);
1019        let two = run("gzip", &Input::Memory(vec![b]), &p);
1020        assert_eq!(both.bytes_out, one.bytes_out + two.bytes_out);
1021        let stream = run(
1022            "gzip",
1023            &Input::Memory(vec![text(20_000, 0), text(30_000, 7)]),
1024            &plan(Mode::OneStream, 1),
1025        );
1026        assert_ne!(
1027            both.bytes_out, stream.bytes_out,
1028            "per-file is not one stream"
1029        );
1030        // Per file: encoder creation, one write (each file is under 64 KiB)
1031        // and `finish` to compress, one fill to decompress; summed.
1032        assert_eq!(both.compress_ns, 2 * 30);
1033        assert_eq!(both.decompress_ns, 2 * 10);
1034    }
1035
1036    /// Streamed: compression is encoder creation, one region per input
1037    /// chunk's `write_all` (each 64 KiB read falls between them) and
1038    /// `finish`; decompression one region per 64 KiB filled, the decoder
1039    /// made in the first.
1040    #[test]
1041    fn a_streamed_run_times_each_write_and_not_the_reads() {
1042        let _lock = armed_lock();
1043        let s = Scratch::new();
1044        // Two chunks: 64 KiB, then 36 KiB.
1045        let paths = vec![s.file("a", &text(100 * 1024, 6))];
1046        let files = load_input(&paths, 0).unwrap().0;
1047        let mut clock = FakeClock(0);
1048        let m = ok(measure(
1049            "gzip",
1050            None,
1051            &files,
1052            &plan(Mode::OneStream, 1),
1053            &mut clock,
1054            &mut |_, _, _| {},
1055        ));
1056        let regions = 1 + 2 + 1;
1057        // 64 KiB, then 36 KiB: the short fill is the end.
1058        let fills = 2;
1059        assert_eq!(m.compress_ns, 10 * regions);
1060        assert_eq!(m.decompress_ns, 10 * fills);
1061        assert_eq!(clock.0, 10 * 2 * (regions + fills), "no other clock reads");
1062    }
1063
1064    #[test]
1065    fn no_threads_means_no_governor() {
1066        // Pinned structurally rather than by exporting STUFFR_THREADS: the
1067        // lib tests share one environment (edition 2024 `set_var` is
1068        // unsafe for exactly that reason), and `governor_for` returns
1069        // before `resolved_budget`, the only reader of the variable, is
1070        // called.
1071        assert!(governor_for(&plan(Mode::OneStream, 1)).is_none());
1072        let two = Plan {
1073            threads: Some(2),
1074            ..plan(Mode::OneStream, 1)
1075        };
1076        assert!(governor_for(&two).is_some());
1077    }
1078
1079    /// In memory, as streamed: encoder creation, one piece per 64 KiB
1080    /// `write_all` and `finish`; then one piece per 64 KiB the decoder
1081    /// fills, the decoder made in the first. Each piece is exactly two
1082    /// clock reads, and there are no others.
1083    #[test]
1084    fn the_fake_clock_still_times_only_codec_work() {
1085        // Three chunks: 64, 64 and 32 KiB.
1086        let input = Input::Memory(vec![text(160 * 1024, 3)]);
1087        let mut clock = FakeClock(0);
1088        let m = ok(measure(
1089            "gzip",
1090            None,
1091            &input,
1092            &plan(Mode::OneStream, 1),
1093            &mut clock,
1094            &mut |_, _, _| {},
1095        ));
1096        let compress_pieces = 1 + 3 + 1;
1097        // 64, 64, then 32 KiB: the short fill is the end.
1098        let decompress_pieces = 3;
1099        assert_eq!(m.compress_ns, 10 * compress_pieces);
1100        assert_eq!(m.decompress_ns, 10 * decompress_pieces);
1101        assert_eq!(
1102            clock.0,
1103            10 * 2 * (compress_pieces + decompress_pieces),
1104            "two reads per piece, nothing else"
1105        );
1106    }
1107
1108    /// A clock whose reads a progress callback can see: the callback runs
1109    /// inside a timed region exactly when the count is odd.
1110    struct SeenClock(std::rc::Rc<std::cell::Cell<u64>>);
1111
1112    impl Clock for SeenClock {
1113        fn now(&mut self) -> u64 {
1114            let n = self.0.get();
1115            self.0.set(n + 1);
1116            10 * n
1117        }
1118    }
1119
1120    /// Splits `seen` into repeats at each `(Compress, 0, _)`, the one report
1121    /// that opens a repeat, and checks each: both phases rise to `total`,
1122    /// compression strictly after the opening report.
1123    fn check_reports(seen: &[(Phase, u64, u64)], total: u64, repeats: usize, what: &str) {
1124        let opens: Vec<usize> = (0..seen.len())
1125            .filter(|&i| matches!(seen[i], (Phase::Compress, 0, _)))
1126            .collect();
1127        assert_eq!(opens.len(), repeats, "{what}: {seen:?}");
1128        assert_eq!(opens[0], 0, "{what}: the first report opens a repeat");
1129        for (i, &start) in opens.iter().enumerate() {
1130            let end = opens.get(i + 1).copied().unwrap_or(seen.len());
1131            let one = &seen[start..end];
1132            assert!(one.iter().all(|r| r.2 == total), "{what}: {one:?}");
1133            let fed: Vec<u64> = one
1134                .iter()
1135                .filter(|r| r.0 == Phase::Compress)
1136                .map(|r| r.1)
1137                .collect();
1138            let decoded: Vec<u64> = one
1139                .iter()
1140                .filter(|r| r.0 == Phase::Decompress)
1141                .map(|r| r.1)
1142                .collect();
1143            assert!(fed.windows(2).all(|w| w[0] < w[1]), "{what}: {fed:?}");
1144            assert!(
1145                decoded.windows(2).all(|w| w[0] <= w[1]),
1146                "{what}: {decoded:?}"
1147            );
1148            assert_eq!(fed.last(), Some(&total), "{what}: compression");
1149            assert_eq!(decoded.last(), Some(&total), "{what}: decompression");
1150            assert!(fed.len() > 2 && decoded.len() > 1, "{what}: {one:?}");
1151        }
1152    }
1153
1154    #[test]
1155    fn progress_is_reported_between_pieces_only() {
1156        use std::{cell::Cell, rc::Rc};
1157        let _lock = armed_lock();
1158        let s = Scratch::new();
1159        let data = text(160 * 1024, 5);
1160        let paths = vec![
1161            s.file("a", &data[..100 * 1024]),
1162            s.file("b", &data[100 * 1024..]),
1163        ];
1164        let memory = load_input(&paths, 1 << 30).unwrap().0;
1165        let files = load_input(&paths, 0).unwrap().0;
1166        for (input, disk, name) in [
1167            (&memory, false, "memory"),
1168            (&memory, true, "memory, temp file"),
1169            (&files, false, "streamed"),
1170            (&files, true, "streamed, temp file"),
1171        ] {
1172            for mode in [Mode::OneStream, Mode::PerFile] {
1173                let reads = Rc::new(Cell::new(0));
1174                let mut clock = SeenClock(reads.clone());
1175                let mut seen = Vec::new();
1176                let mut report = |phase, done, total| {
1177                    assert_eq!(reads.get() % 2, 0, "a report inside a timed region");
1178                    seen.push((phase, done, total));
1179                };
1180                let p = plan(mode, 2);
1181                let mut go = || measure("gzip", None, input, &p, &mut clock, &mut report);
1182                let o = if disk { on_disk(go) } else { go() };
1183                ok(o);
1184                check_reports(&seen, data.len() as u64, 2, &format!("{name} {mode:?}"));
1185            }
1186        }
1187    }
1188
1189    #[test]
1190    fn repeat_reports_the_median() {
1191        let input = Input::Memory(vec![text(4096, 1)]);
1192        // Per run: encoder creation, one write, `finish`, then one fill.
1193        // Compression 30, 10, 20; decompression 7, 5, 9.
1194        let mut clock = ScriptClock {
1195            durations: vec![10, 10, 10, 7, 2, 3, 5, 5, 5, 5, 10, 9],
1196            reads: 0,
1197            t: 0,
1198        };
1199        let m = ok(measure(
1200            "gzip",
1201            None,
1202            &input,
1203            &plan(Mode::OneStream, 3),
1204            &mut clock,
1205            &mut |_, _, _| {},
1206        ));
1207        assert_eq!(m.compress_ns, 20);
1208        assert_eq!(m.decompress_ns, 7);
1209        assert_eq!(clock.reads, 24, "two reads per piece, nothing else");
1210    }
1211
1212    #[test]
1213    fn an_even_repeat_reports_the_lower_median() {
1214        let input = Input::Memory(vec![text(4096, 1)]);
1215        let mut clock = ScriptClock {
1216            // Compression 40, 10, 30, 20; decompression 1, 2, 3, 4.
1217            durations: vec![20, 10, 10, 1, 4, 3, 3, 2, 10, 10, 10, 3, 5, 5, 10, 4],
1218            reads: 0,
1219            t: 0,
1220        };
1221        let m = ok(measure(
1222            "gzip",
1223            None,
1224            &input,
1225            &plan(Mode::OneStream, 4),
1226            &mut clock,
1227            &mut |_, _, _| {},
1228        ));
1229        assert_eq!(m.compress_ns, 20);
1230        assert_eq!(m.decompress_ns, 2);
1231    }
1232
1233    #[test]
1234    fn a_rejected_level_is_a_failed_row() {
1235        let input = Input::Memory(vec![text(1024, 0)]);
1236        let mut clock = FakeClock(0);
1237        match measure(
1238            "gzip",
1239            Some(99),
1240            &input,
1241            &plan(Mode::OneStream, 1),
1242            &mut clock,
1243            &mut |_, _, _| {},
1244        ) {
1245            Outcome::Failed(msg) => assert!(msg.contains("99"), "{msg}"),
1246            Outcome::Ok(m) => panic!("level 99 accepted: {m:?}"),
1247        }
1248        assert_eq!(clock.0, 0, "nothing was timed");
1249    }
1250
1251    #[test]
1252    fn an_unknown_format_is_a_failed_row() {
1253        let input = Input::Memory(vec![text(1024, 0)]);
1254        let o = measure(
1255            "nope",
1256            None,
1257            &input,
1258            &plan(Mode::OneStream, 1),
1259            &mut FakeClock(0),
1260            &mut |_, _, _| {},
1261        );
1262        assert!(matches!(o, Outcome::Failed(_)), "{o:?}");
1263    }
1264
1265    #[test]
1266    fn streamed_input_matches_memory_input() {
1267        let _lock = armed_lock();
1268        let s = Scratch::new();
1269        let (a, b) = (text(32 * 1024, 2), text(32 * 1024, 9));
1270        let paths = vec![s.file("a", &a), s.file("b", &b)];
1271        let files = load_input(&paths, 0).unwrap().0;
1272        assert!(matches!(files, Input::Files(_)));
1273        warm(&files).unwrap();
1274        let memory = load_input(&paths, 1 << 30).unwrap().0;
1275        assert!(matches!(memory, Input::Memory(_)));
1276        for mode in [Mode::OneStream, Mode::PerFile] {
1277            let p = plan(mode, 1);
1278            for codec in ["gzip", "xz", "lz4"] {
1279                let want = run(codec, &memory, &p).bytes_out;
1280                assert_eq!(run(codec, &files, &p).bytes_out, want, "{codec}");
1281                let disk = on_disk(|| run(codec, &files, &p)).bytes_out;
1282                assert_eq!(disk, want, "{codec} via the temp file");
1283            }
1284        }
1285    }
1286
1287    #[test]
1288    fn the_output_temp_file_is_removed() {
1289        let _lock = armed_lock();
1290        let s = Scratch::new();
1291        let paths = vec![s.file("a", &text(16 * 1024, 4))];
1292        let files = load_input(&paths, 0).unwrap().0;
1293        let made = temps_made(|| {
1294            on_disk(|| run("gzip", &files, &plan(Mode::OneStream, 2)));
1295        });
1296        assert_eq!(made, 1);
1297        assert_eq!(bench_temps(), Vec::<String>::new());
1298    }
1299
1300    #[test]
1301    fn the_temp_guard_removes_on_panic() {
1302        let _lock = armed_lock();
1303        let seen = std::sync::Mutex::new(None::<PathBuf>);
1304        let r = std::panic::catch_unwind(|| {
1305            let guard = TempOut::create().unwrap();
1306            *seen.lock().unwrap() = Some(guard.path().to_path_buf());
1307            assert!(guard.path().exists());
1308            panic!("mid-run");
1309        });
1310        assert!(r.is_err());
1311        let path = seen.into_inner().unwrap().expect("the guard was created");
1312        assert!(!path.exists(), "{} survived the panic", path.display());
1313        #[cfg(unix)]
1314        assert!(
1315            crate::compat::armed_path().is_none(),
1316            "still armed after the panic"
1317        );
1318    }
1319
1320    #[cfg(unix)]
1321    #[test]
1322    fn the_temp_file_is_armed_for_signals() {
1323        let _lock = armed_lock();
1324        let guard = TempOut::create().unwrap();
1325        let path = guard.path().to_path_buf();
1326        let name = path.file_name().unwrap().to_str().unwrap().to_owned();
1327        assert!(
1328            name.starts_with(&format!("stuffr-bench-{}-", std::process::id())),
1329            "{name}"
1330        );
1331        assert_eq!(crate::compat::armed_path().as_deref(), Some(path.as_path()));
1332        drop(guard);
1333        assert!(
1334            crate::compat::armed_path().is_none(),
1335            "armed after the drop"
1336        );
1337        assert!(!path.exists());
1338    }
1339
1340    #[test]
1341    fn empty_input_measures_without_error() {
1342        for mode in [Mode::OneStream, Mode::PerFile] {
1343            for codec in crate::bench::levels::benchable_codecs() {
1344                let m = run(codec, &Input::Memory(vec![Vec::new()]), &plan(mode, 1));
1345                assert!(m.bytes_out < 1024, "{codec}: {m:?}");
1346            }
1347        }
1348    }
1349
1350    #[test]
1351    fn every_benchable_codec_round_trips_in_both_modes() {
1352        let _lock = armed_lock();
1353        let s = Scratch::new();
1354        let paths = vec![s.file("a", &text(8 * 1024, 5)), s.file("b", &text(4096, 8))];
1355        let files = load_input(&paths, 0).unwrap().0;
1356        let memory = load_input(&paths, 1 << 30).unwrap().0;
1357        for mode in [Mode::OneStream, Mode::PerFile] {
1358            for codec in crate::bench::levels::benchable_codecs() {
1359                for input in [&files, &memory] {
1360                    run(codec, input, &plan(mode, 1));
1361                    on_disk(|| run(codec, input, &plan(mode, 1)));
1362                }
1363            }
1364        }
1365    }
1366
1367    /// Counts the reads that reach the wrapped reader.
1368    struct CountReads<R> {
1369        inner: R,
1370        reads: Arc<std::sync::atomic::AtomicUsize>,
1371    }
1372
1373    impl<R: Read> Read for CountReads<R> {
1374        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
1375            self.reads.fetch_add(1, Ordering::Relaxed);
1376            self.inner.read(buf)
1377        }
1378    }
1379
1380    impl<R: Seek> Seek for CountReads<R> {
1381        fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
1382            self.inner.seek(pos)
1383        }
1384    }
1385
1386    /// LZW's decoder reads one byte per call. Through the temp file's
1387    /// source those calls must not each reach the file (one system call
1388    /// per compressed byte made streamed decoding 17-34x slower than the
1389    /// same decode from memory).
1390    #[test]
1391    fn the_temp_file_is_read_in_large_pieces_by_a_byte_reader() {
1392        let input = text(64 * 1024, 4);
1393        let codec = stuffr::registry()
1394            .codec(stuffr::FormatId::new("compress"))
1395            .expect("compress is in every default build");
1396        let out = SharedVec::with_capacity(0);
1397        let mut sink = codec
1398            .encoder(Box::new(out.clone()), &EncodeOpts::default())
1399            .unwrap();
1400        sink.write_all(&input).unwrap();
1401        sink.finish().unwrap();
1402        let bytes = out.take();
1403        let len = bytes.len() as u64;
1404        let reads = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1405        let file = CountReads {
1406            inner: Cursor::new(bytes),
1407            reads: reads.clone(),
1408        };
1409        let src = temp_source(file, len);
1410        // The shape `MemSource` gives the in-memory path.
1411        assert_eq!(
1412            src.caps(),
1413            SourceCaps {
1414                seekable: true,
1415                len: Some(len)
1416            }
1417        );
1418        let mut buf = vec![0u8; CHUNK];
1419        let (decoded, _) = decode(
1420            codec.as_ref(),
1421            &DecodeOpts::default(),
1422            Box::new(src),
1423            &mut buf,
1424            &mut FakeClock(0),
1425            &mut Progress::new(&mut |_, _, _| {}, input.len() as u64),
1426        )
1427        .unwrap();
1428        assert_eq!(decoded, input.len() as u64);
1429        let reads = reads.load(Ordering::Relaxed);
1430        assert!(
1431            reads as u64 <= len / CHUNK as u64 + 4,
1432            "{reads} reads reached the file for {len} compressed bytes"
1433        );
1434    }
1435
1436    /// The same check through `measure`: LZW reads its input a byte at a
1437    /// time, so if the production temp-file path stopped buffering, the
1438    /// reads reaching the file would be about one per compressed byte.
1439    #[test]
1440    fn measure_reads_the_temp_file_in_large_pieces() {
1441        let _lock = armed_lock();
1442        let input = Input::Memory(vec![text(64 * 1024, 4)]);
1443        FILE_READS.with(|r| r.set(0));
1444        let m = on_disk(|| run("compress", &input, &plan(Mode::OneStream, 1)));
1445        let reads = FILE_READS.with(std::cell::Cell::get) as u64;
1446        assert!(m.bytes_out > 4096, "{} compressed bytes", m.bytes_out);
1447        assert!(reads > 0, "no read reached the temp file");
1448        assert!(
1449            reads <= m.bytes_out / CHUNK as u64 + 4,
1450            "{reads} reads reached the file for {} compressed bytes",
1451            m.bytes_out
1452        );
1453    }
1454
1455    /// HOST-9: the temp file is used through the descriptor `create` made.
1456    /// Its name is swapped for a symlink to a victim file mid-run (from the
1457    /// progress callback, between timed pieces); the later repeats still
1458    /// write and read the original file, the run succeeds and the victim is
1459    /// untouched. Reopened by name, the next repeat truncated the victim.
1460    #[cfg(unix)]
1461    #[test]
1462    fn a_symlink_put_at_the_temp_name_is_never_written_through() {
1463        let _lock = armed_lock();
1464        let dir = std::env::temp_dir().join(format!("stuffr-benchvictim-{}", std::process::id()));
1465        fs::create_dir_all(&dir).unwrap();
1466        let victim = dir.join("victim");
1467        fs::write(&victim, b"precious").unwrap();
1468        let input = Input::Memory(vec![text(64 * 1024, 6)]);
1469        let mut swapped = false;
1470        let m = on_disk(|| {
1471            measure(
1472                "gzip",
1473                None,
1474                &input,
1475                &plan(Mode::OneStream, 3),
1476                &mut FakeClock(0),
1477                &mut |_, _, _| {
1478                    if swapped {
1479                        return;
1480                    }
1481                    if let Some(name) = bench_temps().into_iter().next() {
1482                        let path = std::env::temp_dir().join(name);
1483                        fs::remove_file(&path).unwrap();
1484                        std::os::unix::fs::symlink(&victim, &path).unwrap();
1485                        swapped = true;
1486                    }
1487                },
1488            )
1489        });
1490        let m = ok(m);
1491        assert!(swapped, "the progress callback never ran");
1492        assert!(m.on_disk);
1493        assert_eq!(fs::read(&victim).unwrap(), b"precious");
1494        let _ = fs::remove_dir_all(&dir);
1495    }
1496
1497    /// `compress` can write twice its input; near the limit that must not be
1498    /// reserved in memory, over the budget. Each codec gets the destination
1499    /// its own worst case allows: `compress` goes to the temp file and its
1500    /// row says so; `gzip`, whose worst case fits, stays in memory beside it.
1501    #[test]
1502    fn compress_near_the_limit_goes_to_disk_and_is_marked() {
1503        let _lock = armed_lock();
1504        let len = 16 * 1024;
1505        let input = Input::Memory(vec![noise(len)]);
1506        // The input (16384) is between limit / 4 and limit / 2. gzip's
1507        // worst case (len + len / 8 + 1 KiB = 19456) fits in that half
1508        // exactly; compress's (2 * len + 1 KiB = 33792) does not.
1509        let p = Plan {
1510            memory_limit: 2 * 19_456,
1511            ..plan(Mode::OneStream, 1)
1512        };
1513        for (codec, want_temps, want_disk) in [("compress", 1, true), ("gzip", 0, false)] {
1514            let mut m = None;
1515            let made = temps_made(|| m = Some(run(codec, &input, &p)));
1516            assert_eq!(made, want_temps, "{codec}");
1517            assert_eq!(m.unwrap().on_disk, want_disk, "{codec}");
1518        }
1519        // Streamed: the whole limit is the budget.
1520        for (codec, limit, want) in [("gzip", 19_456, 0), ("compress", 33_791, 1)] {
1521            assert_eq!(
1522                destination(codec, len as u64, limit, true) == Dest::TempFile,
1523                want == 1,
1524                "{codec} streamed at {limit}"
1525            );
1526        }
1527        assert_eq!(
1528            destination("compress", len as u64, 33_792, true),
1529            Dest::Memory
1530        );
1531    }
1532
1533    #[test]
1534    fn the_output_goes_to_memory_whenever_its_worst_case_fits() {
1535        // 0 bytes in: worst case 1 KiB, which fits under 2048 / 2.
1536        // Beside input in memory: half the limit.
1537        assert_eq!(destination("gzip", 0, 2048, false), Dest::Memory);
1538        assert_eq!(destination("gzip", 0, 2047, false), Dest::TempFile);
1539        // 8000 + 1000 + 1024.
1540        assert_eq!(destination("gzip", 8000, 2 * 10_024, false), Dest::Memory);
1541        assert_eq!(
1542            destination("gzip", 8000, 2 * 10_024 - 1, false),
1543            Dest::TempFile
1544        );
1545        assert_eq!(
1546            destination("gzip", u64::MAX, u64::MAX, false),
1547            Dest::TempFile
1548        );
1549        // Streamed input holds no copy: the whole limit.
1550        assert_eq!(destination("gzip", 8000, 10_024, true), Dest::Memory);
1551        assert_eq!(destination("gzip", 8000, 10_023, true), Dest::TempFile);
1552        assert_eq!(
1553            destination("gzip", u64::MAX, u64::MAX, true),
1554            Dest::TempFile
1555        );
1556    }
1557
1558    /// Streamed input whose output fits in the whole limit stays off the
1559    /// disk, one stream included; only output past the limit makes a temp
1560    /// file.
1561    #[test]
1562    fn streamed_input_compresses_to_memory_when_the_output_fits() {
1563        let _lock = armed_lock();
1564        let s = Scratch::new();
1565        let (a, b) = (text(16 * 1024, 1), text(16 * 1024, 2));
1566        let paths = vec![s.file("a", &a), s.file("b", &b)];
1567        // 32 KiB is over 40000 / 2, so the input is streamed. One stream's
1568        // worst case (32768 + 4096 + 1024 = 37888) fits under 40000, not
1569        // under 37887; each file's (19456) fits under either.
1570        let files = load_input(&paths, 40_000).unwrap().0;
1571        assert!(matches!(files, Input::Files(_)));
1572        let memory = Input::Memory(vec![a, b]);
1573        for (mode, limit, want_temps) in [
1574            (Mode::PerFile, 40_000, 0),
1575            (Mode::OneStream, 40_000, 0),
1576            (Mode::OneStream, 37_888, 0),
1577            (Mode::OneStream, 37_887, 1),
1578        ] {
1579            let p = Plan {
1580                memory_limit: limit,
1581                ..plan(mode, 1)
1582            };
1583            let mut got = None;
1584            let made = temps_made(|| got = Some(run("gzip", &files, &p)));
1585            assert_eq!(made, want_temps, "{mode:?} at {limit}");
1586            let want = run("gzip", &memory, &plan(mode, 1)).bytes_out;
1587            assert_eq!(got.unwrap().bytes_out, want, "{mode:?}");
1588        }
1589    }
1590
1591    /// In-memory input whose output does not fit goes to the temp file too,
1592    /// for each codec whose own worst case is over the budget.
1593    #[test]
1594    fn memory_input_past_the_output_limit_uses_the_temp_file() {
1595        let _lock = armed_lock();
1596        let input = Input::Memory(vec![text(16 * 1024, 5)]);
1597        // The input fits under 32768 / 2; its worst case does not.
1598        let p = Plan {
1599            memory_limit: 32 * 1024,
1600            ..plan(Mode::OneStream, 1)
1601        };
1602        for codec in ["gzip", "compress"] {
1603            let made = temps_made(|| {
1604                run(codec, &input, &p);
1605            });
1606            assert_eq!(made, 1, "{codec}");
1607        }
1608    }
1609
1610    /// The reserved `Vec` must not reallocate inside a timed region, even
1611    /// on input no codec can shrink, at any level.
1612    #[test]
1613    fn the_reserved_room_holds_incompressible_output() {
1614        let input = noise(64 * 1024);
1615        let reg = stuffr::registry();
1616        for name in crate::bench::levels::benchable_codecs() {
1617            let codec = reg.codec(stuffr::FormatId::new(name)).unwrap();
1618            let sweep = crate::bench::levels::levels_of(name).unwrap().sweep;
1619            let levels: Vec<Option<i32>> = if sweep.is_empty() {
1620                vec![None]
1621            } else {
1622                sweep.iter().map(|&l| Some(l)).collect()
1623            };
1624            for level in levels {
1625                let room = output_room(name, input.len() as u64).unwrap() as usize;
1626                let shared = SharedVec::with_capacity(room);
1627                let cap = shared.capacity();
1628                let dst = BufWriter::with_capacity(CHUNK, shared.clone());
1629                let opts = EncodeOpts {
1630                    level,
1631                    ..Default::default()
1632                };
1633                let mut sink = codec.encoder(Box::new(dst), &opts).unwrap();
1634                sink.write_all(&input).unwrap();
1635                sink.finish().unwrap();
1636                assert_eq!(shared.capacity(), cap, "{name} {level:?} reallocated");
1637                assert!(shared.take().len() <= room, "{name} {level:?}");
1638            }
1639        }
1640    }
1641
1642    /// A log still being written: the rows bench it as it was at load.
1643    #[test]
1644    fn a_file_that_grows_after_load_is_benched_at_its_loaded_length() {
1645        let _lock = armed_lock();
1646        let s = Scratch::new();
1647        let original = text(40 * 1024, 3);
1648        let path = s.file("grow.log", &original);
1649        let (files, bytes_in) = load_input(std::slice::from_ref(&path), 0).unwrap();
1650        assert!(matches!(files, Input::Files(_)));
1651        assert_eq!(bytes_in, original.len() as u64);
1652        let mut f = fs::OpenOptions::new().append(true).open(&path).unwrap();
1653        f.write_all(&text(10 * 1024, 9)).unwrap();
1654        drop(f);
1655        let snapshot = Input::Memory(vec![original]);
1656        for mode in [Mode::OneStream, Mode::PerFile] {
1657            let p = plan(mode, 1);
1658            for codec in ["gzip", "compress"] {
1659                // Ok: the round trip gave back exactly the loaded length.
1660                let grown = run(codec, &files, &p);
1661                assert_eq!(
1662                    grown.bytes_out,
1663                    run(codec, &snapshot, &p).bytes_out,
1664                    "{codec}: the appended bytes were benched"
1665                );
1666            }
1667        }
1668    }
1669
1670    #[test]
1671    fn load_input_splits_at_half_the_memory_limit() {
1672        let s = Scratch::new();
1673        let paths = vec![s.file("a", &[1; 100]), s.file("b", &[2; 100])];
1674        match load_input(&paths, 400).unwrap() {
1675            (Input::Memory(parts), 200) => {
1676                assert_eq!(parts, vec![vec![1; 100], vec![2; 100]])
1677            }
1678            other => panic!("200 bytes fit in 400 / 2: {other:?}"),
1679        }
1680        assert!(matches!(
1681            load_input(&paths, 399).unwrap(),
1682            (Input::Files(_), 200)
1683        ));
1684    }
1685
1686    /// Restores a file's mode, so `Scratch` can remove it.
1687    #[cfg(unix)]
1688    struct Restore(PathBuf);
1689
1690    #[cfg(unix)]
1691    impl Drop for Restore {
1692        fn drop(&mut self) {
1693            use std::os::unix::fs::PermissionsExt;
1694            let _ = fs::set_permissions(&self.0, fs::Permissions::from_mode(0o600));
1695        }
1696    }
1697
1698    /// A file that cannot be opened is a usage error naming it, whether
1699    /// the input is loaded into memory or streamed (opened by `warm`).
1700    #[cfg(unix)]
1701    #[test]
1702    fn an_input_that_cannot_be_opened_is_a_usage_error_naming_it() {
1703        use std::os::unix::fs::PermissionsExt;
1704        let s = Scratch::new();
1705        let locked = s.file("locked.log", b"secret");
1706        let _restore = Restore(locked.clone());
1707        fs::set_permissions(&locked, fs::Permissions::from_mode(0o000)).unwrap();
1708        // Root opens it anyway; there is nothing to test.
1709        if File::open(&locked).is_ok() {
1710            return;
1711        }
1712        let want = format!("{}: cannot read it (", locked.display());
1713        let usage = |r: stuffr::Result<()>| match r {
1714            Err(Error::Usage(m)) => m,
1715            other => panic!("expected a usage error, got {other:?}"),
1716        };
1717        let paths = vec![locked.clone()];
1718        let m = usage(load_input(&paths, 1 << 30).map(drop));
1719        assert!(m.starts_with(&want), "in memory: {m}");
1720        // Streamed: loading only reads the size; the warm-up opens it.
1721        let (files, _) = load_input(&paths, 0).unwrap();
1722        assert!(matches!(files, Input::Files(_)));
1723        let m = usage(warm(&files));
1724        assert!(m.starts_with(&want), "streamed: {m}");
1725    }
1726}