smarm/tests/stack_shrink.rs

//! RFC 019 commit 3 — park-path stack shrink, observed from the outside.
//!
//! The one integration-level claim of the shrink machinery: an actor that
//! spikes deep, returns shallow, and then parks past the cooldown gets its
//! dead span MADV_FREE'd — visible as `LazyFree` in `/proc/self/smaps`
//! within the stack's address range — while everything live survives.
//!
//! The high-water mark is *sampled* at context-save, so the spike yields
//! once at max depth to guarantee a sample there (in production, preemption
//! provides the quasi-random samples; a test must not rely on luck).

use smarm::runtime::{Config, SHRINK_COOLDOWN, SHRINK_THRESHOLD};
use smarm::{actor_info, channel, spawn, spawn_with, yield_now, ActorState, SpawnOpts};

/// Burn ~`frames` × 4 KiB of stack, yielding once at the bottom so the
/// context-save samples `sp` at max depth.
#[inline(never)]
fn burn_stack_yielding(frames: usize) -> u64 {
    let mut local = [0u8; 4096];
    local[0] = frames as u8;
    let below = if frames == 0 {
        yield_now();
        0
    } else {
        burn_stack_yielding(frames - 1)
    };
    std::hint::black_box(&mut local);
    below.wrapping_add(local[0] as u64)
}

/// Sum the `LazyFree:` kB of every smaps mapping intersecting [lo, hi).
fn lazy_free_bytes_in(lo: usize, hi: usize) -> usize {
    let smaps = std::fs::read_to_string("/proc/self/smaps").unwrap();
    let mut total_kb = 0usize;
    let mut in_range = false;
    for line in smaps.lines() {
        if let Some((range, _)) = line.split_once(' ') {
            if let Some((a, b)) = range.split_once('-') {
                if let (Ok(start), Ok(end)) =
                    (usize::from_str_radix(a, 16), usize::from_str_radix(b, 16))
                {
                    in_range = start < hi && end > lo;
                    continue;
                }
            }
        }
        if in_range {
            if let Some(rest) = line.strip_prefix("LazyFree:") {
                let kb: usize = rest.trim().trim_end_matches(" kB").trim().parse().unwrap();
                total_kb += kb;
            }
        }
    }
    total_kb * 1024
}

#[test]
fn spike_then_parks_marks_lazyfree_and_keeps_live_data() {
    // Single scheduler: the controller can gate on the worker being Parked.
    let rt = smarm::runtime::init(Config::exact(1));
    rt.run(|| {
        let (park_tx, park_rx) = channel::<()>();
        let (done_tx, done_rx) = channel::<(usize, u64)>();

        let spike = 768 * 4096; // ~3 MiB, well past SHRINK_THRESHOLD
        assert!(spike > SHRINK_THRESHOLD);

        let worker = spawn_with(
            SpawnOpts { stack_reserve: Some(8 * 1024 * 1024), ..SpawnOpts::default() },
            move || {
                // Live data that must survive the shrink, and an anchor
                // address inside the stack for the smaps scan.
                let live = [0xA5u8; 64];
                let anchor = live.as_ptr() as usize;

                // Spike: ~3 MiB deep, sampled at the bottom, unwound.
                std::hint::black_box(burn_stack_yielding(768));

                // Park past the cooldown. Each recv on the drained inbox is
                // one park; the controller sends only when it sees us Parked.
                for _ in 0..(SHRINK_COOLDOWN + 8) {
                    park_rx.recv().unwrap();
                }

                // Measure from inside: the stack spans ≤ 8 MiB below anchor.
                let lazy = lazy_free_bytes_in(anchor - 8 * 1024 * 1024, anchor + 4096);
                let checksum = live.iter().map(|&b| b as u64).sum();
                done_tx.send((lazy, checksum)).unwrap();
            },
        );

        let wpid = worker.pid();
        for _ in 0..(SHRINK_COOLDOWN + 8) {
            // Gate: send only once the worker is genuinely parked so every
            // round is a real park-on-empty-mailbox.
            loop {
                match actor_info(wpid) {
                    Some(info) if info.state == ActorState::Parked => break,
                    Some(_) => yield_now(),
                    None => panic!("worker died early"),
                }
            }
            park_tx.send(()).unwrap();
        }

        let (lazy, checksum) = done_rx.recv().unwrap();
        // The spike was ~3 MiB; demand at least 2 MiB marked to leave slack
        // for the redzone, rounding, and pages the unwind re-dirtied.
        assert!(
            lazy >= 2 * 1024 * 1024,
            "expected ≥ 2 MiB LazyFree in the stack range, got {} bytes",
            lazy
        );
        assert_eq!(checksum, 64 * 0xA5u64, "live stack data corrupted by shrink");
        worker.join().unwrap();
    });
}

/// Steady-state actors must never pay the syscall: an actor that parks a lot
/// but never spikes past the threshold ends with zero LazyFree in its stack.
#[test]
fn shallow_actor_never_shrinks() {
    let rt = smarm::runtime::init(Config::exact(1));
    rt.run(|| {
        let (park_tx, park_rx) = channel::<()>();
        let (done_tx, done_rx) = channel::<usize>();

        let worker = spawn(move || {
            let probe = 0u8;
            let anchor = &probe as *const u8 as usize;
            for _ in 0..(SHRINK_COOLDOWN + 8) {
                park_rx.recv().unwrap();
            }
            done_tx.send(lazy_free_bytes_in(anchor - 64 * 1024, anchor + 4096)).unwrap();
        });

        let wpid = worker.pid();
        for _ in 0..(SHRINK_COOLDOWN + 8) {
            loop {
                match actor_info(wpid) {
                    Some(info) if info.state == ActorState::Parked => break,
                    Some(_) => yield_now(),
                    None => panic!("worker died early"),
                }
            }
            park_tx.send(()).unwrap();
        }

        assert_eq!(done_rx.recv().unwrap(), 0, "steady-state actor was shrunk");
        worker.join().unwrap();
    });
}