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

//! What a publisher sends right behind its SUBSCRIBE_OK reaches the subscriber the relay
//! subscribed for. The relay answers that subscriber once it has the publisher's answer, and
//! the publisher may send from the instant it gave it: its first datagrams can reach the
//! relay before the SUBSCRIBE_OK that names their alias is read, and its first objects
//! before the subscriber is attached to the track. They went to nobody.

mod common;

use common::*;
use joule_moq::draft22::data::{DatagramContent, ObjectDatagram, ObjectStatus, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{ControlMessage, LocationFilter, Parameter, Parameters, SubscribeOk};
use joule_moq_relay::client::{Client, ClientOptions};
use std::time::Duration;
use tokio::time::timeout;

const ROUNDS: u64 = 12;
const EACH: u64 = 4;

fn datagram(alias: u64, object: u64) -> bytes::Bytes {
    let d = ObjectDatagram { track_alias: alias, group_id: 0, object_id: object, publisher_priority: Some(1), properties: None, end_of_group: false, content: DatagramContent::Payload(format!("datagram {object}").into_bytes()) };
    bytes::Bytes::from(d.encode().expect("a datagram"))
}

/// Answer the relay's SUBSCRIBE on `p` under `alias`, and send `first` with nothing awaited
/// in between but the answer's own write.
async fn answer_and_send(p: &Client, alias: u64, first: impl AsyncFnOnce(&Client)) -> (joule_moq_relay::wire::ReqTx, joule_moq_relay::wire::ReqRx) {
    let (tx, rx, asked) = p.next_request(STEP).await.expect("the relay's SUBSCRIBE");
    assert!(matches!(asked, ControlMessage::Subscribe(_)), "{}", asked.name());
    tx.send(&ControlMessage::SubscribeOk(SubscribeOk { track_alias: alias, params: Parameters::new(), track_properties: Default::default() })).await.expect("SUBSCRIBE_OK");
    first(p).await;
    (tx, rx)
}

async fn next_datagram(c: &Client, what: &str) -> ObjectDatagram {
    let b = timeout(STEP, c.conn.recv_datagram()).await.unwrap_or_else(|_| panic!("{what}: no datagram in time")).expect("a datagram");
    match joule_moq::draft22::data::Datagram::decode(&b).expect("a datagram that reads") {
        joule_moq::draft22::data::Datagram::Object(d) => d,
        _ => panic!("not an object datagram"),
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn datagrams_sent_behind_subscribe_ok_reach_the_subscriber() {
    timeout(Duration::from_secs(180), async {
        let server = relay(|_| {}).await;
        let space = ns(&["first", "datagrams"]);
        let (p, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let _ns = p.publish_namespace(space.clone(), Parameters::new()).await.expect("PUBLISH_NAMESPACE");
        for round in 0..ROUNDS {
            let (s, _es) = connect(&server, Via::Quic, ClientOptions::default()).await;
            let (alias, track) = (40 + round, name(&format!("track-{round}")));
            let send = async |p: &Client| {
                for o in 0..EACH {
                    p.conn.send_datagram(datagram(alias, o)).expect("sent");
                }
            };
            let (sub, _held) = tokio::join!(s.subscribe(space.clone(), track, Parameters::new()), answer_and_send(&p, alias, send));
            let (ok, _stx, _srx) = sub.expect("SUBSCRIBE");
            for o in 0..EACH {
                let d = next_datagram(&s, &format!("round {round}, datagram {o}")).await;
                assert_eq!((d.track_alias, d.group_id, d.object_id), (ok.track_alias, 0, o), "round {round}: each datagram, once, in the order sent");
            }
            assert!(timeout(Duration::from_millis(100), s.conn.recv_datagram()).await.is_err(), "round {round}: a datagram was sent twice");
            s.close(0, "done");
        }
        p.close(0, "done");
        server.shutdown().await;
    })
    .await
    .expect("in time");
}

fn object(id: u64) -> SubgroupObject {
    SubgroupObject { object_id: id, properties: Default::default(), status: ObjectStatus::Normal, payload: format!("object {id}").into_bytes() }
}

/// A subgroup's stream opened right behind SUBSCRIBE_OK, some of it before the subscriber is
/// attached and the rest after: the subscriber is sent one stream that begins with the
/// subgroup's first object, holds every object once and in order, and ends with a FIN. A
/// subscription with the Next Object filter is ranged from the Largest Location before what
/// it is sent, and its SUBSCRIBE_OK states that one.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_subgroup_sent_behind_subscribe_ok_reaches_the_subscriber_whole() {
    timeout(Duration::from_secs(180), async {
        let server = relay(|_| {}).await;
        let space = ns(&["first", "streams"]);
        let (p, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
        let _ns = p.publish_namespace(space.clone(), Parameters::new()).await.expect("PUBLISH_NAMESPACE");
        for round in 0..ROUNDS {
            let (s, _es) = connect(&server, Via::Quic, ClientOptions::default()).await;
            let (alias, track) = (70 + round, name(&format!("track-{round}")));
            let header = SubgroupHeader { track_alias: alias, group_id: 3, subgroup_id: SubgroupId::Explicit(0), publisher_priority: Some(1), properties: true, end_of_group: false, first_object: true };
            // The stream is held open across the answer: its first half is written at once,
            // its second after the subscriber has its SUBSCRIBE_OK.
            let (half_tx, half_rx) = tokio::sync::oneshot::channel::<()>();
            let send = async |p: &Client| {
                let mut b = Vec::new();
                header.encode(&mut b);
                let mut c = joule_moq::draft22::data::SubgroupCodec::new(header);
                for o in 0..EACH {
                    c.encode_object(&object(o), &mut b).expect("an object");
                }
                let mut st = p.conn.open_uni().await.expect("a stream");
                st.write_all(&b).await.expect("written");
                tokio::spawn(async move {
                    let _ = half_rx.await;
                    let mut b = Vec::new();
                    for o in EACH..2 * EACH {
                        c.encode_object(&object(o), &mut b).expect("an object");
                    }
                    st.write_all(&b).await.expect("written");
                    st.finish().expect("finished");
                });
            };
            let params = if round % 2 == 1 { Parameters::new().with(Parameter::LocationFilter(LocationFilter::NextObject)) } else { Parameters::new() };
            let (sub, _held) = tokio::join!(s.subscribe(space.clone(), track, params), answer_and_send(&p, alias, send));
            let (ok, _stx, _srx) = sub.expect("SUBSCRIBE");
            assert_eq!(ok.params.largest_object(), None, "round {round}: the Largest Location before what the subscription is sent");
            half_tx.send(()).expect("the second half");
            let got = collect(&s, STEP, |g| !g.is_empty()).await;
            assert_eq!(got.len(), 1, "round {round}: one stream");
            let g = &got[0];
            assert_eq!((g.header.track_alias, g.header.group_id, g.reset), (ok.track_alias, 3, None), "round {round}");
            assert!(g.header.first_object, "round {round}: the stream begins with the subgroup's first object");
            assert_eq!(g.objects.iter().map(|o| o.object_id).collect::<Vec<_>>(), (0..2 * EACH).collect::<Vec<_>>(), "round {round}: every object once, in order");
            assert!(g.objects.iter().all(|o| o.payload == format!("object {}", o.object_id).into_bytes()), "round {round}");
            assert!(s.next_data(Duration::from_millis(100)).await.is_none(), "round {round}: a second stream of the subgroup");
            s.close(0, "done");
        }
        p.close(0, "done");
        server.shutdown().await;
    })
    .await
    .expect("in time");
}