open-standards/proof/proof-rs/crates/proof-cascade/tests/process_energy.rs

//! This process's CPU energy counter as a meter, on the real hardware.
//!
//! Every test here checks the meter against an independent oracle: the tests'
//! own binding of `proc_pid_rusage(getpid(), 6)`, read into a raw buffer at
//! byte 336, with their own loop and their own literals. The oracle watches
//! the counter for one second while a helper thread burns CPU and calls it
//! moving only if it changed at no fewer than 5 of its reads and at 80% of
//! them. Where the oracle sees the counter move, the meter must have selected
//! it and must meter with it; where it does not, the meter must not. No branch
//! passes without asserting.
//!
//! The process counter counts everything this process's threads burn, so the
//! tests in this file run one at a time (`SERIAL`), and nothing else here runs
//! while one of them meters.

use proof_cascade::meter::{measure_marginal, meter_source, process_energy_status, CounterStatus, MeterSource};
use proof_schema::Provenance;
use std::process::Child;
use std::sync::{Mutex, MutexGuard, OnceLock};
use std::time::{Duration, Instant};

static SERIAL: Mutex<()> = Mutex::new(());

fn serial() -> MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(|e| e.into_inner())
}

/// What the oracle saw of the counter over its window.
#[derive(Debug, Clone, Copy)]
struct Oracle {
    moved: bool,
    reads: u32,
    changes: u32,
    advanced_nj: u64,
}

#[cfg(target_os = "macos")]
mod ffi {
    use std::os::raw::c_int;
    unsafe extern "C" {
        fn proc_pid_rusage(pid: c_int, flavor: c_int, buffer: *mut [u64; 58]) -> c_int;
        fn getpid() -> c_int;
    }

    /// `ri_energy_nj` of `proc_pid_rusage(getpid(), RUSAGE_INFO_V6 = 6)`: a
    /// 464-byte buffer, read at byte 336 (its 43rd `uint64_t`).
    pub fn energy_nj() -> Option<u64> {
        let mut buf = [0u64; 58];
        // SAFETY: buf is 464 writable bytes, the size of rusage_info_v6.
        (unsafe { proc_pid_rusage(getpid(), 6, &mut buf) } == 0).then_some(buf[42])
    }
}

/// The oracle's verdict, taken once for this test binary (under `SERIAL`).
fn oracle() -> Oracle {
    static ORACLE: OnceLock<Oracle> = OnceLock::new();
    *ORACLE.get_or_init(watch_counter)
}

#[cfg(target_os = "macos")]
fn watch_counter() -> Oracle {
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::Arc;
    let Some(first) = ffi::energy_nj() else {
        return Oracle { moved: false, reads: 0, changes: 0, advanced_nj: 0 };
    };
    let stop = Arc::new(AtomicBool::new(false));
    let burner = {
        let stop = stop.clone();
        std::thread::spawn(move || {
            let mut acc = 0u64;
            while !stop.load(Ordering::Relaxed) {
                acc = std::hint::black_box(acc.wrapping_add(1));
            }
        })
    };
    let started = Instant::now();
    let (mut prev, mut reads, mut changes) = (first, 0u32, 0u32);
    while started.elapsed() < Duration::from_secs(1) {
        std::thread::sleep(Duration::from_millis(25));
        let now = ffi::energy_nj().expect("the oracle's read failed after its first succeeded");
        reads += 1;
        changes += u32::from(now != prev);
        prev = now;
    }
    stop.store(true, Ordering::Relaxed);
    let _ = burner.join();
    let moved = changes >= 5 && changes * 5 >= reads * 4;
    Oracle { moved, reads, changes, advanced_nj: prev - first }
}

#[cfg(not(target_os = "macos"))]
fn watch_counter() -> Oracle {
    Oracle { moved: false, reads: 0, changes: 0, advanced_nj: 0 }
}

/// How long [`Running`] waits for a process it has killed to be reaped. `SIGKILL` ends a
/// process at once unless it is stuck in the kernel; a drop that waited on such a one without
/// a bound would hang the test binary instead of failing it.
const REAPED_WITHIN: Duration = Duration::from_secs(10);

/// A process a test started, killed and reaped when it is dropped: when the test is done with
/// it, or while a panic unwinds through it. A [`Child`] is not killed when it is dropped, so a
/// busy child started before a later one failed to start was left spinning a core after the
/// test binary had ended.
struct Running(Child);

impl Running {
    /// Kills the process and reaps it, waiting up to [`REAPED_WITHIN`]: whether it was reaped.
    /// A process already reaped is not signalled again (`Child::kill` does nothing then, as its
    /// pid may be another process's by now).
    fn stop(&mut self) -> bool {
        let _ = self.0.kill();
        let t = Instant::now();
        loop {
            match self.0.try_wait() {
                Ok(Some(_)) => return true,
                Ok(None) if t.elapsed() < REAPED_WITHIN => std::thread::sleep(Duration::from_millis(10)),
                _ => return false,
            }
        }
    }
}

impl Drop for Running {
    fn drop(&mut self) {
        if !self.stop() {
            eprintln!("process {} was killed and not reaped within {} s", self.0.id(), REAPED_WITHIN.as_secs());
        }
    }
}

