open-standards/joule-code/joule-code-rs/crates/joule-moq-relay/tests/relay_energy.rs

//! The relay's own energy in a room's per-hop session record (WAI `live-metering.md` §8,
//! `energy-measurement.md` §2.3, §3.1), on a host whose counter the relay reads: macOS's task
//! counter, or Linux's package counter where the configuration states the host runs the
//! relay alone (this test states it, to read the counter; the record is not published).
//!
//! A publisher sends a room's video for over a second to one subscriber. The record then
//! carries `OnChipCounter` with its evidence, reads back and verifies; the same room beside a
//! second room's session carries the unmetered marker, as does a relay whose configuration
//! turns the meter off. A record over a window of sessions still open carries the relay's
//! figure for that window too, a quiet session beside the room sharing nothing, and with another
//! room's objects in the window the room's share of the window's figure by the bytes each room
//! moved, `ModelBased` (WAI decision D79; `tests/room_energy_shares.rs` holds the rule on a
//! counter of its own, on every host). Where the host offers no counter (a
//! container, another platform, macOS on an Intel processor, whose task counter reads 0 while
//! the relay works) the relay says why, its record carries the unmetered marker and never a
//! figure of 0 J, and the test prints SKIPPED with the reason and passes only if
//! `JOULE_MOQ_RELAY_ENERGY_REQUIRE` is unset: a skip is never a pass where the counter is required.
//! A relay told to hold joule ceilings ends the sessions of a grant its own shares pass, sealing
//! on its own interval with nobody flushing it, and refuses the grant's next SETUP; one not told
//! to states the ceiling and ends nothing.

mod common;

use common::*;
use joule_moq::draft22::{SessionError, Token};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::config::{EnergyMode, MeteredGrant};
use joule_moq_relay::metering::HopRecord;
use joule_moq_relay::transport::{quic, TransportError};
use std::time::Duration;

const ROOM: &str = "5f3c2b1a09f8e7d6c5b4a39281706f5e4d3c2b1a09f8e7d6c5b4a39281706f5e";
const OTHER_ROOM: &str = "a0b1c2d3e4f5061728394a5b6c7d8e9fa0b1c2d3e4f5061728394a5b6c7d8e9f";
const GROUPS: u64 = 40;

fn key(i: u8) -> ed25519_dalek::SigningKey {
    ed25519_dalek::SigningKey::from_bytes(&[i; 32])
}

fn required() -> bool {
    std::env::var_os("JOULE_MOQ_RELAY_ENERGY_REQUIRE").is_some()
}

