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

//! Probe tracks served by the relay under `.session` (`draft-gurel-moq-probing-00`): the
//! catalog, a paced and a greedy downlink probe, an uplink probe, and each refusal with the
//! code the draft's Table 6 gives. `joule_moq::probe` holds the arithmetic and its corpus;
//! this holds what the relay does with it.
//!
//! Planted, each with its own outcome: a probe past a per-probe limit; a second probe while
//! one runs, and one too soon after it; a probe past the relay-wide limit from another
//! session; an unknown preset; a FETCH of a probe track; a SUBSCRIBE in the uplink namespace;
//! an uplink probe with no FORWARD 0; a preset named in an uplink probe; an uplink probe that
//! sends more than its `b` (ended, the session standing). And a relay with probes off
//! answers `DOES_NOT_EXIST`.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupCodec, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{ControlMessage, Parameter, Parameters, Publish, PublishDone, PublishDoneCode, RequestErrorCode, StreamHeader, TrackRequest};
use joule_moq::probe::{self, Arrival, Direction, Preset};
use joule_moq::{TrackName, TrackNamespace};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::config::ProbePreset;
use joule_moq_relay::wire::ReqRx;
use joule_moq_relay::{RelayConfig, RelayServer};
use std::time::Duration;
use tokio::time::timeout;

const LIMIT: Duration = Duration::from_secs(120);

fn down() -> TrackNamespace {
    TrackNamespace(vec![probe::SESSION_FIELD.to_vec(), probe::PROBE_FIELD.to_vec()])
}

fn up() -> TrackNamespace {
    TrackNamespace(vec![probe::SESSION_FIELD.to_vec(), probe::SINK_FIELD.to_vec()])
}

fn presets() -> Vec<ProbePreset> {
    vec![
        ProbePreset { name: "sd-check".into(), label: Some("400 kbit/s, video-like, 400 ms".into()), probe: "d400-f50-g10-k3-r400000".into() },
        ProbePreset { name: "burst".into(), label: None, probe: "b100000".into() },
    ]
}

async fn probing(f: impl FnOnce(&mut RelayConfig)) -> RelayServer {
    relay(|c| {
        c.probes.enabled = true;
        c.probes.presets = presets();
        c.probes.min_gap_ms = 0;
        f(c);
    })
    .await
}

/// A request's first answer: the REQUEST_ERROR's code, Retry Interval and reason, or the
/// message that was not one.
async fn refusal(c: &Client, m: ControlMessage) -> (RequestErrorCode, u64, String) {
    let (_tx, mut rx) = c.request(&m).await.expect("the request stream");
    match Client::response(&mut rx).await.expect("a response") {
        ControlMessage::RequestError(e) => (e.code, e.retry_interval, String::from_utf8_lossy(&e.reason).into_owned()),
        other => panic!("not refused: {}", other.name()),
    }
}

fn subscribe(c: &Client, ns: TrackNamespace, name: &str) -> ControlMessage {
    ControlMessage::Subscribe(TrackRequest { request_id: c.next_request_id(), namespace: ns, name: TrackName(name.as_bytes().to_vec()), params: Parameters::new().with(Parameter::SubscriberPriority(probe::SUBSCRIBER_PRIORITY)) })
}

/// What a downlink probe delivered: each stream's group and objects, the arrivals as the
/// draft's §11 measures them, and the PUBLISH_DONE.
struct Probed {
    streams: Vec<Got>,
    arrivals: Vec<Arrival>,
    done: PublishDone,
    alias: u64,
}

