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

//! Two publishers of one track, and the one a subgroup is taken from goes while the other
//! stays (`draft-ietf-moq-transport-23` §7.3, §2.2).
//!
//! The relay forwards a subgroup from the first publisher's stream and counts the second's
//! copy of it as duplicates. When the first one's stream ends without a FIN (its session
//! closes, or it resets the stream) the subscriber's stream of that subgroup is reset, and the
//! track does not end, since a publisher remains. What the remaining publisher then sends of
//! the subgroup, past the last object forwarded, reaches the subscriber on a new stream of the
//! same Group and Subgroup, which ends with that publisher's FIN; what it sends of the
//! objects already forwarded (it was behind the publisher that went) is not sent again. Before, the subgroup stayed the departed
//! stream's and the subscriber was sent nothing more of it.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupCodec, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq::draft22::{Location, Parameters, StreamResetCode, UniEvent};
use joule_moq_relay::client::{Client, ClientOptions};
use joule_moq_relay::transport::{BoxSend, TransportError};
use joule_moq_relay::wire::{StreamError, UniReader};
use std::time::Duration;
use tokio::time::timeout;

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

/// How the publisher the subgroup is taken from goes.
#[derive(Clone, Copy, Debug)]
enum Goes {
    /// Its session closes, with no PUBLISH_DONE.
    SessionCloses,
    /// It resets the subgroup's stream with this code, and stays.
    ResetsItsStream(u64),
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_remaining_publisher_takes_a_subgroup_up_when_its_owners_session_closes() {
    timeout(LIMIT, run(Goes::SessionCloses)).await.expect("in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_remaining_publisher_takes_a_subgroup_up_when_its_owner_resets_the_stream() {
    timeout(LIMIT, run(Goes::ResetsItsStream(StreamResetCode::INTERNAL_ERROR.0))).await.expect("in time");
}

fn body(o: u64) -> Vec<u8> {
    format!("object {o} of the subgroup").into_bytes()
}

/// A stream of group 0, subgroup 0 of the publisher's track (alias 1), left open.
async fn open(c: &Client) -> (BoxSend, SubgroupCodec) {
    let h = SubgroupHeader { track_alias: 1, group_id: 0, subgroup_id: SubgroupId::Zero, publisher_priority: Some(0), properties: false, end_of_group: false, first_object: true };
    let mut b = Vec::new();
    h.encode(&mut b);
    let mut st = timeout(STEP, c.conn.open_uni()).await.expect("a stream in time").expect("a stream");
    timeout(STEP, st.write_all(&b)).await.expect("in time").expect("written");
    (st, SubgroupCodec::new(h))
}

async fn send(st: &mut (BoxSend, SubgroupCodec), o: u64) {
    let mut b = Vec::new();
    st.1.encode_object(&SubgroupObject::normal(o, body(o)), &mut b).unwrap();
    timeout(STEP, st.0.write_all(&b)).await.expect("in time").expect("written");
}

/// The next object of a subscriber's stream, or how the stream ended.
async fn next(r: &mut UniReader) -> Result<Option<(u64, Vec<u8>)>, u64> {
    loop {
        match timeout(STEP, r.next()).await.expect("a read in time") {
            Ok(Some(UniEvent::SubgroupObject(o))) => return Ok(Some((o.object_id, o.payload))),
            Ok(Some(_)) => continue,
            Ok(None) => return Ok(None),
            Err(StreamError::Transport(TransportError::Reset(code))) => return Err(code),
            Err(e) => panic!("a data stream: {e}"),
        }
    }
}

async fn run(goes: Goes) {
    let server = relay(|_| {}).await;
    let relay = server.relay().clone();
    let (space, track) = (ns(&["handover"]), name("video"));
    let key = (space.clone(), track.clone());
    let (first, _e1) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let (second, _e2) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let _p1 = first.publish(space.clone(), track.clone(), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH");
    let _p2 = second.publish(space.clone(), track.clone(), 1, Parameters::new(), track_props(0, None)).await.expect("PUBLISH by a second publisher");
    let (subscriber, _e3) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let (ok, _stx, mut srx) = subscriber.subscribe(space.clone(), track.clone(), Parameters::new()).await.expect("SUBSCRIBE");

    // The first publisher's stream is the one the subgroup is taken from: objects 0, 1 and 2.
    let mut a = open(&first).await;
    send(&mut a, 0).await;
    assert!(until(STEP, || relay.largest(&key) == Some(Location::new(0, 0))).await, "the first publisher's object is in");
    send(&mut a, 1).await;
    send(&mut a, 2).await;
    let mut r1 = subscriber.next_data(STEP).await.expect("the subgroup's stream");
    for o in 0..3 {
        assert_eq!(next(&mut r1).await, Ok(Some((o, body(o)))));
    }
    // The second publisher is behind it. Its copy of object 0: a duplicate, forwarded to nobody.
    let mut b = open(&second).await;
    send(&mut b, 0).await;
    assert!(subscriber.next_data(Duration::from_millis(400)).await.is_none(), "a duplicate subgroup opened no stream");

    // The first publisher goes; the subscriber's stream is reset, and its subscription stands.
    let want = match goes {
        Goes::SessionCloses => {
            first.close(0, "gone");
            StreamResetCode::SESSION_CLOSED.0
        }
        Goes::ResetsItsStream(code) => {
            a.0.reset(code);
            code
        }
    };
    assert_eq!(next(&mut r1).await, Err(want), "{goes:?}: the stream taken from the publisher that went is reset");
    assert!(timeout(Duration::from_millis(400), srx.next()).await.is_err(), "{goes:?}: a publisher remains, so the track has not ended");

    // What the remaining publisher sends of the subgroup next. Objects 1 and 2 were forwarded
    // already and are not sent again; objects 3 and 4 arrive on a new stream of the same
    // subgroup, which ends with its FIN.
    send(&mut b, 1).await;
    send(&mut b, 2).await;
    assert!(subscriber.next_data(Duration::from_millis(400)).await.is_none(), "{goes:?}: an object already forwarded was sent again");
    send(&mut b, 3).await;
    let mut r2 = subscriber.next_data(STEP).await.unwrap_or_else(|| panic!("{goes:?}: nothing more of the subgroup reached the subscriber once its first publisher went"));
    let joule_moq::draft22::StreamHeader::Subgroup(h2) = r2.header.clone().expect("its header") else { panic!("not a subgroup stream") };
    assert_eq!((h2.track_alias, h2.group_id, h2.resolve_subgroup_id(3), h2.first_object), (ok.track_alias, 0, 0, false), "{goes:?}: the same subgroup, on a new stream that does not claim the subgroup's first object");
    assert_eq!(next(&mut r2).await, Ok(Some((3, body(3)))), "{goes:?}: from the object after the last one forwarded");
    send(&mut b, 4).await;
    assert_eq!(next(&mut r2).await, Ok(Some((4, body(4)))));
    b.0.finish().expect("a FIN");
    assert_eq!(next(&mut r2).await, Ok(None), "{goes:?}: the remaining publisher's FIN ends it");

    // A later group is the remaining publisher's from its start.
    let h = SubgroupHeader { track_alias: 1, group_id: 1, subgroup_id: SubgroupId::Zero, publisher_priority: Some(0), properties: false, end_of_group: true, first_object: true };
    second.send_subgroup(h, &[SubgroupObject::normal(0, body(9))]).await.expect("group 1");
    let got = collect(&subscriber, STEP, |g| !g.is_empty()).await;
    assert_eq!((got[0].header.group_id, got[0].objects.len(), got[0].reset), (1, 1, None));
    assert!(timeout(Duration::from_millis(200), srx.next()).await.is_err(), "{goes:?}: the subscription stands");

    // The second publisher's summary: its copies of objects 0, 1 and 2 counted duplicates.
    second.close(0, "done");
    assert!(until(STEP, || relay.summaries().iter().any(|s| s.summary.tally.publisher.deduplicated > 0)).await);
    let s = relay.summaries().into_iter().find(|s| s.summary.tally.publisher.deduplicated > 0).unwrap();
    assert_eq!((s.summary.tally.publisher.deduplicated, s.summary.tally.publisher.objects_received), (3, 6), "{goes:?}");
    for c in [&first, &subscriber] {
        c.close(0, "done");
    }
    server.shutdown().await;
}