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

//! The relay's per-hop session record of a room (WAI `live-metering.md` §3, §8): two
//! publishers and three subscribers of one WAI room, one subscriber slow, and the record the
//! relay signs from its session summaries recomputes from their counts.
//!
//! The room is the corpus's `relay` room (`wai/live-conformance/metering/`): its id, the
//! principals' keys and their grants are the ones `wai_live_metering_vectors relay-room`
//! prints, every key SHA-256 of capture binding's corpus label and the principal's index,
//! and the relay signs with the room's clock key. Tracks are named as live-rooms §6.1 names
//! them: each publisher's video, layer 0, with its receipts track, and the first
//! publisher's events track. The second publisher's video has a 100 ms delivery timeout,
//! and the slow subscriber reads nothing until the load has arrived and a second and a
//! half more, so some of its objects are dropped at this hop while their receipts still
//! reach it. Every subscriber reads its streams concurrently, as slow_subscriber.rs's do
//! (README, "A subscriber that does not read"). The subscribers' sessions end before the
//! publishers'. As a session closes, the relay ends its subscriptions one at a time, and a
//! moment after its last media subscription ended and before its receipts or events one
//! did is time on no media track, class 8: the test requires such time to fall within the
//! session's close, and no time above layer 0 at all.
//!
//! With `WAI_METERING_FIXTURE_DIR` set, the test also writes the record and its ledger there
//! as `relay-record.wmh` and `relay-ledger.wml`: the corpus holds the pair one run wrote.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{Parameter, Parameters, Token};
use joule_moq_relay::client::ClientOptions;
use joule_moq_relay::config::MeteredGrant;
use joule_moq_relay::metering::{classify, HopRecord, Kind, CLASSES};
use std::collections::BTreeMap;
use std::time::Duration;

/// The corpus room's id and `t0` (2026-10-18T00:00:00Z).
const ROOM: &str = "89a7be137d7f3c51d15b0d4ece614d60743d7542b7693c96ac7633793e9c102f";
const T0_US: u64 = 1_792_281_600_000_000;
/// Each principal's index under the corpus's key label, its grant's digest and ceilings.
const GRANTS: [(u8, &str, u64, u64, &str); 5] = [
    (30, "510124baa39d99a07d9d6575698f70183759e3da8b458f8edb09f09215e14318", u64::MAX, u64::MAX, ""),
    (31, "4fc2bb567c1e950b92bb2dd6796df52eb67f3e5c2e5753b3ee9bae483600f17e", u64::MAX, u64::MAX, ""),
    (32, "52786b7331441935c2143544988bccb4a4db6f5c76d734292a843a0b5e20e83d", 5_000_000, 500, "EUR"),
    (33, "fec5bc703c31e2600d5a51cceb45eb2e5011ec2c1890f8dc5b96fd683ef85362", u64::MAX, u64::MAX, ""),
    (34, "6b4272b140a49c75d1dbedfca74343ad6c13a8a9d3a20429e7d457f30d2afd82", u64::MAX, u64::MAX, ""),
];
const RELAY: u8 = 1;
const GROUPS: u64 = 30;
const EVENT_GROUPS: u64 = 6;
const STALL: Duration = Duration::from_millis(1500);

fn sha256(parts: &[&[u8]]) -> [u8; 32] {
    let mut c = ring::digest::Context::new(&ring::digest::SHA256);
    for p in parts {
        c.update(p);
    }
    c.finish().as_ref().try_into().unwrap()
}

/// A principal's key: SHA-256 of capture binding's corpus label and its index.
fn key(i: u8) -> ed25519_dalek::SigningKey {
    ed25519_dalek::SigningKey::from_bytes(&sha256(&[b"wai:capture-take/corpus/key\x00", &[i]]))
}

fn pk_hex(i: u8) -> String {
    hex::encode(key(i).verifying_key().to_bytes())
}

fn token(i: u8) -> Vec<u8> {
    format!("room-metering-token-{i}").into_bytes()
}

