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

//! A track switch through the relay (`joule-moq`'s `switching`: `draft-ugur-moq-fast-switching-00`
//! planned as the two requests `draft-ietf-moq-transport-23` has; `draft-gurel-moq-track-switching-01`'s
//! aim, no overlap and no gap between two alternatives).
//!
//! A publisher sends two renditions of one source, group for group. A subscriber on the lower
//! one plans a switch from the object it has just received: a SUBSCRIBE of the target starting
//! at the Switch Point's Group, object 0, and a REQUEST_UPDATE ending the current subscription
//! at the expected End Group. The relay then forwards the current track through its End Group
//! and no further, and the target from the Switch Point on: released in Frame Number order,
//! every frame is presented once, none skipped and none twice, and the handoff is on a Group
//! boundary of both tracks.
//!
//! Planted: a switch mid-group. The plan's own check refuses a start that is not object 0 of a
//! Switch Point (`switch_start_mid_group`). A subscriber that subscribes mid-group all the same
//! is forwarded exactly what it asked for, a Group without its first object: a Location filter
//! is an inclusive range and "a publisher MUST NOT send objects from outside the requested
//! range" (`draft-ietf-moq-transport-23` §3.3.1), and delivery "can begin mid-group" (§3.5), so
//! a relay that rounded the start down to the Group's boundary would be wrong. The release
//! buffer then takes each object forwarded and holds it, presenting none: it presents in Frame
//! Number order and the Group's first frame never comes. It refuses nothing until its bound,
//! where the object that would pass it is refused `switch_buffer_full`. Not one frame of the
//! track is presented, before or after.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{ControlMessage, Location, LocationFilter, Parameter, Parameters, RequestUpdate};
use joule_moq::switching::{Push, Quality, ReleaseBuffer, SwitchRefusal, SwitchSet, SwitchTrack};
use joule_moq_relay::client::{Client, ClientOptions};
use std::time::Duration;
use tokio::time::timeout;

const SPACING: u64 = 2;
const LO: usize = 0;
const HI: usize = 1;

fn set() -> SwitchSet {
    let track = |name: &str, bitrate| SwitchTrack { name: name.into(), spacing: SPACING, quality: Quality { bitrate: Some(bitrate), ..Quality::default() } };
    SwitchSet::new(1, vec![track("lo", 500_000), track("hi", 3_000_000)]).expect("a Switch Set")
}

/// Group `g` of both renditions: `SPACING` objects each, on one subgroup stream a track.
async fn publish_group(p: &Client, g: u64) {
    for (alias, tag) in [(10u64, "lo"), (11, "hi")] {
        let h = SubgroupHeader { track_alias: alias, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(128), properties: false, end_of_group: true, first_object: true };
        let objects: Vec<SubgroupObject> = (0..SPACING).map(|o| SubgroupObject::normal(o, format!("{tag}:{g}:{o}").into_bytes())).collect();
        p.send_subgroup(h, &objects).await.expect("a group");
    }
}

/// Every object the subscriber is forwarded until `quiet` passes with none: (track, Location).
async fn drain(c: &Client, aliases: &[(u64, usize)], quiet: Duration) -> Vec<(usize, Location)> {
    let mut out = Vec::new();
    while let Some(r) = c.next_data(quiet).await {
        let Some(got) = read_stream(r).await else { continue };
        let track = aliases.iter().find(|(a, _)| *a == got.header.track_alias).expect("a subscribed track").1;
        for o in &got.objects {
            let tag = if track == LO { "lo" } else { "hi" };
            assert_eq!(o.payload, format!("{tag}:{}:{}", got.header.group_id, o.object_id).into_bytes());
            out.push((track, Location::new(got.header.group_id, o.object_id)));
        }
    }
    out
}

fn from(group: u64, object: u64) -> Parameters {
    Parameters::new().with(Parameter::LocationFilter(LocationFilter::AbsoluteStart { start: Location::new(group, object) }))
}

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

