open-standards/wai/wai-pack/tests/pack_alloc.rs

//! What the packer holds does not grow with the source (packaging §11.4): a source eight
//! times as long, read from a file and packed into a directory, allocates at its peak what
//! the short one does, give or take the records of its pieces, though it is millions of
//! bytes longer. Its own binary: it installs a counting global allocator.

use std::alloc::{GlobalAlloc, Layout, System};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};

struct Peak;
static LIVE: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Peak {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(l) };
        if !p.is_null() {
            let now = LIVE.fetch_add(l.size(), Ordering::SeqCst) + l.size();
            PEAK.fetch_max(now, Ordering::SeqCst);
        }
        p
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) };
        LIVE.fetch_sub(l.size(), Ordering::SeqCst);
    }
}

#[global_allocator]
static A: Peak = Peak;

fn corpus() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("../pack-conformance")
}

/// The golden `video-base-01`'s fixture and supplied floor played `k` times over, written
/// below `dir`: the source as one Y4M, and the floor as one stream of the golden's units
/// repeated, so its key frames fall on the repeated plan; the golden's policy pinned to
/// that floor. Returns the policy, the source's path and the floor's directory.
fn repeated(dir: &Path, k: usize) -> (Vec<u8>, PathBuf, PathBuf) {
    let y = wai_pack::fixture::generate("video-base-01").unwrap();
    let head = y.iter().position(|&c| c == b'\n').unwrap() + 1;
    let mut src = y[..head].to_vec();
    for _ in 0..k {
        src.extend_from_slice(&y[head..]);
    }
    let units: Vec<Vec<u8>> = wai_pack::av1::temporal_units(&std::fs::read(corpus().join("packs/video-base-01/floor/v.obu")).unwrap()).unwrap().iter().map(|t| t.bytes()).collect();
    let floor_dir = dir.join(format!("floor-{k}"));
    std::fs::create_dir_all(&floor_dir).unwrap();
    std::fs::write(floor_dir.join("v.obu"), units.iter().cycle().take(k * units.len()).flatten().copied().collect::<Vec<u8>>()).unwrap();
    let src_path = dir.join(format!("v-{k}.y4m"));
    std::fs::write(&src_path, &src).unwrap();
    let listing = wai_pack::floor::listing_sha256(&wai::pin::read_dir_set(&floor_dir).unwrap()).unwrap();
    let p = String::from_utf8(std::fs::read(corpus().join("packs/video-base-01/policy.toml")).unwrap()).unwrap();
    let at = p.find("supplied_listing = \"").unwrap() + "supplied_listing = \"".len();
    (format!("{}{listing}{}", &p[..at], &p[at + 64..]).into_bytes(), src_path, floor_dir)
}

/// Pack into `out` from files, and give the peak allocated beyond what was live before.
fn pack_peak(policy: &[u8], src: &Path, floor: &Path, out: &Path) -> (wai::package::PackageReceiptV2, usize) {
    let mut o = wai_pack::pack::PackOpts::with_floor(ed25519_dalek::SigningKey::from_bytes(&[0x5a; 32]), "did:key:test", wai_pack::files::FileSet::Dir(floor.to_path_buf()));
    let mut sources = BTreeMap::new();
    sources.insert("v".to_string(), wai_pack::files::Input::Path(src.to_path_buf()));
    let base = LIVE.load(Ordering::SeqCst);
    PEAK.store(base, Ordering::SeqCst);
    let mut sink = wai_pack::files::DirSink::new(out).unwrap();
    let r = wai_pack::pack::pack_into(policy, &sources, &mut o, &mut sink).unwrap();
    sink.keep();
    (r, PEAK.load(Ordering::SeqCst) - base)
}

#[test]
fn a_longer_source_takes_no_more_memory() {
    let dir = std::env::temp_dir().join(format!("wai-pack-alloc-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let (p1, s1, f1) = repeated(&dir, 1);
    let (p8, s8, f8) = repeated(&dir, 8);
    let (r1, peak1) = pack_peak(&p1, &s1, &f1, &dir.join("out-1"));
    let (r8, peak8) = pack_peak(&p8, &s8, &f8, &dir.join("out-8"));
    // The long pack is the short one eight times: eight times the entries, its repeated
    // pieces stored once.
    assert_eq!(r8.total_entries, 8 * r1.total_entries);
    assert_eq!(r8.unique_objects, r1.unique_objects, "every repetition's pieces are the first's objects");
    wai_pack::verify::verify_dir(&dir.join("out-8")).unwrap();
    let grown = std::fs::metadata(&s8).unwrap().len() - std::fs::metadata(&s1).unwrap().len();
    assert!(grown > 4_000_000, "the long source is {grown} bytes longer");
    // Its records grow with its pieces (a few hundred bytes each); nothing grows with its
    // frames.
    eprintln!("peak allocated: {} entries {peak1} B; {} entries {peak8} B", r1.total_entries, r8.total_entries);
    let slack = 256 * 1024;
    assert!(peak8 < peak1 + slack, "the long pack's peak {peak8} B is past the short one's {peak1} B by more than {slack} B");
    let _ = std::fs::remove_dir_all(&dir);
}