/// Real, cheap, pure CPU work: a fixed fold.
fn cheap() {
    std::hint::black_box((0..20_000u64).fold(0u64, |a, b| a.wrapping_mul(31).wrapping_add(b)));
}

/// 2 ms of computing, by the wall clock.
fn compute_2ms() {
    let t = Instant::now();
    let mut acc = 0u64;
    while t.elapsed() < Duration::from_millis(2) {
        acc = std::hint::black_box(acc.wrapping_add(1));
    }
}

/// 2 ms of sleeping.
fn sleep_2ms() {
    std::thread::sleep(Duration::from_millis(2));
}

/// The counter is selected exactly when the oracle sees it keep moving: live
/// and first in line where it moves, never selected where it does not.
#[test]
fn the_process_counter_is_live_exactly_when_the_oracle_sees_it_move() {
    let _serial = serial();
    let status = process_energy_status();
    let o = oracle();
    let source = meter_source();
    eprintln!("process liveness: status={status:?} source={source:?} oracle={o:?}");
    if cfg!(target_os = "macos") {
        assert_eq!(
            status == CounterStatus::Live,
            o.moved,
            "the probe says {status:?}; the oracle saw the counter change at {} of {} reads (+{} nJ)",
            o.changes,
            o.reads,
            o.advanced_nj
        );
        if o.moved {
            assert_eq!(source, Some(MeterSource::ProcessEnergy), "a live process counter was not selected first");
        } else {
            assert_ne!(source, Some(MeterSource::ProcessEnergy), "a process counter the oracle saw hold was selected");
        }
    } else {
        assert_eq!(status, CounterStatus::Unavailable, "a process counter off macOS");
        assert_ne!(source, Some(MeterSource::ProcessEnergy));
    }
}

/// A fixed cheap closure meters as a non-zero `HwShunt` figure from the
/// process counter, roughly stable: three calls within a factor of 3.
#[test]
fn a_fixed_cheap_closure_meters_as_a_stable_hwshunt_figure() {
    let _serial = serial();
    if meter_source() == Some(MeterSource::ProcessEnergy) {
        let figures: Vec<_> = (0..3).map(|_| measure_marginal(cheap).expect("a live process counter gave no figure")).collect();
        eprintln!("process stability: cheap closure, three calls: {figures:?}");
        for m in &figures {
            assert_eq!(m.source, MeterSource::ProcessEnergy, "{m:?}");
            assert!(m.uj > 0, "a live process counter metered real work as zero: {m:?}");
            let e = m.energy();
            assert_eq!((e.measured_uj, e.provenance), (Some(m.uj), Provenance::HwShunt), "{m:?}");
            assert!(e.check_provenance().is_ok(), "{e:?}");
        }
        let (lo, hi) = (figures.iter().map(|m| m.uj).min().unwrap(), figures.iter().map(|m| m.uj).max().unwrap());
        assert!(hi <= 3 * lo, "three calls on the same closure spread from {lo} to {hi} uJ, more than 3x");
    } else {
        let o = oracle();
        assert!(!o.moved, "the oracle saw the process counter move ({o:?}), yet it was not selected");
        let m = measure_marginal(cheap);
        assert_ne!(m.map(|m| m.source), Some(MeterSource::ProcessEnergy), "a process-counter figure from a counter not selected");
    }
}

/// Sleeping is not computing: a closure that sleeps 2 ms meters at under a
/// tenth of one that computes for 2 ms — and stays there while four busy
/// child processes burn CPU beside it, which the counter excludes.
#[test]
fn a_closure_that_sleeps_meters_far_less_than_one_that_computes() {
    let _serial = serial();
    if meter_source() == Some(MeterSource::ProcessEnergy) {
        let compute = measure_marginal(compute_2ms).expect("computing gave no figure");
        let sleep = measure_marginal(sleep_2ms).expect("sleeping gave no figure");
        // Each held as it starts: if a later one cannot start, those already running are
        // killed as the panic unwinds through them.
        let busy: Vec<_> = (0..4)
            .map(|_| {
                Running(
                    std::process::Command::new("/bin/sh")
                        .args(["-c", "while :; do :; done"])
                        .spawn()
                        .expect("a busy child process starts"),
                )
            })
            .collect();
        std::thread::sleep(Duration::from_millis(100));
        let beside_children = measure_marginal(sleep_2ms);
        // Each killed and reaped.
        drop(busy);
        let beside_children = beside_children.expect("sleeping beside busy children gave no figure");
        eprintln!(
            "process sleep vs compute: compute 2 ms = {compute:?}; sleep 2 ms = {sleep:?}; sleep 2 ms beside 4 busy \
             children = {beside_children:?}"
        );
        assert!(compute.uj > 0, "computing metered as zero: {compute:?}");
        assert!(sleep.uj * 10 < compute.uj, "sleeping metered {} uJ a call, computing {} uJ", sleep.uj, compute.uj);
        assert!(
            beside_children.uj * 10 < compute.uj,
            "sleeping beside busy child processes metered {} uJ a call, computing {} uJ: other processes were counted",
            beside_children.uj,
            compute.uj
        );
    } else {
        let o = oracle();
        assert!(!o.moved, "the oracle saw the process counter move ({o:?}), yet it was not selected");
        assert_ne!(measure_marginal(sleep_2ms).map(|m| m.source), Some(MeterSource::ProcessEnergy));
    }
}