open-standards/wai/wai-quantum/examples/qec_memory/receipt.rs

//! Energy, time, peak memory and the signed receipt each run writes.

use serde_json::{Value, json};
use std::path::Path;
use std::time::Instant;
use wai_quantum::ed25519_dalek::SigningKey;

/// Run `f`, timing it and reading an energy counter around it: the result,
/// the wall time in seconds, and the energy figure (or why there is none).
/// On macOS the counter is this process's own CPU energy; on Linux, the
/// processor packages' counters, which count every process on the host (the
/// figure says so, and gives the host's load when the run began).
pub fn metered<T>(f: impl FnOnce() -> T) -> (T, f64, Value) {
    if cfg!(target_os = "linux") {
        return metered_packages(f);
    }
    metered_process(f)
}

fn metered_packages<T>(f: impl FnOnce() -> T) -> (T, f64, Value) {
    use std::sync::Mutex;
    use std::sync::atomic::{AtomicBool, Ordering};
    use wai_quantum::engine::energy::{Accumulated, Packages, package_figure};
    let load = std::fs::read_to_string("/proc/loadavg").ok().and_then(|s| s.split_whitespace().next().map(str::to_string));
    let packages = match Packages::open() {
        Ok(p) => p,
        Err(why) => {
            let started = Instant::now();
            let out = f();
            return (out, started.elapsed().as_secs_f64(), json!({ "unmetered": why }));
        }
    };
    let Some(first) = packages.read() else {
        let started = Instant::now();
        let out = f();
        return (out, started.elapsed().as_secs_f64(), json!({ "unmetered": "a package counter did not read" }));
    };
    let acc = Mutex::new(Accumulated::new(&packages, first));
    let done = AtomicBool::new(false);
    let started = Instant::now();
    let (out, ok) = std::thread::scope(|s| {
        // Read every 250 ms, so that no counter wraps between two reads.
        s.spawn(|| {
            while !done.load(Ordering::Relaxed) {
                std::thread::sleep(std::time::Duration::from_millis(250));
                let _ = acc.lock().unwrap().read();
            }
        });
        let out = f();
        done.store(true, Ordering::Relaxed);
        let ok = acc.lock().unwrap().read().is_some();
        (out, ok)
    });
    let secs = started.elapsed().as_secs_f64();
    let mut figure = if ok { package_figure(Some(&acc.lock().unwrap().reading(0)), "") } else { json!({ "unmetered": "a package counter did not read at the end of the run" }) };
    if figure.get("joules_micro").is_some() {
        figure["covers"] = json!("the energy of the host's processor packages over the run: every process on the host in that window, this one among them, and the packages' idle draw");
        figure["source"] = json!("powercap energy_uj of every package zone, read before the run, every 250 ms during it and after it, each wrap corrected at the zone's max_energy_range_uj");
        figure["host_load_1m_at_start"] = json!(load);
        if let Some(m) = figure.as_object_mut() {
            m.remove("concurrent_jobs");
        }
    }
    (out, secs, figure)
}

fn metered_process<T>(f: impl FnOnce() -> T) -> (T, f64, Value) {
    use wai_quantum::engine::energy::own;
    let before = own();
    let started = Instant::now();
    let out = f();
    let secs = started.elapsed().as_secs_f64();
    let after = own();
    let figure = match (before, after) {
        (Some(a), Some(b)) => {
            let nj = |r: &wai_quantum::engine::energy::Reading| r.energy_nj.saturating_add(r.billed_nj).saturating_sub(r.serviced_nj);
            let joules_micro = nj(&b).saturating_sub(nj(&a)).saturating_add(500) / 1000;
            if joules_micro == 0 {
                json!({ "unmetered": "the process's energy counter did not advance over the run" })
            } else {
                json!({
                    "joules_micro": joules_micro,
                    "energy_provenance": "OnChipCounter",
                    "energy_counter_filtering": "Unknown",
                    "attribution": "process",
                    "coverage": "partial",
                    "covers": "the CPU energy of this process's own threads over the run, with what system services billed to it",
                    "excludes": ["memory", "the host's idle draw", "energy of threads still on a core at the second read"],
                    "cpu_ms": b.cpu_ns.saturating_sub(a.cpu_ns) / 1_000_000,
                    "source": "proc_pid_rusage ri_energy_nj + ri_billed_energy − ri_serviced_energy of this process, read before and after the run",
                })
            }
        }
        _ => json!({ "unmetered": "this platform gives no per-process energy counter here" }),
    };
    (out, secs, figure)
}

