open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-energy/examples/energy_probe.rs

//! Cross-platform energy probe — read the best available hardware counter on
//! this host and print its energy delta with its honest provenance.
//!
//! The point of the probe: show whether on real hardware JouleClaw reports a
//! `HwShunt` / `ModelBased` reading or an `Estimator` fallback, and label the
//! figure accordingly. Build with the feature that matches the host's counter:
//!
//!   # NVIDIA discrete GPU (Linux or Windows — loads libnvidia-ml / nvml.dll):
//!   cargo run -p jouleclaw-energy --release --example energy_probe --features nvml
//!
//!   # Intel/AMD x86 on Linux (powercap sysfs RAPL):
//!   cargo run -p jouleclaw-energy --release --example energy_probe --features rapl
//!
//!   # Apple Silicon (IOReport Energy Model when its CPU rail is live, else the
//!   # load-avg estimator, with the reason printed):
//!   cargo run -p jouleclaw-energy --release --example energy_probe --features ioreport
//!
//!   # Mac, whole machine (the SMC's `PSTR` integrated over time, ModelBased;
//!   # with `ioreport` too, after a live IOReport counter):
//!   cargo run -p jouleclaw-energy --release --example energy_probe --features smc
//!
//! It samples the counter across a ~2 s window (driving a small CPU workload so
//! the package/SoC is not fully idle) and prints the integrated microjoules. For
//! the integrating counters (NVML) the repeated reads during the window are what
//! build the energy integral.

use jouleclaw_energy::EnergyCounter;
use std::time::{Duration, Instant};

/// Pick the first available real counter for the features this binary was built
/// with. Returns the counter and a short label.
fn pick_counter() -> Option<(&'static str, Box<dyn EnergyCounter>)> {
    #[cfg(feature = "nvml")]
    {
        let c = jouleclaw_energy::nvml::NvmlCounter::new();
        if c.is_available() {
            // Report which NVML path reads — the cumulative-energy counter,
            // only where it was observed live (HwShunt), or the power-draw
            // integral (ModelBased) — and the counter's liveness.
            println!(
                "nvml path  : {} (counter liveness: {:?})",
                if c.uses_true_counter() {
                    "nvmlDeviceGetTotalEnergyConsumption (live cumulative counter, HwShunt)"
                } else {
                    "power_usage() integration (ModelBased)"
                },
                c.liveness()
            );
            return Some(("nvml (NVIDIA GPU shunt)", Box::new(c)));
        }
    }
    #[cfg(feature = "rapl")]
    {
        let c = jouleclaw_energy::rapl::RaplCounter::new();
        println!("rapl       : liveness {:?}", c.liveness());
        if c.is_available() {
            return Some(("rapl (x86 package MSR)", Box::new(c)));
        }
    }
    #[cfg(all(feature = "ioreport", target_os = "macos"))]
    {
        // Prefer the real IOReport energy counter (HwShunt, like RAPL and NVML).
        // It is unavailable unless IOReport opens AND its CPU energy rail
        // advances under load — a counter that is present but not live gives no
        // hardware-graded figure; say which before falling back to the load-avg
        // estimator.
        let c = jouleclaw_energy::apple_ioreport::AppleIoReportCounter::new();
        if c.is_available() {
            return Some(("apple-ioreport (Energy Model)", Box::new(c)));
        }
        println!("ioreport   : declined ({:?})", c.status());
    }
    #[cfg(all(feature = "smc", target_os = "macos"))]
    {
        // The whole machine's power from the SMC, integrated over time: a vendor
        // power reading, so ModelBased, over the System domain.
        let c = jouleclaw_energy::apple_smc::SmcSystemCounter::new();
        if c.is_available() {
            println!("smc        : PSTR {:.1} W now", c.watts().unwrap_or(0.0));
            return Some(("apple-smc (SMC PSTR, whole machine)", Box::new(c)));
        }
        println!("smc        : unavailable ({})", c.unavailable_reason().unwrap_or("?"));
    }
    #[cfg(all(feature = "ioreport", target_os = "macos"))]
    {
        let c = jouleclaw_energy::apple::AppleSiliconMeter::new();
        if c.is_available() {
            return Some(("apple (load-avg estimator)", Box::new(c)));
        }
    }
    None
}

/// A small CPU workload so the host is not fully idle during the window.
fn churn(deadline: Instant) -> u64 {
    let mut acc: u64 = 0;
    let mut x: u64 = 0x9E3779B97F4A7C15;
    while Instant::now() < deadline {
        for _ in 0..200_000 {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            acc = acc.wrapping_add(x);
        }
    }
    acc
}

fn main() {
    println!("== JouleClaw energy probe ==");
    let Some((label, counter)) = pick_counter() else {
        eprintln!(
            "no hardware energy counter available on this host/build.\n\
             rebuild with --features nvml|rapl|ioreport|smc for this platform."
        );
        std::process::exit(1);
    };

    println!("counter    : {label}");
    println!("provenance : {}", counter.provenance());
    println!("domain     : {:?}", counter.domain());
    println!("resolution : {} uj", counter.resolution_uj());
    println!("min window : {} ns", counter.min_window_ns());

    // Prime the counter (first read seeds the integrator / cumulative baseline).
    let _ = counter.read();
    let before = match counter.read() {
        Ok(r) => r,
        Err(e) => {
            eprintln!("counter read failed: {e}");
            std::process::exit(1);
        }
    };

    // Integrate over a ~2 s window, or the counter's minimum window if longer
    // (plus a margin): churn the CPU and re-read every ~200 ms so the
    // integrating counters accumulate the wall-power integral.
    let floor = Duration::from_nanos(counter.min_window_ns());
    let window = Duration::from_secs(2).max(floor + Duration::from_millis(500));
    let start = Instant::now();
    let mut checksum: u64 = 0;
    while start.elapsed() < window {
        let step = (Instant::now() + Duration::from_millis(200)).min(start + window);
        checksum = checksum.wrapping_add(churn(step));
        let _ = counter.read();
    }
    let after = match counter.read() {
        Ok(r) => r,
        Err(e) => {
            eprintln!("counter read failed: {e}");
            std::process::exit(1);
        }
    };

    let elapsed = start.elapsed().as_secs_f64();
    let delta_uj = after.uj.saturating_sub(before.uj);
    let joules = delta_uj as f64 / 1_000_000.0;
    // An estimator's figure is not a measurement; say which one this is.
    if after.provenance == jouleclaw_energy::Provenance::Estimator {
        println!("--- estimated ---");
    } else {
        println!("--- measured ---");
    }
    println!("window     : {:.3} s", elapsed);
    println!(
        "energy     : {} uj  ({:.3} J)  avg {:.2} W  [{}]",
        delta_uj,
        joules,
        if elapsed > 0.0 { joules / elapsed } else { 0.0 },
        after.provenance,
    );
    // Keep the workload from being optimized away.
    if checksum == 0 {
        eprintln!("(workload checksum was zero — unexpected)");
    }
}