fn opts(i: u8) -> ClientOptions {
    ClientOptions { tokens: vec![Token::UseValue { token_type: 0, value: token(i) }], ..ClientOptions::default() }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn two_publishers_three_subscribers_one_slow_and_the_record_recomputes() {
    tokio::time::timeout(Duration::from_secs(180), run()).await.expect("the test finishes in time");
}

async fn run() {
    let seed_path = std::env::temp_dir().join(format!("joule-moq-relay-room-metering-{}.seed", std::process::id()));
    std::fs::write(&seed_path, hex::encode(sha256(&[b"wai:capture-take/corpus/key\x00", &[RELAY]]))).unwrap();
    let server = relay(|c| {
        c.limits.quic_send_window_bytes = 64 * 1024;
        c.summary.key_path = Some(seed_path.clone());
        // The corpus holds this record, so its bytes must not depend on the host: its energy
        // is the unmetered marker here, and tests/relay_energy.rs reads the meter.
        c.metering.energy = joule_moq_relay::config::EnergyMode::Off;
        c.metering.grants = GRANTS
            .iter()
            .map(|&(i, g, jc, fc, cur)| MeteredGrant {
                token_sha256: hex::encode(sha256(&[&token(i)])),
                grant: g.into(),
                joules_ceiling_micro: jc,
                funds_ceiling: fc,
                currency: cur.into(),
            })
            .collect();
    })
    .await;
    let _ = std::fs::remove_file(&seed_path);
    assert_eq!(server.relay().summary_key().to_bytes(), key(RELAY).verifying_key().to_bytes(), "the relay signs with the room's clock key");
    let (ka, kb) = (key(30), key(31));
    let video = |i: u8, alias: u64, receipt_alias: u64, prio: u8| Tracks::new(ns(&[ROOM, &pk_hex(i), "video"]), "0", alias, receipt_alias, prio);
    let a = video(30, 10, 1, 0);
    let b = video(31, 20, 2, 200);
    let events_ns = ns(&[ROOM, &pk_hex(30), "events"]);

    let (pa, _ea) = connect(&server, Via::Quic, opts(30)).await;
    let (pb, _eb) = connect(&server, Via::Quic, opts(31)).await;
    let _pa = publish_tracks(&pa, &a, None).await;
    let _pb = publish_tracks(&pb, &b, Some(100)).await;
    let _pe = pa.publish(events_ns.clone(), name("0"), 30, Parameters::new(), track_props(64, None)).await.expect("PUBLISH the events track");

    let (s1, _e1) = connect(&server, Via::Quic, opts(32)).await;
    let (s2, _e2) = connect(&server, Via::WebTransport, opts(33)).await;
    let mut small = quinn::TransportConfig::default();
    small.stream_receive_window(quinn::VarInt::from_u32(32 * 1024));
    small.receive_window(quinn::VarInt::from_u32(96 * 1024));
    let (s3, _e3) = connect_with(&server, Via::Quic, ClientOptions { lazy_data: true, ..opts(34) }, Some(small)).await;
    let next = || Parameters::new().with(Parameter::LocationFilter(joule_moq::draft22::LocationFilter::NextObject));
    let mut views = Vec::new();
    for s in [&s1, &s2, &s3] {
        let va = subscribe_both(s, &a, next()).await;
        let vb = subscribe_both(s, &b, next()).await;
        let (ve, etx, erx) = s.subscribe(events_ns.clone(), name("0"), next()).await.expect("SUBSCRIBE the events track");
        views.push((va, vb, ve.track_alias, (etx, erx)));
    }

    let (published_tx, published_rx) = tokio::sync::oneshot::channel::<()>();
    let publish = async {
        for g in 0..GROUPS {
            publish_group(&pa, &ka, &a, g, &payload("a", g, 900)).await;
            publish_group(&pb, &kb, &b, g, &payload("b", g, 24 * 1024)).await;
            if g < EVENT_GROUPS {
                // A group of an MSF event timeline: the join set, then two events.
                let h = SubgroupHeader { track_alias: 30, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(64), properties: false, end_of_group: true, first_object: true };
                let objects: Vec<SubgroupObject> = [b"[]".to_vec(), format!("[\"event {g}.1\"]").into_bytes(), format!("[\"event {g}.2\"]").into_bytes()]
                    .into_iter()
                    .enumerate()
                    .map(|(k, p)| SubgroupObject::normal(k as u64, p))
                    .collect();
                pa.send_subgroup(h, &objects).await.expect("an events group");
            }
            tokio::time::sleep(Duration::from_millis(15)).await;
        }
        let _ = published_tx.send(());
    };
    let whole = |got: &[Got], v: &(Subscribed, Subscribed, u64, _)| {
        let n = |al: u64| got.iter().filter(|x| x.header.track_alias == al && x.reset.is_none()).count() as u64;
        n(v.0.media_alias) == GROUPS && n(v.0.receipts_alias) == GROUPS && n(v.1.receipts_alias) == GROUPS && n(v.2) == EVENT_GROUPS
    };
    // Each reads until A's objects, every receipt and every event group are whole, then
    // until a second passes with nothing new, so every B stream is read or seen reset.
    let fast = |i: usize| {
        let (s, v) = (if i == 0 { &s1 } else { &s2 }, &views[i]);
        async move { collect_concurrent(s, STEP * 2, Duration::from_secs(1), |got| whole(got, v)).await }
    };
    let slow = async {
        let _ = tokio::time::timeout(STEP * 2, published_rx).await;
        tokio::time::sleep(STALL).await;
        collect_concurrent(&s3, STEP * 2, Duration::from_secs(1), |got| whole(got, &views[2])).await
    };
    let ((), g1, g2, g3) = tokio::join!(publish, fast(0), fast(1), slow);
    for (g, v) in [(&g1, &views[0]), (&g2, &views[1]), (&g3, &views[2])] {
        assert!(whole(g, v), "every A object, every receipt and every event group reached each subscriber");
    }

    let relay = server.relay().clone();
    let closing = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_micros() as u64;
    for s in [&s1, &s2, &s3] {
        s.close(0, "done");
    }
    assert!(until(STEP, || relay.summaries().len() >= 3).await, "the subscribers' three summaries, before a publisher leaves");
    pa.close(0, "done");
    pb.close(0, "done");
    assert!(until(STEP, || relay.summaries().len() >= 5).await, "five summaries");
    let room: [u8; 32] = hex::decode(ROOM).unwrap().try_into().unwrap();
    let rec = relay.hop_record(room, T0_US);
    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.hop, key(RELAY).verifying_key().to_bytes());
    assert_eq!((rec.joules, rec.class_tag), (0, 0), "the unmetered marker: the meter is off for the corpus's record");
    assert_eq!(rec.to_bytes().len(), 281 + 90 * rec.entries.len(), "no class evidence, and nothing per object or per viewer");
    assert!(rec.start < rec.end);

    // The record recomputes from the summaries, summed here track by track.
    let summaries: Vec<_> = relay.summaries().into_iter().map(|s| s.summary).collect();
    let mut want: BTreeMap<(u8, [u8; 32], u8), [u64; 7]> = BTreeMap::new();
    let kind_of = |n: &str| match classify(n) {
        Some((r, k)) if r == room => Some(match k {
            Kind::Media { digest, .. } => (1u8, digest),
            Kind::Events => (2, [0; 32]),
            Kind::Other => (3, [0; 32]),
        }),
        _ => None,
    };
    let (mut events, mut sub_us) = (0u64, 0u64);
    for s in &summaries {
        for (n, t) in &s.tally.subscriber.tracks {
            let Some((k, d)) = kind_of(n) else { continue };
            let e = want.entry((k, d, 2)).or_default();
            e[0] += t.objects_forwarded + t.fetch_objects;
            e[1] += t.object_bytes_forwarded + t.fetch_bytes;
            e[2] += t.dropped_delivery_timeout;
            e[5] += t.undelivered_at_close;
            e[6] += t.dropped_receipts_track_subscribed;
            if k == 2 {
                events += t.incremental_objects;
            }
        }
        for (n, t) in &s.tally.publisher.tracks {
            let Some((k, d)) = kind_of(n) else { continue };
            let e = want.entry((k, d, 1)).or_default();
            e[0] += t.objects_received;
            e[1] += t.object_bytes_received;
            e[3] += t.refused_at_door;
            e[4] += t.deduplicated;
        }
        sub_us += s.tally.subscriber.layer_spans.iter().filter(|x| x.room == ROOM).map(|x| x.to_unix_us - x.from_unix_us).sum::<u64>();
    }
    let got: BTreeMap<(u8, [u8; 32], u8), [u64; 7]> = rec
        .entries
        .iter()
        .map(|e| ((e.kind, e.track, e.dir), [e.objects, e.bytes, e.dropped_timeout, e.dropped_refused, e.dropped_duplicate, e.dropped_closed, e.receipted]))
        .collect();
    assert_eq!(got, want, "each track entry is the sum of the summaries' counts");
    let digest_of = |i: u8| sha256(&[format!("{ROOM}-{}-video", pk_hex(i)).as_bytes(), &[0], b"0"]);
    let (da, db) = (digest_of(30), digest_of(31));
    assert_eq!(got[&(1, da, 1)][0], GROUPS, "the relay received each of A's objects once");
    assert_eq!(got[&(1, db, 1)][0], GROUPS);
    let b_out = got[&(1, db, 2)];
    assert!(b_out[2] > 0, "the slow subscriber lost some of B's objects to their delivery timeout");
    assert_eq!(b_out[6], b_out[2], "the receipts of every dropped object were forwarded");
    assert_eq!(b_out[0] + b_out[2] + b_out[5], 3 * GROUPS, "each of B's objects to each subscriber sent or dropped");
    assert_eq!(got[&(1, da, 2)][0], 3 * GROUPS, "A's track, with no timeout, reached all three whole");
    assert_eq!(rec.events, events);
    assert_eq!(rec.events, 3 * EVENT_GROUPS * 2, "two events a group, to each subscriber; the join sets are not events");
    assert_eq!(rec.sessions, 3, "three subscriber sessions; a publisher's session subscribes to nothing");
    assert_eq!(rec.subscriber_us.iter().sum::<u64>(), sub_us);
    // Each subscriber session's layer spans, relative to the instant the test began closing
    // the subscribers, for a report on failure.
    let spans: Vec<(u64, u8, u64, u64)> = summaries
        .iter()
        .flat_map(|s| s.tally.subscriber.layer_spans.iter().filter(|x| x.room == ROOM).map(move |x| (s.session, x.class, x.from_unix_us, x.to_unix_us)))
        .collect();
    let report = || {
        spans
            .iter()
            .map(|(id, c, a, z)| format!("session {id}: class {c} from {} to {} us after closing began", *a as i64 - closing as i64, *z as i64 - closing as i64))
            .collect::<Vec<_>>()
            .join("; ")
    };
    if spans.iter().any(|x| x.1 != 0) {
        eprintln!("TEARDOWN-CLASS: time outside class 0: {}", report());
    }
    assert!(rec.subscriber_us[0] > 0, "every subscriber held layer 0 of the room's video: {}", report());
    assert!(rec.subscriber_us[1..8].iter().all(|&x| x == 0), "no subscriber held a layer above 0: {}", report());
    assert!(
        spans.iter().filter(|x| x.1 != 0).all(|x| x.2 >= closing),
        "time on no media track only while a session closed, its subscriptions ending one by one: {}",
        report()
    );

    // The ledger: one entry for each grant, from the token each session presented.
    let mut grants: Vec<&str> = GRANTS.iter().map(|g| g.1).collect();
    grants.sort();
    assert_eq!(rec.ledger.iter().map(|e| hex::encode(e.grant)).collect::<Vec<_>>(), grants);
    let s1_entry = rec.ledger.iter().find(|e| hex::encode(e.grant) == GRANTS[2].1).unwrap();
    assert_eq!((s1_entry.joules_ceiling, s1_entry.funds_ceiling, &s1_entry.currency), (5_000_000, 500, b"EUR"));
    for i in 0..CLASSES {
        assert_eq!(rec.ledger.iter().map(|e| e.subscriber_us[i]).sum::<u64>(), rec.subscriber_us[i]);
    }
    assert!(rec.ledger.iter().all(|e| e.sessions == 1));
    let ina: u64 = got.iter().filter(|(k, _)| k.2 == 1).map(|(_, v)| v[1]).sum();
    assert_eq!(rec.ledger.iter().map(|e| e.bytes_in).sum::<u64>(), ina, "the bytes in are the publishers'");

    if let Ok(dir) = std::env::var("WAI_METERING_FIXTURE_DIR") {
        let d = std::path::Path::new(&dir);
        std::fs::write(d.join("relay-record.wmh"), rec.to_bytes()).unwrap();
        std::fs::write(d.join("relay-ledger.wml"), rec.ledger_bytes()).unwrap();
    }
    server.shutdown().await;
}