smarm/tests/stack_recycle.rs

//! RFC 019 commit 5 — pool recycle zaps a dead stack down to its retained
//! entry end, observed from the outside.
//!
//! A default-shaped stack that spiked deep and then died must not carry its
//! spike into the pool as resident RSS: `recycle_stack` DONTNEEDs everything
//! below the top `RECYCLE_RETAIN` bytes before pushing. The zap is
//! synchronous on the death path, so the drop is immediate — but the death
//! path itself races the observer's `join` return, hence the brief poll.
//!
//! Residency is measured with `mincore`, not smaps: a neighboring rw anon
//! mapping can land flush against the stack top and the kernel merges the
//! VMAs (observed under the full test run), so per-mapping smaps fields
//! over-count. The PROT_NONE guard below can never merge, so the usable
//! base is exactly the anchor VMA's start, and `mincore` counts pages
//! within [usable_base, usable_base + reserve) regardless of merging.

use smarm::runtime::{Config, RECYCLE_RETAIN};
use smarm::{channel, spawn, yield_now};

const RESERVE: usize = 4 * 1024 * 1024;

/// Burn ~`frames` × 4 KiB of stack, dirtying every frame.
#[inline(never)]
fn burn_stack(frames: usize) -> u64 {
    let mut local = [0u8; 4096];
    local[0] = frames as u8;
    let below = if frames == 0 {
        0
    } else {
        burn_stack(frames - 1)
    };
    std::hint::black_box(&mut local);
    below.wrapping_add(local[0] as u64)
}

/// Resident-page count over [lo, lo + len) via mincore (len page-aligned).
fn resident_pages(lo: usize, len: usize) -> usize {
    let page = 4096;
    let mut vec = vec![0u8; len / page];
    let ret = unsafe { libc::mincore(lo as *mut libc::c_void, len, vec.as_mut_ptr()) };
    assert_eq!(
        ret,
        0,
        "mincore failed: {}",
        std::io::Error::last_os_error()
    );
    vec.iter().filter(|&&b| b & 1 != 0).count()
}

/// The [start, end) of the VMA containing `addr`.
fn vma_containing(addr: usize) -> (usize, usize) {
    let maps = std::fs::read_to_string("/proc/self/maps").unwrap();
    for line in maps.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))
                {
                    if start <= addr && addr < end {
                        return (start, end);
                    }
                }
            }
        }
    }
    panic!("no VMA contains {addr:#x}");
}

fn vma_exists(addr: usize) -> bool {
    let maps = std::fs::read_to_string("/proc/self/maps").unwrap();
    for line in maps.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))
                {
                    if start <= addr && addr < end {
                        return true;
                    }
                }
            }
        }
    }
    false
}

#[test]
fn recycle_zaps_dead_stack_down_to_retain() {
    // Default reserve raised so the pool holds big stacks (default-shaped ⇒
    // pooled) and the zap has something to bite; single scheduler.
    let rt = smarm::runtime::init(Config::exact(1).stack_reserve(RESERVE));
    rt.run(|| {
        let (tx, rx) = channel::<usize>();

        let h = spawn(move || {
            let probe = 0u8;
            let anchor = &probe as *const u8 as usize;
            // The guard below is PROT_NONE and can never merge with the
            // usable region, so the anchor VMA's start IS the usable base.
            let (vlo, _) = vma_containing(anchor);
            // Dirty ~3 MiB of the 4 MiB reserve, then die.
            std::hint::black_box(burn_stack(768));
            tx.send(vlo).unwrap();
        });

        let usable_base = rx.recv().unwrap();
        h.join().unwrap();

        // The zap span is everything below the retained entry end. DONTNEED
        // on private anon discards synchronously and unconditionally, so
        // this must go to exactly zero resident pages; the poll only covers
        // the death path racing join's return.
        let zap_len = RESERVE - RECYCLE_RETAIN;
        let mut resident = usize::MAX;
        for _ in 0..10_000 {
            resident = resident_pages(usable_base, zap_len);
            if resident == 0 {
                break;
            }
            yield_now();
        }
        assert_eq!(
            resident, 0,
            "recycled stack's zap span still resident: {resident} pages in \
             [{usable_base:#x}, +{zap_len:#x})"
        );
        // Pooled, not munmapped: the mapping must still be there.
        assert!(
            vma_exists(usable_base),
            "default-shaped stack was unmapped instead of pooled"
        );
    });
}