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

//! A publisher's PUBLISH_DONE can reach the relay before the data streams its Stream
//! Count says the publisher opened: the request stream and the data streams are not
//! ordered with each other (draft-ietf-moq-transport-22 §9.9). The relay waits for those
//! streams, so a subscriber gets every object and then PUBLISH_DONE, not PUBLISH_DONE
//! alone; and a count no stream arrives for ends the track after a bounded wait.
//!
//! A track its publisher ended with TRACK_ENDED and another publisher then publishes is the
//! same track published anew (-22 §2.1: an Object's "contents MUST NOT change over time"):
//! the relay keeps its Largest Location and what it cached, the new publisher's objects reach
//! subscribers, and are not taken for duplicates of a session that has gone
//! (`tests/republished_track.rs` has one that sends other bytes at a Location).

mod common;

use common::*;
use joule_moq::draft::Draft;
use joule_moq::draft22::data::{SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{ControlMessage, Parameters, PublishDone, PublishDoneCode};
use joule_moq_relay::client::ClientOptions;
use joule_moq_relay::relay::PUBLISH_DONE_STREAM_WAIT;
use std::time::Duration;
use tokio::time::{timeout, Instant};

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

/// PUBLISH a track, have a subscriber take it, then send PUBLISH_DONE saying `count`
/// streams and, `stream_after` later (if any), the one stream. Returns what the
/// subscriber read and the PUBLISH_DONE it was sent, with when it came.
async fn done_then_stream(via: Via, count: u64, stream_after: Option<Duration>) -> (Vec<Got>, PublishDone, Duration) {
    let server = relay(|_| {}).await;
    let (publisher, _e1) = connect(&server, via, ClientOptions::default()).await;
    let (subscriber, _e2) = connect(&server, via, ClientOptions::default()).await;
    let space = ns(&["done", "streams"]);
    let (ptx, _prx) = publisher.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH");
    let (ok, _stx, mut srx) = subscriber.subscribe(space, name("track"), Parameters::new()).await.expect("SUBSCRIBE");
    let start = Instant::now();
    // When PUBLISH_DONE reaches the subscriber, read beside its data streams.
    let done_rx = tokio::spawn(async move {
        let d = timeout(STEP, srx.next()).await;
        (d, Instant::now())
    });
    ptx.send(&ControlMessage::PublishDone(PublishDone { status: PublishDoneCode::TRACK_ENDED, stream_count: count, reason: Vec::new() })).await.expect("PUBLISH_DONE");
    let got = match stream_after {
        Some(after) => {
            tokio::time::sleep(after).await;
            let h = SubgroupHeader { track_alias: 1, group_id: 0, subgroup_id: SubgroupId::Zero, publisher_priority: Some(0), properties: false, end_of_group: true, first_object: true };
            publisher.send_subgroup(h, &[SubgroupObject::normal(0, b"after the PUBLISH_DONE".to_vec())]).await.expect("the stream");
            collect(&subscriber, STEP, |g| !g.is_empty()).await
        }
        None => Vec::new(),
    };
    let (d, at) = done_rx.await.expect("the reader");
    let done = match d {
        Ok(Ok(Some(ControlMessage::PublishDone(d)))) => d,
        other => panic!("{via:?}: no PUBLISH_DONE downstream: {other:?}"),
    };
    assert!(got.iter().all(|g| g.header.track_alias == ok.track_alias), "{via:?}: the subscriber's alias");
    (got, done, at - start)
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_stream_after_the_publish_done_still_reaches_the_subscriber() {
    timeout(LIMIT, async {
        for via in [Via::Quic, Via::QuicDraft(Draft::D18), Via::WtDraft(Draft::D18), Via::QuicDraft(Draft::D16)] {
            let (got, done, _) = done_then_stream(via, 1, Some(Duration::from_millis(300))).await;
            let objects: Vec<&SubgroupObject> = got.iter().flat_map(|g| &g.objects).collect();
            assert_eq!(objects.len(), 1, "{via:?}: the object sent after PUBLISH_DONE");
            assert_eq!(objects[0].payload, b"after the PUBLISH_DONE", "{via:?}");
            assert_eq!((done.status, done.stream_count), (PublishDoneCode::TRACK_ENDED, 1), "{via:?}: the relay's own count of what it sent");
        }
    })
    .await
    .expect("in time");
}

/// A count that no stream ever meets: the track still ends, after the bounded wait, and
/// the relay's PUBLISH_DONE counts the streams it sent (none).
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_count_no_stream_meets_ends_after_the_wait() {
    timeout(LIMIT, async {
        let (got, done, took) = done_then_stream(Via::Quic, 1, None).await;
        assert!(got.is_empty());
        assert_eq!((done.status, done.stream_count), (PublishDoneCode::TRACK_ENDED, 0));
        assert!(took >= PUBLISH_DONE_STREAM_WAIT, "ended after the wait, not before: {took:?}");
    })
    .await
    .expect("in time");
}

/// Stream Count 0 and the unknown count end the track at once.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn zero_and_unknown_counts_end_at_once() {
    timeout(LIMIT, async {
        for count in [0, u64::MAX] {
            let (got, done, took) = done_then_stream(Via::Quic, count, None).await;
            assert!(got.is_empty());
            assert_eq!(done.stream_count, 0, "count {count}");
            assert!(took < PUBLISH_DONE_STREAM_WAIT, "count {count}: ended without the wait, in {took:?}");
        }
    })
    .await
    .expect("in time");
}

/// One publisher publishes a track, sends group 0 and ends it with TRACK_ENDED, and
/// leaves; a second publishes the same track and sends group 0 again, the same object, and
/// group 1. A subscriber of the second from {0, 0} is told the track's Largest Location and
/// gets both groups from the second publisher: the subgroup the first sent is owned by no
/// session that has gone.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_track_published_again_after_it_ended_is_the_same_track() {
    timeout(LIMIT, async {
        for via in [Via::Quic, Via::WtDraft(Draft::D18)] {
            let server = relay(|_| {}).await;
            let space = ns(&["done", "again"]);
            let h = SubgroupHeader { track_alias: 1, group_id: 0, subgroup_id: SubgroupId::Zero, publisher_priority: Some(0), properties: false, end_of_group: true, first_object: true };
            {
                let (first, _e) = connect(&server, via, ClientOptions::default()).await;
                let (tx, _rx) = first.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH");
                first.send_subgroup(h, &[SubgroupObject::normal(0, b"the first instance".to_vec())]).await.expect("the first group 0");
                let relay = server.relay().clone();
                let key = (space.clone(), name("track"));
                assert!(until(STEP, || relay.largest(&key).is_some()).await, "{via:?}: the relay took the first group 0");
                tx.send(&ControlMessage::PublishDone(PublishDone { status: PublishDoneCode::TRACK_ENDED, stream_count: 1, reason: Vec::new() })).await.expect("PUBLISH_DONE");
                assert!(until(STEP, || relay.track_ended(&key) == Some(PublishDoneCode::TRACK_ENDED)).await, "{via:?}: the first publisher ended its track");
                first.close(0, "done");
            }
            let (second, _e2) = connect(&server, via, ClientOptions::default()).await;
            let (subscriber, _e3) = connect(&server, via, ClientOptions::default()).await;
            let _p = second.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH again");
            let from_start = Parameters::new().with(joule_moq::draft22::Parameter::LocationFilter(joule_moq::draft22::LocationFilter::AbsoluteStart { start: joule_moq::draft22::Location::new(0, 0) }));
            let (ok, _stx, _srx) = subscriber.subscribe(space.clone(), name("track"), from_start).await.expect("SUBSCRIBE");
            assert_eq!(ok.params.largest_object(), Some(joule_moq::draft22::Location::new(0, 0)), "{via:?}: the track's Largest Location, kept");
            second.send_subgroup(h, &[SubgroupObject::normal(0, b"the first instance".to_vec())]).await.expect("group 0 again");
            second.send_subgroup(SubgroupHeader { group_id: 1, ..h }, &[SubgroupObject::normal(0, b"group 1".to_vec())]).await.expect("group 1");
            let got = collect(&subscriber, STEP, |g| g.len() >= 2).await;
            let mut payloads: Vec<&[u8]> = got.iter().flat_map(|g| &g.objects).map(|o| o.payload.as_slice()).collect();
            payloads.sort();
            assert_eq!(payloads, vec![&b"group 1"[..], &b"the first instance"[..]], "{via:?}");
        }
    })
    .await
    .expect("in time");
}