/// `struct rusage`: two `struct timeval` (user and system time, 16 bytes
/// each on these 64-bit targets), then fourteen longs, `ru_maxrss` the first.
#[repr(C)]
struct Rusage {
    times: [i64; 4],
    longs: [i64; 14],
}

fn rusage() -> Option<Rusage> {
    unsafe extern "C" {
        fn getrusage(who: i32, usage: *mut Rusage) -> i32;
    }
    if !(cfg!(target_os = "macos") || cfg!(target_os = "linux")) || cfg!(not(target_pointer_width = "64")) {
        return None;
    }
    let mut r = Rusage { times: [0; 4], longs: [0; 14] };
    // SAFETY: `r` is a writable `struct rusage` of the platform's layout,
    // and RUSAGE_SELF (0) names this process.
    let rc = unsafe { getrusage(0, &mut r) };
    (rc == 0).then_some(r)
}

/// This process's peak resident memory in bytes, as `getrusage` reports it
/// (bytes on macOS, kilobytes on Linux).
pub fn peak_rss_bytes() -> Option<u64> {
    let r = rusage()?;
    if r.longs[0] <= 0 {
        return None;
    }
    let v = r.longs[0] as u64;
    Some(if cfg!(target_os = "macos") { v } else { v * 1024 })
}

/// This process's CPU time so far (user and system, every thread), in
/// nanoseconds.
pub fn cpu_ns() -> Option<u64> {
    let r = rusage()?;
    let t = |s: i64, us: i64| (s.max(0) as u64) * 1_000_000_000 + (us.max(0) as u64) * 1000;
    Some(t(r.times[0], r.times[1]) + t(r.times[2], r.times[3]))
}

/// The signing key: 64 hex digits in the file `QEC_SIGNING_KEY` names, or,
/// with none, a key every reader knows, which the receipt then says it was
/// sealed with.
pub fn signing_key() -> (SigningKey, &'static str) {
    match std::env::var("QEC_SIGNING_KEY") {
        Ok(path) => {
            let text = std::fs::read_to_string(&path).unwrap_or_else(|e| panic!("reading the key {path}: {e}"));
            let text = text.trim();
            assert_eq!(text.len(), 64, "the key file holds 64 hex digits");
            let mut b = [0u8; 32];
            for (i, x) in b.iter_mut().enumerate() {
                *x = u8::from_str_radix(&text[2 * i..2 * i + 2], 16).expect("hex digits");
            }
            (SigningKey::from_bytes(&b), "the operator's key")
        }
        Err(_) => (SigningKey::from_bytes(&[42u8; 32]), "a demonstration key every reader knows (set QEC_SIGNING_KEY to sign with your own)"),
    }
}

/// Seal `result` of `request` and write the three, as one JSON document, to
/// `out`. Returns the document.
pub fn seal(request: Value, mut result: Value, out: &Path) -> Value {
    let (key, which) = signing_key();
    result["signed_with"] = json!(which);
    let receipt = wai_quantum::engine::receipt::seal(&key, &request, &result);
    let doc = json!({ "request": request, "result": result, "receipt": receipt });
    if let Some(dir) = out.parent() {
        std::fs::create_dir_all(dir).unwrap_or_else(|e| panic!("creating {}: {e}", dir.display()));
    }
    std::fs::write(out, serde_json::to_string(&doc).unwrap() + "\n").unwrap_or_else(|e| panic!("writing {}: {e}", out.display()));
    doc
}