open-standards/joule-ui/joule-ui-rs/crates/joule-ui-interaction/examples/support/span_meter.rs

//! The Sandbox `PowerMeter` as a joule-ui `Meter` (SPEC §5 edge) over managed
//! spans, shared by the `real_meter` and `jcp_metered` examples. It is the
//! adapter `joule-ui-serve` seals its served interactions with.
//!
//! Every `sample_uj` closes the span the previous sample opened, adds that
//! span's figure to a running counter, and opens the next. So the difference
//! of two consecutive samples is exactly one span's figure, and `provenance`
//! is that span's own, which `run_interaction` asks for right after its second
//! sample. A span is `HwShunt` only when it was a hardware reading of the CPU
//! the handler ran on, whose CPU counter advanced. Otherwise it is the meter's
//! `Estimator` figure (wall-clock × TDP, never 0 for work that ran): on a host
//! with no live energy counter — every macOS 27.0.1 host checked, whose CPU
//! energy rail is frozen — when the span's CPU counter held while another
//! class moved, and when the span has no CPU reading at all (an NVML-only
//! host reads the GPU alone: see [`interaction_figure`]).
//!
//! The examples' old adapter sampled `PowerMeter::sample_uj`, a sum of raw
//! counters under the trait's default `HwShunt`. Where no hardware probe
//! existed that was `None`, and each interaction that ran was sealed at
//! `0 µJ [Estimator]`.
//!
//! `tests/span_meter.rs` includes this file too, and drives the adapter over
//! scripted spans.

use joule_ui_interaction::{Meter, Provenance};
use sandbox_meter::probe::Span;
use sandbox_meter::{Measurement, PowerMeter};
use std::cell::RefCell;
use std::time::Duration;

/// A meter that measures managed spans: the Sandbox `PowerMeter` does, and a
/// test can script one.
pub trait Spans {
    type Open;
    fn open(&mut self) -> Self::Open;
    fn close(&mut self, span: Self::Open) -> Measurement;
    /// The meter's estimate for `wall` of work, with no reading.
    fn estimate(&self, wall: Duration) -> Measurement;
}

impl Spans for PowerMeter {
    type Open = Span;
    fn open(&mut self) -> Span {
        self.start()
    }
    fn close(&mut self, span: Span) -> Measurement {
        span.stop(self)
    }
    fn estimate(&self, wall: Duration) -> Measurement {
        PowerMeter::estimate(self, wall)
    }
}

/// `RefCell` because sampling mutates the meter while the `Meter` trait
/// samples through `&self`. Each example drives one interaction at a time, so
/// a span is exclusive to one interaction.
pub struct SpanMeter<M: Spans = PowerMeter>(RefCell<State<M>>);

struct State<M: Spans> {
    meter: M,
    open: Option<M::Open>,
    counter_uj: u64,
    last: Provenance,
}

impl<M: Spans> SpanMeter<M> {
    pub fn new(meter: M) -> Self {
        SpanMeter(RefCell::new(State { meter, open: None, counter_uj: 0, last: Provenance::Estimator }))
    }
}

impl<M: Spans> Meter for SpanMeter<M> {
    fn sample_uj(&self) -> Option<u64> {
        let mut s = self.0.borrow_mut();
        if let Some(span) = s.open.take() {
            let m = s.meter.close(span);
            let m = interaction_figure(m, |wall| s.meter.estimate(wall));
            s.counter_uj = s.counter_uj.saturating_add(m.total_uj());
            s.last = span_provenance(&m);
        }
        let next = s.meter.open();
        s.open = Some(next);
        Some(s.counter_uj)
    }

    /// The provenance of the span the last sample closed.
    fn provenance(&self) -> Provenance {
        self.0.borrow().last
    }
}

/// The figure an interaction is sealed with, from its span. joule-ui seals
/// one figure per interaction, with no field to say that figure covers part
/// of the work, so a span with no hardware reading of the CPU is not a
/// reading of the interaction: an NVML-only probe set reads the GPU alone,
/// and its figure would be sealed `HwShunt` for work that ran on the CPU. Such
/// a span is sealed as the meter's estimate for its wall time (`Estimator`,
/// never 0 for work that ran); a span that reads the CPU passes through,
/// with any GPU record beside it.
fn interaction_figure(m: Measurement, estimate: impl FnOnce(Duration) -> Measurement) -> Measurement {
    if m.reads_cpu() { m } else { estimate(m.wall) }
}

/// A span's provenance: the weakest of its records, because a sum is only as
/// good as its weakest part.
fn span_provenance(m: &Measurement) -> Provenance {
    use sandbox_meter::Provenance as Grade;
    if m.devices.is_empty() || m.devices.iter().any(|d| d.provenance == Grade::Estimator) {
        Provenance::Estimator
    } else if m.devices.iter().all(|d| d.provenance == Grade::HwShunt && d.measured_uj.is_some()) {
        Provenance::HwShunt
    } else {
        Provenance::ModelBased
    }
}