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

//! Fan-out order under load (`draft-ietf-moq-transport-22` §5.1.2): two publishers send
//! four tracks of different publisher priorities to three subscribers, each subscribing
//! with its own subscriber priorities and group orders, while the subscribers' transports
//! are held shut. When they open, every object reaches every subscriber in the order
//! §5.1.2 gives: subscriber priority, then publisher priority, then group order.
//!
//! The sessions run over the in-process multiplexer, the WebSocket binding's own, with a
//! gate on what the relay sends each subscriber; the relay writes one object at a time.
//! Each loop here is bounded by a hard timeout.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{GroupOrder, Parameter, Parameters};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::transport::mux::{memory_pair, Gate, MuxConfig};
use joule_moq_relay::{Relay, RelayConfig};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::timeout;

const GROUPS: u64 = 25;

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

async fn attach(relay: &Relay, gated: bool, opts: ClientOptions) -> (Client, Option<Gate>) {
    let mux = MuxConfig { send_high_water: 1, ..MuxConfig::default() };
    let (c, s, gate) = memory_pair(mux, gated);
    if let Some(g) = &gate {
        g.open();
    }
    tokio::spawn(joule_moq_relay::session::serve(relay.clone(), Arc::new(s)));
    let client = timeout(STEP, Client::connect(Arc::new(c), opts)).await.expect("in time").expect("SETUP in memory");
    (client, gate)
}

/// A track: its name, its publisher's alias for it, its publisher priority.
struct T {
    name: &'static str,
    alias: u64,
    priority: u8,
}

const TRACKS: [T; 4] = [
    T { name: "t1", alias: 1, priority: 0 },
    T { name: "t2", alias: 2, priority: 50 },
    T { name: "t3", alias: 3, priority: 0 },
    T { name: "t4", alias: 4, priority: 200 },
];

/// Each subscriber's (subscriber priority, group order) for t1..t4.
fn plan(i: usize) -> [(u8, GroupOrder); 4] {
    use GroupOrder::*;
    match i {
        0 => [(10, Ascending), (10, Ascending), (200, Descending), (5, Ascending)],
        1 => [(128, Descending), (128, Ascending), (128, Ascending), (128, Descending)],
        _ => [(100, Ascending), (100, Descending), (0, Descending), (100, Ascending)],
    }
}

/// Every group of the tracks `js`, one object to a group, each send bounded.
async fn send(c: &Client, js: &[usize]) {
    for g in 0..GROUPS {
        for &j in js {
            let t = &TRACKS[j];
            let h = SubgroupHeader { track_alias: t.alias, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(t.priority), properties: false, end_of_group: true, first_object: true };
            timeout(STEP, c.send_subgroup(h, &[SubgroupObject::normal(0, payload(t.name, g, 200))])).await.expect("a send in time").expect("a send");
        }
    }
}

async fn run() {
    let mut cfg = RelayConfig::loopback();
    cfg.quic_listen = None;
    cfg.websocket_listen = None;
    cfg.limits.write_concurrency = 1;
    let relay = Relay::new(cfg).unwrap();
    let space = ns(&["load"]);

    let (p1, _) = attach(&relay, false, ClientOptions::default()).await;
    let (p2, _) = attach(&relay, false, ClientOptions::default()).await;
    let pubs = [&p1, &p1, &p2, &p2];
    let mut held = Vec::new();
    for (t, p) in TRACKS.iter().zip(pubs) {
        held.push(p.publish(space.clone(), name(t.name), t.alias, Parameters::new(), track_props(t.priority, None)).await.expect("PUBLISH"));
    }

    // Three subscribers; each read in arrival order (lazily, so nothing reorders them).
    let mut subs = Vec::new();
    for i in 0..3 {
        let (c, gate) = attach(&relay, true, ClientOptions { lazy_data: true, ..ClientOptions::default() }).await;
        let mut keys: HashMap<u64, (u8, u8, GroupOrder)> = HashMap::new();
        for (j, t) in TRACKS.iter().enumerate() {
            let (sp, order) = plan(i)[j];
            let params = Parameters::new().with(Parameter::SubscriberPriority(sp)).with(Parameter::GroupOrder(order));
            let (ok, tx, rx) = c.subscribe(space.clone(), name(t.name), params).await.expect("SUBSCRIBE");
            held.push((tx, rx));
            keys.insert(ok.track_alias, (sp, t.priority, order));
        }
        subs.push((c, gate.expect("a gate"), keys));
    }
    for (_, g, _) in &subs {
        g.close();
    }

    // Load: both publishers send every group of their tracks, concurrently.
    tokio::join!(send(&p1, &[0, 1]), send(&p2, &[2, 3]));
    let total = GROUPS * TRACKS.len() as u64;
    assert!(until(STEP, || relay.stats().objects_in >= total).await, "the relay received every object");
    // Settle: everything received has reached the subscribers' schedulers.
    let mut stable = 0;
    assert!(
        until(STEP, || {
            stable = if relay.stats().queued == 0 { stable + 1 } else { 0 };
            stable >= 5
        })
        .await
    );

    for (_, g, _) in &subs {
        g.open();
    }
    let reads = futures::future::join_all(subs.iter().map(|(c, _, _)| collect(c, STEP, |got| got.len() as u64 >= total))).await;
    for (i, (got, (_, _, keys))) in reads.iter().zip(&subs).enumerate() {
        assert_eq!(got.len() as u64, total, "subscriber {i}: every object");
        let key = |x: &Got| {
            let (sp, pp, order) = keys[&x.header.track_alias];
            let rank = match order {
                GroupOrder::Ascending => x.header.group_id,
                GroupOrder::Descending => u64::MAX - x.header.group_id,
            };
            (sp, pp, rank)
        };
        // At most two objects were handed over before the gate shut: one in the pipe and
        // one write in flight. Everything after them waited in the scheduler.
        let ordered: Vec<_> = got.iter().skip(2).map(key).collect();
        let out_of_order = ordered.windows(2).filter(|w| w[0] > w[1]).count();
        assert_eq!(out_of_order, 0, "subscriber {i}: fan-out in §5.1.2 order: {:?}", got.iter().map(|x| (x.header.track_alias, x.header.group_id)).collect::<Vec<_>>());
    }
    drop(held);
}