async fn run_probe(c: &Client, ns: TrackNamespace, name: &str) -> Probed {
    let (ok, _tx, mut rx) = c.subscribe(ns, TrackName(name.as_bytes().to_vec()), Parameters::new().with(Parameter::SubscriberPriority(probe::SUBSCRIBER_PRIORITY))).await.unwrap_or_else(|e| panic!("the probe {name}: {e}"));
    let done = async {
        match timeout(STEP, rx.next()).await {
            Ok(Ok(Some(ControlMessage::PublishDone(d)))) => d,
            other => panic!("the probe {name} did not end with PUBLISH_DONE: {other:?}"),
        }
    };
    let read = async {
        let mut streams = Vec::new();
        let mut arrivals = Vec::new();
        // Until the PUBLISH_DONE's count of streams is read; the count comes after them.
        loop {
            let Some(mut r) = c.next_data(Duration::from_millis(1500)).await else { break };
            let Some(StreamHeader::Subgroup(header)) = r.header.clone() else { panic!("not a subgroup stream") };
            let mut objects = Vec::new();
            loop {
                match timeout(STEP, r.next()).await.expect("a read in time").expect("a read") {
                    Some(joule_moq::draft22::UniEvent::SubgroupObject(o)) => {
                        let sent_us = o.properties.find(probe::TIMESTAMP_PROPERTY).ok().flatten().and_then(|k| k.as_varint()).expect("a LOC Timestamp on every probe object");
                        let arrived_us = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_micros() as u64;
                        arrivals.push(Arrival { location: joule_moq::draft22::Location::new(header.group_id, o.object_id), bytes: o.payload.len() as u64, sent_us, arrived_us });
                        objects.push(o);
                    }
                    Some(_) => {}
                    None => break,
                }
            }
            streams.push(Got { header, objects, reset: None, arrived: tokio::time::Instant::now() });
        }
        (streams, arrivals)
    };
    let (done, (streams, arrivals)) = tokio::join!(done, read);
    Probed { streams, arrivals, done, alias: ok.track_alias }
}