/// A publisher of `room` sends `GROUPS` groups over about 1.6 s to one subscriber; with
/// `other`, a second room's publisher stays connected across it. Both close before the
/// record is signed, and the sampler has passed the span's end.
async fn session(server: &joule_moq_relay::RelayServer, other: bool) -> HopRecord {
    let relay = server.relay().clone();
    assert!(until(STEP, || relay.energy_sampled_until().is_some()).await, "the sampler has a first sample");
    let k = key(7);
    let owner = hex::encode(k.verifying_key().to_bytes());
    let t = Tracks::new(ns(&[ROOM, &owner, "video"]), "0", 10, 1, 0);
    let (p, _ep) = connect(server, Via::Quic, Default::default()).await;
    let _pt = publish_tracks(&p, &t, None).await;
    let second = if other {
        let (q, eq) = connect(server, Via::Quic, Default::default()).await;
        let u = Tracks::new(ns(&[OTHER_ROOM, &owner, "video"]), "0", 20, 2, 0);
        let pu = publish_tracks(&q, &u, None).await;
        Some((q, eq, pu))
    } else {
        None
    };
    let (s, _es) = connect(server, Via::Quic, Default::default()).await;
    let next = joule_moq::draft22::Parameters::new().with(joule_moq::draft22::Parameter::LocationFilter(joule_moq::draft22::LocationFilter::NextObject));
    let v = subscribe_both(&s, &t, next).await;
    let publish = async {
        for g in 0..GROUPS {
            publish_group(&p, &k, &t, g, &payload("e", g, 4096)).await;
            tokio::time::sleep(Duration::from_millis(40)).await;
        }
    };
    let read = collect_concurrent(&s, STEP, Duration::from_secs(1), |got| got.iter().filter(|x| x.header.track_alias == v.media_alias).count() as u64 == GROUPS);
    let ((), got) = tokio::join!(publish, read);
    assert_eq!(got.iter().filter(|x| x.header.track_alias == v.media_alias).count() as u64, GROUPS, "every group reached the subscriber");
    s.close(0, "done");
    p.close(0, "done");
    let want = if other { 3 } else { 2 };
    if let Some((q, _, _)) = &second {
        q.close(0, "done");
    }
    assert!(until(STEP, || relay.summaries().len() >= want).await, "the sessions' summaries");
    let closed_us = relay.summaries().iter().map(|s| (s.summary.closed_unix_ms + 1) * 1000).max().unwrap();
    assert!(until(STEP, || relay.energy_sampled_until().is_some_and(|t| t > closed_us + 200_000)).await, "the sampler passed the span's end");
    let room: [u8; 32] = hex::decode(ROOM).unwrap().try_into().unwrap();
    relay.hop_record(room, 0)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_room_served_alone_carries_the_relays_own_energy() {
    let server = relay(|c| c.metering.dedicated_host = true).await;
    let Some(scope) = server.relay().energy_scope() else {
        let why = server.relay().energy_unavailable().expect("a relay with no counter says why").to_string();
        assert!(!required(), "JOULE_MOQ_RELAY_ENERGY_REQUIRE is set, and this host offers the relay no counter: {why}");
        // No counter: the record carries the unmetered marker and the reason is stated, never a
        // figure of 0 J (on macOS on an Intel processor the task counter reads 0 while work runs).
        unmetered_run(&server, 11).await;
        println!("SKIPPED: this host offers the relay no energy counter (macOS task counter, or a readable RAPL package counter); nothing was measured: {why}");
        server.shutdown().await;
        return;
    };
    let rec = tokio::time::timeout(Duration::from_secs(120), session(&server, false)).await.expect("in time");
    let bytes = rec.to_bytes();
    let back = HopRecord::decode(&bytes, &rec.ledger_bytes()).expect("the record reads back");
    assert_eq!(back, rec);
    rec.verify().expect("the record verifies under the relay's key");
    assert_eq!(rec.class_tag, 4, "OnChipCounter: the counter was live, the room alone ({} µJ)", rec.joules);
    assert!(rec.joules > 0);
    let ppm = u32::from_be_bytes(rec.class_evidence[..4].try_into().unwrap());
    let window = u64::from_be_bytes(rec.class_evidence[4..12].try_into().unwrap());
    assert_eq!(rec.class_evidence[12], 0, "filtering Unknown: an unprivileged reader cannot see one");
    assert_eq!(window, rec.end - rec.start, "the window is the record's span");
    assert!(rec.joules * u64::from(ppm) >= 288_676, "the resolution floor");
    assert_eq!(bytes.len(), 281 + 13 + 90 * rec.entries.len(), "tag 4's evidence, and nothing per object or per viewer");
    println!(
        "measured: {} µJ over {} µs, scope {}, relative standard uncertainty {} ppm ({:.3} W)",
        rec.joules,
        window,
        scope.label(),
        ppm,
        rec.joules as f64 / window as f64
    );
    server.shutdown().await;
    // On this counter too: a record over a window of sessions still open, and a hop told to hold
    // a joule ceiling. Beside the tests of this file, not new ones: the counter's CI job holds
    // the file to their number.
    a_window_of_open_sessions_carries_the_relays_own_energy().await;
    tokio::time::timeout(Duration::from_secs(120), ceiling_run(true)).await.expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_room_beside_another_rooms_session_carries_the_unmetered_marker() {
    let server = relay(|c| c.metering.dedicated_host = true).await;
    if server.relay().energy_scope().is_none() {
        let why = server.relay().energy_unavailable().expect("a relay with no counter says why");
        assert!(!required(), "JOULE_MOQ_RELAY_ENERGY_REQUIRE is set, and this host offers the relay no counter: {why}");
        println!("SKIPPED: this host offers the relay no energy counter; nothing was measured: {why}");
        server.shutdown().await;
        return;
    }
    let rec = tokio::time::timeout(Duration::from_secs(120), session(&server, true)).await.expect("in time");
    assert_eq!((rec.joules, rec.class_tag, rec.class_evidence.len()), (0, 0, 0), "another room's work in the span: the relay's energy is not this room's");
    server.shutdown().await;
    a_window_beside_another_rooms_objects_carries_the_rooms_share().await;
    tokio::time::timeout(Duration::from_secs(120), ceiling_run(false)).await.expect("in time");
}

/// The record over a window of `ROOM`, its sessions open throughout: a publisher sends for over
/// a second to a subscriber, beside a session that counts nothing; with `other`, a second room's
/// publisher sends its objects in the window too. With the window's statement of every session.
async fn window_of_open_sessions(server: &joule_moq_relay::RelayServer, other: bool) -> (HopRecord, joule_moq_relay::window::SignedWindow) {
    let relay = server.relay().clone();
    assert!(until(STEP, || relay.energy_sampled_until().is_some()).await, "the sampler has a first sample");
    let k = key(7);
    let owner = hex::encode(k.verifying_key().to_bytes());
    let t = Tracks::new(ns(&[ROOM, &owner, "video"]), "0", 10, 1, 0);
    let (p, _ep) = connect(server, Via::Quic, Default::default()).await;
    let _pt = publish_tracks(&p, &t, None).await;
    let (quiet, _eq) = connect(server, Via::Quic, Default::default()).await;
    let (s, _es) = connect(server, Via::Quic, Default::default()).await;
    let next = joule_moq::draft22::Parameters::new().with(joule_moq::draft22::Parameter::LocationFilter(joule_moq::draft22::LocationFilter::NextObject));
    let v = subscribe_both(&s, &t, next).await;
    let u = Tracks::new(ns(&[OTHER_ROOM, &owner, "video"]), "0", 20, 2, 0);
    let (q, _eo) = connect(server, Via::Quic, Default::default()).await;
    let _pu = publish_tracks(&q, &u, None).await;
    // The window opens here, with every session up and the sampler running.
    let (w, _, _) = relay.flush();
    assert_eq!(relay.seal(w), vec![w]);
    let publish = async {
        for g in 0..GROUPS {
            publish_group(&p, &k, &t, g, &payload("w", g, 4096)).await;
            if other {
                publish_group(&q, &k, &u, g, &payload("o", g, 64)).await;
            }
            tokio::time::sleep(Duration::from_millis(40)).await;
        }
    };
    let read = collect_concurrent(&s, STEP, Duration::from_secs(1), |got| got.iter().filter(|x| x.header.track_alias == v.media_alias).count() as u64 == GROUPS);
    let ((), got) = tokio::join!(publish, read);
    assert_eq!(got.iter().filter(|x| x.header.track_alias == v.media_alias).count() as u64, GROUPS, "every group reached the subscriber");
    let window = relay.flush_and_seal(Duration::from_millis(300)).await;
    assert_eq!(relay.stats().sessions_open, 4, "nobody left");
    assert!(until(STEP, || relay.energy_sampled_until().is_some_and(|t| t > window.to_unix_us + 200_000)).await, "the sampler passed the window's end");
    let room: [u8; 32] = hex::decode(ROOM).unwrap().try_into().unwrap();
    let rec = relay.window_record(room, 0, window.index).expect("the window's record");
    let statement = relay.window_statement(window.index, None).expect("the window's statement");
    for c in [&p, &quiet, &s, &q] {
        c.close(0, "done");
    }
    (rec, statement)
}

async fn a_window_of_open_sessions_carries_the_relays_own_energy() {
    let server = relay(|c| c.metering.dedicated_host = true).await;
    let (rec, statement) = tokio::time::timeout(Duration::from_secs(120), window_of_open_sessions(&server, false)).await.expect("in time");
    statement.check_shares().expect("the window's figure, the room's whole");
    statement.check_record(&rec, 0).expect("the record carries what the statement states");
    let back = HopRecord::decode(&rec.to_bytes(), &rec.ledger_bytes()).expect("the record reads back");
    assert_eq!(back, rec);
    rec.verify().expect("the record verifies under the relay's key");
    assert_eq!(rec.class_tag, 4, "OnChipCounter: the counter was live, the room alone in the window, a quiet session beside it ({} µJ)", rec.joules);
    assert!(rec.joules > 0);
    let window = u64::from_be_bytes(rec.class_evidence[4..12].try_into().unwrap());
    assert_eq!(window, rec.end - rec.start, "the evidence's window is the window's span");
    println!("measured over a window of open sessions: {} µJ over {window} µs", rec.joules);
    server.shutdown().await;
}

async fn a_window_beside_another_rooms_objects_carries_the_rooms_share() {
    let server = relay(|c| c.metering.dedicated_host = true).await;
    let (rec, statement) = tokio::time::timeout(Duration::from_secs(120), window_of_open_sessions(&server, true)).await.expect("in time");
    statement.check_shares().expect("every share recomputes from the window's parts");
    statement.check_record(&rec, 0).expect("the record carries the room's share the statement states");
    let e = statement.statement.energy.as_ref().expect("the window's figure and shares");
    let j = e.joules_micro.expect("the counter was live across the window");
    assert_eq!(rec.class_tag, joule_moq_relay::metering::TAG_MODEL_BASED, "another room's objects in the window: the relay's figure is not this room's, and its share is a model's, not the counter's ({} µJ of {j})", rec.joules);
    assert!(rec.joules > 0 && rec.joules < j && rec.class_evidence.is_empty());
    assert_eq!(rec.joules, joule_moq_relay::window::share(j, rec.entry_bytes(), e.bytes).unwrap(), "⌊J × b / B⌋");
    println!("shared over a window with another room's objects: {} µJ of {j} µJ, {} of {} bytes", rec.joules, rec.entry_bytes(), e.bytes);
    server.shutdown().await;
}

/// A publisher under a grant with no ceiling sends for over a second to a subscriber under a
/// grant whose joule ceiling is one microjoule, and the window closes on their counts.
async fn ceiling_run(hold: bool) {
    let (free, ceiled) = (b"no-ceiling".to_vec(), b"one-microjoule".to_vec());
    let bearer = |t: &Vec<u8>| ClientOptions { tokens: vec![Token::UseValue { token_type: 0, value: t.clone() }], ..ClientOptions::default() };
    let metered = |t: &Vec<u8>, grant: String, joules: u64| MeteredGrant { token_sha256: hex::encode(joule_moq_relay::auth::sha256(t)), grant, joules_ceiling_micro: joules, funds_ceiling: u64::MAX, currency: String::new() };
    let tiny = "c1".repeat(32);
    let server = relay(|c| {
        c.metering.dedicated_host = true;
        c.metering.grants = vec![metered(&free, "c0".repeat(32), u64::MAX), metered(&ceiled, tiny.clone(), 1)];
        c.control.hold_joule_ceilings = hold;
        // The grace outlasts the sampler's interval many times over: a seal shares the window's
        // figure, which needs a sample past its end.
        if hold {
            c.control.flush_interval_ms = Some(500);
            c.control.flush_grace_ms = 1_000;
        }
    })
    .await;
    let relay = server.relay().clone();
    assert!(until(STEP, || relay.energy_sampled_until().is_some()).await, "the sampler has a first sample");
    let k = key(7);
    let t = Tracks::new(ns(&[ROOM, &hex::encode(k.verifying_key().to_bytes()), "video"]), "0", 10, 1, 0);
    let (p, _ep) = connect(&server, Via::Quic, bearer(&free)).await;
    let _pt = publish_tracks(&p, &t, None).await;
    let (s, _es) = connect(&server, Via::Quic, bearer(&ceiled)).await;
    let next = joule_moq::draft22::Parameters::new().with(joule_moq::draft22::Parameter::LocationFilter(joule_moq::draft22::LocationFilter::NextObject));
    let v = subscribe_both(&s, &t, next).await;
    if !hold {
        let (w, _, _) = relay.flush();
        assert_eq!(relay.seal(w), vec![w]);
    }
    let publish = async {
        for g in 0..GROUPS {
            publish_group(&p, &k, &t, g, &payload("c", g, 4096)).await;
            tokio::time::sleep(Duration::from_millis(40)).await;
        }
    };
    let read = collect_concurrent(&s, STEP, Duration::from_secs(1), |got| got.iter().filter(|x| x.header.track_alias == v.media_alias).count() as u64 == GROUPS);
    let ((), _) = tokio::join!(publish, read);
    if hold {
        // Nothing here flushes: the relay's own interval seals the window, and holds the ceiling.
        match tokio::time::timeout(Duration::from_secs(60), s.conn.closed()).await.expect("closed in time") {
            TransportError::Closed { code, reason, .. } => {
                assert_eq!(code, SessionError::UNAUTHORIZED.0);
                assert!(reason.starts_with(joule_moq_relay::relay::PAST_CEILING), "{reason}");
            }
            other => panic!("{other:?}"),
        }
        assert_eq!(relay.grants_past_ceiling(), vec![tiny], "its share of the relay's figure passed one microjoule");
        assert!(until(STEP, || relay.stats().sessions_open == 1).await, "the grant with no ceiling goes on");
        // The grant's next SETUP is refused: the relay closes the session at its door, with the
        // bare reason. The relay reads the client's SETUP, sends its own saying it refuses, and
        // closes the session once the client has left or its wait has passed (session.rs); this
        // client stays, as one that does not know GRANT_REFUSED does, so `Client::connect`'s Ok
        // is not an admission. What is held is the relay's own close, within STEP of the
        // client's SETUP, and its reason exactly the door's: a session admitted and then ended at the relay's next seal closes with the
        // same code and a longer reason, which says the grant's joules passed its ceiling.
        let ep = quic::client_endpoint(loopback(), vec![server.identity.chain[0].clone()], vec![joule_moq::draft22::ALPN.to_vec()], None).unwrap();
        let conn: joule_moq_relay::transport::SharedConnection = std::sync::Arc::new(quic::connect(&ep, server.quic_addr.unwrap(), "localhost").await.expect("a QUIC session"));
        let client = Client::connect(conn.clone(), ClientOptions { ignore_refusal: true, ..bearer(&ceiled) }).await;
        match tokio::time::timeout(STEP, conn.closed()).await {
            Ok(TransportError::Closed { code, reason, by_peer: true }) => {
                assert_eq!(code, SessionError::UNAUTHORIZED.0, "{reason}");
                assert_eq!(reason, joule_moq_relay::relay::PAST_CEILING, "a grant past its ceiling here was admitted, and its session ended after");
            }
            Ok(other) => panic!("the grant's next session ended otherwise: {other:?}"),
            Err(_) => panic!("a grant past its ceiling here was admitted: its session was still open {STEP:?} after its SETUP (connect: {:?})", client.as_ref().map(|_| "Ok")),
        }
        drop(client);
    } else {
        // The grace outlasts the sampler's interval, as the relay's own seal's does above.
        relay.flush_and_seal(Duration::from_secs(2)).await;
        assert!(relay.grants_past_ceiling().is_empty(), "the ceiling is stated, and the deployment's meter holds it");
        assert_eq!(relay.stats().sessions_open, 2, "nobody ended");
        s.close(0, "done");
    }
    p.close(0, "done");
    server.shutdown().await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_relay_whose_meter_is_off_carries_the_unmetered_marker() {
    let server = relay(|c| c.metering.energy = EnergyMode::Off).await;
    assert!(server.relay().energy_scope().is_none());
    assert_eq!(server.relay().energy_unavailable(), Some("metering.energy is off"));
    unmetered_run(&server, 9).await;
    server.shutdown().await;
}

/// A relay with no counter: a room's record over a session carries the unmetered marker, and
/// the session's share is never read.
async fn unmetered_run(server: &joule_moq_relay::RelayServer, seed: u8) {
    let relay = server.relay().clone();
    let k = key(seed);
    let t = Tracks::new(ns(&[ROOM, &hex::encode(k.verifying_key().to_bytes()), "video"]), "0", 10, 1, 0);
    let (p, _ep) = connect(server, Via::Quic, Default::default()).await;
    let _pt = publish_tracks(&p, &t, None).await;
    publish_group(&p, &k, &t, 0, &payload("e", 0, 64)).await;
    p.close(0, "done");
    assert!(until(STEP, || !relay.summaries().is_empty()).await);
    let rec = relay.hop_record(hex::decode(ROOM).unwrap().try_into().unwrap(), 0);
    assert_eq!((rec.joules, rec.class_tag), (0, 0), "the unmetered marker, never a figure of 0 J");
    assert_eq!(relay.summaries()[0].summary.relay_energy.acquisition, "unmetered", "a session's share is never read");
}