async fn run() {
    let set = set();
    let space = ns(&["live", "ladder"]);
    let server = relay(|_| {}).await;
    let (p, _pe) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let _lo = p.publish(space.clone(), name("lo"), 10, Parameters::new(), track_props(128, None)).await.unwrap();
    let _hi = p.publish(space.clone(), name("hi"), 11, Parameters::new(), track_props(128, None)).await.unwrap();

    let (c, _ce) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let start = Location::new(1, 0);
    let (lo_ok, lo_tx, _lo_rx) = c.subscribe(space.clone(), name("lo"), from(1, 0)).await.unwrap();
    let mut buffer = ReleaseBuffer::new(set.clone(), LO, start).expect("a buffer");
    let mut presented: Vec<u64> = Vec::new();
    let mut arrived: Vec<(usize, Location)> = Vec::new();

    // Groups 1 to 3 on the current track.
    for g in 1..=3 {
        publish_group(&p, g).await;
    }
    let mut aliases = vec![(lo_ok.track_alias, LO)];
    let got = drain(&c, &aliases, Duration::from_millis(700)).await;
    assert_eq!(got.len(), 6);
    for (t, at) in &got {
        presented.extend(buffer.push(*t, *at).expect("the current track's object").released.iter().map(|r| r.frame));
    }
    arrived.extend(got);

    // The subscriber has just received the last object of Group 3, and plans.
    let at = Location::new(3, 1);
    let plan = set.plan(LO, at, start, HI).expect("a plan");
    assert_eq!((plan.target_group, plan.end_group), (4, 3), "the Switch Point is the next Group; the current track ends with the Group in hand");
    assert_eq!(plan.subscribe, LocationFilter::AbsoluteStart { start: Location::new(4, 0) });
    buffer.begin(&plan).expect("the switch begins");
    let (hi_ok, _hi_tx, _hi_rx) = c.subscribe(space.clone(), name("hi"), Parameters::new().with(Parameter::LocationFilter(plan.subscribe))).await.unwrap();
    lo_tx.send(&ControlMessage::RequestUpdate(RequestUpdate { request_id: c.next_request_id(), params: Parameters::new().with(Parameter::LocationFilter(plan.update)) })).await.unwrap();
    aliases.push((hi_ok.track_alias, HI));
    tokio::time::sleep(Duration::from_millis(300)).await;

    // The publisher goes on with both renditions.
    for g in 4..=6 {
        publish_group(&p, g).await;
    }
    let got = drain(&c, &aliases, Duration::from_millis(700)).await;
    for (t, at) in &got {
        presented.extend(buffer.push(*t, *at).expect("an object of the switch").released.iter().map(|r| r.frame));
    }
    arrived.extend(got);

    // No overlap and no gap on the wire: the current track through its End Group, the target
    // from the Switch Point, object 0.
    assert!(arrived.iter().filter(|(t, _)| *t == LO).all(|(_, l)| l.group <= plan.end_group), "nothing of the current track past its End Group");
    assert_eq!(arrived.iter().filter(|(t, _)| *t == HI).map(|(_, l)| *l).min(), Some(Location::new(4, 0)), "the target starts on its Group boundary");
    assert_eq!(arrived.len(), 12, "six frames of each, none sent on both tracks");
    // And presented once each, in order.
    assert_eq!(presented, (1..=12).collect::<Vec<u64>>(), "every frame once, in Frame Number order");
    assert_eq!(buffer.current(), HI);

    // ── Planted: a switch mid-group ──
    assert_eq!(set.check_start(HI, Location::new(8, 1)).unwrap_err(), SwitchRefusal::all().into_iter().find(|r| r.code() == "switch_start_mid_group").unwrap());
    let (m, _me) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let (m_lo, _t1, _r1) = m.subscribe(space.clone(), name("lo"), from(7, 0)).await.unwrap();
    let (m_hi, _t2, _r2) = m.subscribe(space.clone(), name("hi"), from(8, 1)).await.unwrap();
    for g in 7..=8 {
        publish_group(&p, g).await;
    }
    let got = drain(&m, &[(m_lo.track_alias, LO), (m_hi.track_alias, HI)], Duration::from_millis(700)).await;
    let hi: Vec<Location> = got.iter().filter(|(t, _)| *t == HI).map(|(_, l)| *l).collect();
    assert_eq!(hi, [Location::new(8, 1)], "the relay forwards from where it was asked: a Group without its first object");
    let lo: Vec<Location> = got.iter().filter(|(t, _)| *t == LO).map(|(_, l)| *l).collect();
    assert_eq!(lo.len(), 4, "the track subscribed on its Group boundary is forwarded whole: {lo:?}");
    // The buffer of a subscriber that took the target from Group 8, fed what the relay
    // forwarded: the object is taken and held, and nothing is presented.
    let first = set.frame_number(HI, Location::new(8, 0)).expect("a frame");
    let mut naive = ReleaseBuffer::new(set.clone(), HI, Location::new(8, 0)).expect("a buffer");
    for at in &hi {
        assert_eq!(naive.push(HI, *at), Ok(Push { released: vec![], dropped: false }), "an object ahead of the frame awaited is held, not refused and not presented");
    }
    assert_eq!((naive.held(), naive.awaiting()), (1, first), "the Group's first frame is still awaited");
    // The publisher goes on, and so does the subscription: every later object is held behind
    // the frame that never came, up to the buffer's bound, and the one past it is refused.
    for g in 9..=11 {
        publish_group(&p, g).await;
    }
    let got = drain(&m, &[(m_lo.track_alias, LO), (m_hi.track_alias, HI)], Duration::from_millis(700)).await;
    let mut later: Vec<Location> = got.iter().filter(|(t, _)| *t == HI).map(|(_, l)| *l).collect();
    later.sort();
    assert_eq!(later, (9..=11).flat_map(|g| (0..SPACING).map(move |o| Location::new(g, o))).collect::<Vec<_>>(), "the Groups after it are forwarded whole");
    let outcomes: Vec<_> = later.iter().map(|at| naive.push(HI, *at)).collect();
    let taken = Ok(Push { released: vec![], dropped: false });
    // 2 * (Gm + 1) = 6 objects are held: the one of Group 8 and five more.
    assert_eq!(outcomes[..5], [taken.clone(), taken.clone(), taken.clone(), taken.clone(), taken], "held, none presented");
    assert_eq!(outcomes[5], Err(SwitchRefusal::BufferFull), "the object past the buffer's bound is refused");
    assert_eq!((naive.held(), naive.awaiting(), naive.current()), (6, first, HI), "not one frame was presented");

    for c in [&c, &m, &p] {
        c.close(0, "done");
    }
    server.shutdown().await;
}