/// A publisher publishes a track, sends group 0 and closes its session with no PUBLISH_DONE
/// (a process that died, a run that ended); another session then publishes the same track
/// and sends group 0 again and group 1. The relay used to keep the closed session as the
/// owner of that subgroup and count the second publisher's objects as duplicates of it, so a
/// new subscriber was sent nothing at all.
///
/// The subscriber that held the first publication's object is told its subscription is over
/// (PUBLISH_DONE GOING_AWAY, the status of a session that closed); subscribing again it is
/// given the track's Largest Location and the second publisher's objects.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_track_published_again_after_its_publishers_session_closed_is_served() {
    timeout(LIMIT, async {
        for via in [Via::Quic, Via::WtDraft(Draft::D18)] {
            let server = relay(|_| {}).await;
            let space = ns(&["done", "closed"]);
            let key = (space.clone(), name("track"));
            let h = SubgroupHeader { track_alias: 1, group_id: 0, subgroup_id: SubgroupId::Zero, publisher_priority: Some(0), properties: false, end_of_group: true, first_object: true };
            let (subscriber, _e3) = connect(&server, via, ClientOptions::default()).await;
            let mut old_request = {
                let (first, _e) = connect(&server, via, ClientOptions::default()).await;
                let _p = first.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH");
                let (_ok, _stx, srx) = subscriber.subscribe(space.clone(), name("track"), Parameters::new()).await.expect("SUBSCRIBE");
                first.send_subgroup(h, &[SubgroupObject::normal(0, b"the closed session's".to_vec())]).await.expect("the first group 0");
                let got = collect(&subscriber, STEP, |g| !g.is_empty()).await;
                assert_eq!(got.iter().flat_map(|g| &g.objects).map(|o| o.payload.as_slice()).collect::<Vec<_>>(), vec![&b"the closed session's"[..]], "{via:?}");
                first.close(0, "gone");
                srx
            };
            match timeout(STEP, old_request.next()).await {
                Ok(Ok(Some(ControlMessage::PublishDone(d)))) => assert_eq!(d.status, PublishDoneCode::GOING_AWAY, "{via:?}: what the old subscription is told"),
                other => panic!("{via:?}: no PUBLISH_DONE for the subscription to the closed instance: {other:?}"),
            }
            let relay = server.relay().clone();
            assert!(until(STEP, || relay.track_ended(&key) == Some(PublishDoneCode::GOING_AWAY)).await, "{via:?}: the track ended with its publisher's session");
            let (second, _e2) = connect(&server, via, ClientOptions::default()).await;
            let _p = second.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH again");
            let (ok, _stx, _srx) = subscriber.subscribe(space.clone(), name("track"), Parameters::new()).await.expect("SUBSCRIBE again");
            second.send_subgroup(h, &[SubgroupObject::normal(0, b"the closed session's".to_vec())]).await.expect("group 0 again");
            second.send_subgroup(SubgroupHeader { group_id: 1, ..h }, &[SubgroupObject::normal(0, b"the second session's".to_vec())]).await.expect("group 1");
            let mine = |g: &Got| g.header.track_alias == ok.track_alias;
            let got = collect(&subscriber, STEP, |g| g.iter().filter(|g| mine(g)).count() >= 2).await;
            let mut payloads: Vec<&[u8]> = got.iter().filter(|g| mine(g)).flat_map(|g| &g.objects).map(|o| o.payload.as_slice()).collect();
            payloads.sort();
            assert_eq!(payloads, vec![&b"the closed session's"[..], &b"the second session's"[..]], "{via:?}: the republished track is served");
            assert_eq!(ok.params.largest_object(), Some(joule_moq::draft22::Location::new(0, 0)), "{via:?}: the track's Largest Location, kept");
        }
    })
    .await
    .expect("in time");
}