/// Hold a probe's delivery to what its name resolves to, object for object (§6, §8).
fn hold(p: &Probed, name: &str) {
    let r = probe::resolve(&probe::parse(name.as_bytes()).unwrap()).unwrap();
    assert_eq!(p.done.status, PublishDoneCode::TRACK_ENDED, "{name}");
    assert_eq!(p.done.stream_count, r.group_count(), "{name}: a stream a group");
    assert_eq!(p.streams.len() as u64, r.group_count(), "{name}");
    assert_eq!(p.arrivals.len() as u64, r.object_count(), "{name}");
    for s in &p.streams {
        assert_eq!((s.header.track_alias, s.header.resolve_subgroup_id(0), s.header.publisher_priority), (p.alias, probe::SUBGROUP_ID, Some(probe::PUBLISHER_PRIORITY)), "{name}: Subgroup 0 at the lowest priority");
        for o in &s.objects {
            assert!(o.payload.iter().all(|b| *b == 0), "{name}: a payload of zeros");
            assert_eq!(o.properties.0.len(), 1, "{name}: the Timestamp and no other property");
        }
    }
    let mut by_location = p.arrivals.clone();
    by_location.sort_by_key(|a| a.location);
    for (i, a) in by_location.iter().enumerate() {
        let want = r.object(i as u64).expect("an object of the probe");
        assert_eq!((a.location, a.bytes), (want.location, want.payload_len), "{name}: object {i}");
    }
    assert_eq!(by_location.iter().map(|a| a.bytes).sum::<u64>(), r.total_bytes(), "{name}");
    let m = probe::measure(&p.arrivals, Some(r.object_count())).expect("a measurement");
    assert_eq!((m.objects, m.missing), (r.object_count(), Some(0)), "{name}");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_catalog_and_downlink_probes() {
    timeout(LIMIT, async {
        let server = probing(|_| {}).await;
        let (c, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;

        // The catalog: the proof of support (§9.2.1), the presets and nothing else.
        let (ok, _ctx, _crx) = c.subscribe(down(), TrackName(probe::CATALOG_TRACK.to_vec()), Parameters::new()).await.expect("the probe catalog");
        let got = collect(&c, STEP, |g| !g.is_empty()).await;
        assert_eq!((got[0].header.track_alias, got[0].header.group_id, got[0].objects.len(), got[0].reset), (ok.track_alias, 0, 1, None));
        let doc: serde_json::Value = serde_json::from_slice(&got[0].objects[0].payload).expect("JSON");
        let want: Vec<Preset> = presets().into_iter().map(|p| Preset { name: p.name, label: p.label, probe_name: p.probe }).collect();
        assert_eq!(probe::read_catalog(&doc).expect("a probe catalog"), want);
        assert_eq!(doc, probe::catalog(&want).unwrap(), "the members the draft names and no limit");

        // A paced probe by its Dynamic Probe Name: 20 objects of 1000 bytes, one every 20 ms.
        let started = std::time::Instant::now();
        let p = run_probe(&c, down(), "d400-r400000").await;
        hold(&p, "d400-r400000");
        let (first, last) = (p.arrivals.iter().map(|a| a.sent_us).min().unwrap(), p.arrivals.iter().map(|a| a.sent_us).max().unwrap());
        assert!(last - first >= 19 * 20_000 - 100_000, "released on the schedule, not at once: {} us from the first object to the last", last - first);
        assert!(started.elapsed() >= Duration::from_millis(350), "{:?}", started.elapsed());

        // The same in the equivalent form: the empty namespace and `.probe/` (§9.1).
        let p = run_probe(&c, TrackNamespace(vec![]), ".probe/d200-r400000").await;
        hold(&p, "d200-r400000");

        // A preset: groups of ten, the first object of each three times the others.
        let p = run_probe(&c, down(), "sd-check").await;
        hold(&p, "d400-f50-g10-k3-r400000");
        assert_eq!(p.streams.len(), 2);

        // A greedy probe: 100 000 bytes in objects of 16 384, the last one cut, one group.
        let p = run_probe(&c, down(), "burst").await;
        hold(&p, "b100000");
        assert_eq!((p.arrivals.len(), p.arrivals.iter().map(|a| a.bytes).min()), (7, Some(100_000 - 6 * 16_384)));

        // Nothing of a probe is a track of the relay's: none cached, none for another session.
        let key = (down(), TrackName(b"d400-r400000".to_vec()));
        assert_eq!(server.relay().largest(&key), None);
        let s = server.relay().stats();
        assert_eq!((s.probes_served, s.probes_refused), (4, 0), "{s:?}");
        let bytes = |name: &str| probe::resolve(&probe::parse(name.as_bytes()).unwrap()).unwrap().total_bytes();
        assert_eq!(s.probe_bytes_sent, bytes("d400-r400000") + bytes("d200-r400000") + bytes("d400-f50-g10-k3-r400000") + bytes("b100000"), "{s:?}");
        c.close(0, "done");
        server.shutdown().await;
    })
    .await
    .expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn each_refusal_has_the_drafts_code() {
    timeout(LIMIT, async {
        let server = probing(|c| {
            c.probes.min_gap_ms = 400;
            c.probes.relay_concurrent = 1;
        })
        .await;
        let (c, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let (other, _eo) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let none = RequestErrorCode::DOES_NOT_EXIST;

        // Invalid, and past a per-probe limit: DOES_NOT_EXIST, not to be sent again.
        assert_eq!(refusal(&c, subscribe(&c, down(), "no-such-preset")).await, (none, 0, "probe_name_unknown".into()));
        assert_eq!(refusal(&c, subscribe(&c, down(), "d1000")).await.0, none, "neither a rate nor bytes");
        assert_eq!(refusal(&c, subscribe(&c, down(), "d6000-r300000")).await, (none, 0, "probe_duration_over_limit".into()));
        assert_eq!(refusal(&c, subscribe(&c, down(), "d1000-r60000000")).await, (none, 0, "probe_rate_over_limit".into()));
        // Another request in the namespace, and a SUBSCRIBE in the uplink one: NOT_SUPPORTED.
        let fetch = ControlMessage::Fetch(TrackRequest { request_id: c.next_request_id(), namespace: down(), name: TrackName(b"d400-r400000".to_vec()), params: Parameters::new() });
        assert_eq!(refusal(&c, fetch).await.0, RequestErrorCode::NOT_SUPPORTED);
        assert_eq!(refusal(&c, subscribe(&c, up(), "d400-r400000")).await.0, RequestErrorCode::NOT_SUPPORTED);
        // Another track under `.session`: the relay's to answer, and it has none.
        let elsewhere = TrackNamespace(vec![probe::SESSION_FIELD.to_vec(), b".other".to_vec()]);
        assert_eq!(refusal(&c, subscribe(&c, elsewhere, "x")).await, (none, 0, "session_track_unknown".into()));

        // While a probe runs: a second on its session, and one from another session past the
        // relay-wide limit of one, are refused EXCESSIVE_LOAD with the time to wait.
        let (ok, _tx, mut rx) = c.subscribe(down(), TrackName(b"d600-r200000".to_vec()), Parameters::new()).await.expect("a probe");
        let (code, retry, why) = refusal(&c, subscribe(&c, down(), "d200-r200000")).await;
        assert_eq!((code, why.as_str()), (RequestErrorCode::EXCESSIVE_LOAD, "probe_session_concurrent"));
        assert!(retry > 400 && retry <= 600 + 400 + 1, "the probe's end and the time between probes, plus one: {retry}");
        let (code, retry, why) = refusal(&other, subscribe(&other, down(), "d200-r200000")).await;
        assert_eq!((code, why.as_str()), (RequestErrorCode::EXCESSIVE_LOAD, "probe_relay_concurrent"));
        assert!(retry > 1 && retry <= 601, "{retry}");
        assert!(matches!(timeout(STEP, rx.next()).await, Ok(Ok(Some(ControlMessage::PublishDone(_))))), "the first probe ends");
        let _ = collect(&c, Duration::from_millis(300), |g| g.iter().any(|g| g.header.track_alias == ok.track_alias)).await;
        // Too soon after it: the time between probes.
        let (code, retry, why) = refusal(&c, subscribe(&c, down(), "d200-r200000")).await;
        assert_eq!((code, why.as_str()), (RequestErrorCode::EXCESSIVE_LOAD, "probe_session_gap"));
        assert!(retry > 1 && retry <= 401, "{retry}");
        tokio::time::sleep(Duration::from_millis(retry)).await;
        let p = run_probe(&c, down(), "d200-r200000").await;
        hold(&p, "d200-r200000");

        // The relay's sessions stand, and it counted each refusal.
        let s = server.relay().stats();
        assert_eq!((s.sessions_open, s.probes_served, s.probes_refused), (2, 2, 10), "{s:?}");
        for x in [&c, &other] {
            x.close(0, "done");
        }
        server.shutdown().await;

        // A relay with probes off answers the catalog DOES_NOT_EXIST (§9.2.1).
        let off = relay(|_| {}).await;
        let (c, _ep) = connect(&off, Via::Quic, ClientOptions::default()).await;
        let e = c.subscribe(down(), TrackName(probe::CATALOG_TRACK.to_vec()), Parameters::new()).await.expect_err("no probes here");
        assert!(e.contains("DOES_NOT_EXIST"), "{e}");
        c.close(0, "done");
        off.shutdown().await;
    })
    .await
    .expect("in time");
}

/// PUBLISH an uplink probe, and the relay's first answer.
async fn publish_probe(c: &Client, name: &str, alias: u64, forward: Option<bool>) -> (joule_moq_relay::wire::ReqTx, ReqRx, ControlMessage) {
    let mut params = Parameters::new();
    if let Some(f) = forward {
        params = params.with(Parameter::Forward(f));
    }
    let m = ControlMessage::Publish(Publish { request_id: c.next_request_id(), namespace: up(), name: TrackName(name.as_bytes().to_vec()), track_alias: alias, params, track_properties: Default::default() });
    let (tx, mut rx) = c.request(&m).await.expect("the request stream");
    let first = Client::response(&mut rx).await.expect("a response");
    (tx, rx, first)
}

/// The relay's REQUEST_UPDATE that resumes an uplink probe (FORWARD 1), answered REQUEST_OK:
/// a publisher sends no object of the probe before it.
async fn resumed(tx: &joule_moq_relay::wire::ReqTx, rx: &mut ReqRx) {
    match timeout(STEP, rx.next()).await {
        Ok(Ok(Some(ControlMessage::RequestUpdate(u)))) => assert_eq!(u.params.forward(), Some(true), "Forward State 1"),
        other => panic!("the relay did not resume the probe: {other:?}"),
    }
    tx.send(&ControlMessage::RequestOk(joule_moq::draft22::RequestOk { params: Parameters::new(), track_properties: Default::default() })).await.expect("REQUEST_OK");
}

/// Send `objects` objects of `bytes` each of an uplink probe, one group, as fast as they go.
async fn send_probe(c: &Client, alias: u64, objects: u64, bytes: usize) {
    let h = SubgroupHeader { track_alias: alias, group_id: 0, subgroup_id: SubgroupId::Zero, publisher_priority: Some(probe::PUBLISHER_PRIORITY), properties: true, end_of_group: true, first_object: true };
    let mut b = Vec::new();
    h.encode(&mut b);
    let mut codec = SubgroupCodec::new(h);
    let mut st = timeout(STEP, c.conn.open_uni()).await.expect("in time").expect("a stream");
    for o in 0..objects {
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_micros() as u64;
        codec.encode_object(&SubgroupObject { object_id: o, properties: probe::object_properties(now), status: joule_moq::draft22::data::ObjectStatus::Normal, payload: vec![0; bytes] }, &mut b).unwrap();
        if timeout(STEP, st.write_all(&b)).await.expect("in time").is_err() {
            return;
        }
        b.clear();
    }
    let _ = st.finish();
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn an_uplink_probe_is_measured_discarded_and_held_to_its_shape() {
    timeout(LIMIT, async {
        let server = probing(|_| {}).await;
        let relay = server.relay().clone();
        let (c, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let (watcher, _ew) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let refused = |m: &ControlMessage| match m {
            ControlMessage::RequestError(e) => (e.code, String::from_utf8_lossy(&e.reason).into_owned()),
            other => panic!("not refused: {}", other.name()),
        };

        // Without FORWARD 0, and with a preset's name: UNINTERESTED, and the session stands.
        let (_t, _r, first) = publish_probe(&c, "b40000", 5, None).await;
        assert_eq!(refused(&first), (RequestErrorCode::UNINTERESTED, "probe_forward_state".into()));
        let (_t, _r, first) = publish_probe(&c, "b40000", 5, Some(true)).await;
        assert_eq!(refused(&first).0, RequestErrorCode::UNINTERESTED);
        let (_t, _r, first) = publish_probe(&c, "burst", 5, Some(false)).await;
        assert_eq!(refused(&first), (RequestErrorCode::UNINTERESTED, "probe_preset_uplink".into()));

        // Accepted: REQUEST_OK, and the relay's REQUEST_UPDATE that sets FORWARD 1, which the
        // publisher answers. The probe's 40 000 bytes arrive and go nowhere.
        let (tx, mut rx, first) = publish_probe(&c, "b40000", 5, Some(false)).await;
        assert!(matches!(first, ControlMessage::RequestOk(_)), "not accepted: {}", first.name());
        resumed(&tx, &mut rx).await;
        // Nobody else can ask for it: the uplink namespace takes PUBLISH and nothing else.
        assert_eq!(refusal(&watcher, subscribe(&watcher, up(), "b40000")).await.0, RequestErrorCode::NOT_SUPPORTED);
        send_probe(&c, 5, 4, 10_000).await;
        assert!(until(STEP, || relay.stats().probe_bytes_received == 40_000).await, "{:?}", relay.stats());
        tx.send(&ControlMessage::PublishDone(PublishDone { status: PublishDoneCode::TRACK_ENDED, stream_count: 1, reason: Vec::new() })).await.expect("PUBLISH_DONE");
        assert!(matches!(timeout(STEP, rx.next()).await, Ok(Ok(None))), "the relay ends the request");
        assert_eq!(relay.largest(&(up(), TrackName(b"b40000".to_vec()))), None, "no track of the relay's");
        let s = relay.stats();
        assert_eq!((s.probes_received, s.probes_ended_for_shape, s.objects_in), (1, 0, 0), "{s:?}");

        // Planted: a probe that sends more than its b. The relay ends it, and the session stands.
        let (tx, mut rx, first) = publish_probe(&c, "b40000", 6, Some(false)).await;
        assert!(matches!(first, ControlMessage::RequestOk(_)), "{}", first.name());
        resumed(&tx, &mut rx).await;
        let sent = std::time::Instant::now();
        send_probe(&c, 6, 6, 10_000).await;
        // At once, for its bytes: well before the end its duration would bring (six seconds).
        match timeout(Duration::from_secs(3), rx.next()).await {
            Ok(Err(joule_moq_relay::wire::StreamError::Transport(_))) => {}
            other => panic!("a probe past its bytes was not ended: {other:?}"),
        }
        assert!(until(STEP, || relay.stats().probes_ended_for_shape == 1).await);
        assert!(sent.elapsed() < Duration::from_secs(3), "{:?}", sent.elapsed());
        // The session stands: its next probe is accepted.
        let (tx, mut rx, first) = publish_probe(&c, "b20000", 7, Some(false)).await;
        assert!(matches!(first, ControlMessage::RequestOk(_)), "{}", first.name());
        resumed(&tx, &mut rx).await;
        send_probe(&c, 7, 2, 10_000).await;
        tx.send(&ControlMessage::PublishDone(PublishDone { status: PublishDoneCode::TRACK_ENDED, stream_count: 1, reason: Vec::new() })).await.expect("PUBLISH_DONE");
        assert!(until(STEP, || relay.stats().probes_received == 3).await);
        assert_eq!(relay.stats().sessions_open, 2);
        let _ = Direction::Uplink;
        for x in [&c, &watcher] {
            x.close(0, "done");
        }
        server.shutdown().await;
    })
    .await
    .expect("in time");
}