/// A publisher that sends PUBLISH_DONE naming a stream the relay has not had, and then
/// closes its session: its subscribers are told what it said (TRACK_ENDED), not
/// GOING_AWAY, and at the close, not after the wait for the stream. One that closes with
/// no PUBLISH_DONE ends its track with GOING_AWAY.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_close_after_publish_done_ends_with_what_the_publisher_said() {
    timeout(LIMIT, async {
        for via in [Via::Quic, Via::WtDraft(Draft::D18), Via::QuicDraft(Draft::D16)] {
            for said in [Some(PublishDoneCode::TRACK_ENDED), None] {
                let server = relay(|_| {}).await;
                let (publisher, _e1) = connect(&server, via, ClientOptions::default()).await;
                let (subscriber, _e2) = connect(&server, Via::Quic, ClientOptions::default()).await;
                let space = ns(&["done", "close"]);
                let (ptx, _prx) = publisher.publish(space.clone(), name("track"), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH");
                let (_ok, _stx, mut srx) = subscriber.subscribe(space, name("track"), Parameters::new()).await.expect("SUBSCRIBE");
                if let Some(status) = said {
                    ptx.send(&ControlMessage::PublishDone(PublishDone { status, stream_count: 1, reason: Vec::new() })).await.expect("PUBLISH_DONE");
                    // Delivered before the close: a close can drop what is still unsent.
                    tokio::time::sleep(Duration::from_millis(200)).await;
                }
                let start = Instant::now();
                publisher.close(0, "done");
                let want = said.unwrap_or(PublishDoneCode::GOING_AWAY);
                match timeout(STEP, srx.next()).await {
                    Ok(Ok(Some(ControlMessage::PublishDone(d)))) => assert_eq!(d.status, want, "{via:?}"),
                    other => panic!("{via:?}: no PUBLISH_DONE downstream: {other:?}"),
                }
                // At the close: well inside the wait the PUBLISH_DONE began.
                assert!(start.elapsed() < PUBLISH_DONE_STREAM_WAIT / 2, "{via:?}, {said:?}: ended at the close, in {:?}", start.elapsed());
            }
        }
    })
    .await
    .expect("in time");
}