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

//! A late joiner served from the relay's cache (`draft-ietf-moq-transport-22` §3.2, §3.4,
//! §7.1), and a FETCH the relay forwards upstream when its cache holds nothing (§3.2.3).

mod common;

use common::*;
use joule_moq::carriage::ReceiptBundle;
use joule_moq::draft22::data::{FetchCodec, FetchEntry, FetchHeader, FetchObject};
use joule_moq::draft22::{
    ControlMessage, FetchOk, GroupOrder, Location, LocationFilter, Parameter, Parameters, RequestErrorCode, StreamHeader, TrackRequest, UniEvent,
};
use joule_moq::receipt::ReceiptVerifier;
use joule_moq_relay::client::{Client, ClientOptions};
use std::time::Duration;
use tokio::time::timeout;

const GROUPS: u64 = 10;

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

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

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

fn objects(entries: &[FetchEntry]) -> Vec<&FetchObject> {
    entries
        .iter()
        .filter_map(|e| match e {
            FetchEntry::Object(o) => Some(o),
            _ => None,
        })
        .collect()
}

async fn from_cache() {
    let server = relay(|_| {}).await;
    let key = ed25519_dalek::SigningKey::from_bytes(&[0x1a; 32]);
    let t = Tracks::new(ns(&["vod", "show"]), "ep1", 10, 3, 128);
    let (publisher, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let _held = publish_tracks(&publisher, &t, None).await;
    for g in 0..GROUPS {
        publish_group(&publisher, &key, &t, g, &payload("ep1", g, 2000)).await;
    }
    // The cache, not the count of objects received, says when it holds the track: an
    // object counted on arrival may not be cached yet, and streams finish out of order.
    let relay = server.relay().clone();
    let (mk, rk) = ((t.ns.clone(), t.media.clone()), (t.ns.clone(), t.receipts.clone()));
    assert!(
        until(STEP, || relay.largest(&mk) == Some(Location::new(GROUPS - 1, 0)) && relay.cache_covers(&mk, Location::new(0, 0)) && relay.cache_covers(&rk, Location::new(0, 0))).await,
        "the cache holds the track and its receipts from group 0"
    );

    // The late joiner comes over the WebSocket fallback.
    let (late, _) = connect(&server, Via::WebSocket, ClientOptions::default()).await;

    // TRACK_STATUS names the Largest Object (§9.13).
    let m = ControlMessage::TrackStatus(TrackRequest { request_id: late.next_request_id(), namespace: t.ns.clone(), name: t.media.clone(), params: Parameters::new() });
    let (_tx, mut rx) = late.request(&m).await.unwrap();
    match Client::response(&mut rx).await.unwrap() {
        ControlMessage::RequestOk(ok) => assert_eq!(ok.params.largest_object(), Some(Location::new(GROUPS - 1, 0))),
        other => panic!("{other:?}"),
    }

    // The whole track and its receipts, by FETCH, from the cache.
    let (ok, media) = late.fetch(t.ns.clone(), t.media.clone(), Parameters::new()).await.expect("FETCH the media track");
    assert_eq!(ok.end_location, Location::new(GROUPS - 1, 0));
    let (_, rec) = late.fetch(t.ns.clone(), t.receipts.clone(), Parameters::new()).await.expect("FETCH the receipts track");
    let media = objects(&media);
    let rec = objects(&rec);
    assert_eq!((media.len(), rec.len()), (GROUPS as usize, GROUPS as usize));
    let verifier = ReceiptVerifier::with_anchors(vec![key.verifying_key()]);
    for (m, r) in media.iter().zip(&rec) {
        assert_eq!((m.group_id, r.group_id), (r.group_id, m.group_id));
        let b = ReceiptBundle::from_payload(&r.payload).unwrap();
        let vg = verifier.verify_group(&b.group).expect("a receipt from the cache verifies");
        verifier.verify_payload(&m.payload, &b.objects[0], &vg).expect("an object from the cache verifies");
        assert_eq!(m.payload, payload("ep1", m.group_id, 2000));
        assert_eq!(m.publisher_priority, 128);
    }
    assert_eq!(media.iter().map(|o| o.group_id).collect::<Vec<_>>(), (0..GROUPS).collect::<Vec<_>>(), "ascending by default (§9.20.8)");

    // A range, groups descending (§3.2.1).
    let range = Parameters::new()
        .with(Parameter::LocationFilter(LocationFilter::AbsoluteRange { start: Location::new(3, 0), end_group_delta: 4, end_object: 0 }))
        .with(Parameter::GroupOrder(GroupOrder::Descending));
    let (ok, part) = late.fetch(t.ns.clone(), t.media.clone(), range).await.expect("FETCH a range");
    assert_eq!(ok.end_location, Location::new(7, 0));
    assert_eq!(objects(&part).iter().map(|o| o.group_id).collect::<Vec<_>>(), vec![7, 6, 5, 4, 3]);

    // A fetch that starts after the Largest Object is refused (§3.2: INVALID_RANGE).
    let after = Parameters::new().with(Parameter::LocationFilter(LocationFilter::AbsoluteStart { start: Location::new(GROUPS + 5, 0) }));
    let refused = late.fetch(t.ns.clone(), t.media.clone(), after).await.expect_err("a range past the Largest Object");
    assert!(refused.contains(&RequestErrorCode::INVALID_RANGE.to_string()), "{refused}");

    // Joining at a past group: SUBSCRIBE from the Next Object with a fill of the last three
    // groups (§3.5), delivered on a fill fetch stream named by the SUBSCRIBE's Request ID.
    let fill = Parameters::new().with(Parameter::LocationFilter(LocationFilter::RelativeStart { start_group: 3 }));
    let id_before = late.next_request_id();
    let params = Parameters::new().with(Parameter::LocationFilter(LocationFilter::NextObject)).with(Parameter::FillParameters(fill));
    let (ok, _stx, _srx) = late.subscribe(t.ns.clone(), t.media.clone(), params).await.expect("SUBSCRIBE with a fill");
    assert_eq!(ok.params.largest_object(), Some(Location::new(GROUPS - 1, 0)));
    let subscribe_id = id_before + 2;
    let mut r = late.next_data(STEP).await.expect("the fill fetch stream");
    match r.header.clone() {
        Some(StreamHeader::Fetch(FetchHeader { request_id })) => assert_eq!(request_id, subscribe_id),
        other => panic!("expected a fill fetch stream, got {other:?}"),
    }
    let mut filled = Vec::new();
    while let Some(UniEvent::FetchEntry(e)) = timeout(STEP, r.next()).await.expect("in time").expect("a fill entry") {
        filled.push(e);
    }
    assert_eq!(objects(&filled).iter().map(|o| o.group_id).collect::<Vec<_>>(), vec![GROUPS - 3, GROUPS - 2, GROUPS - 1]);

    // All of it from the cache: the publisher was asked for nothing.
    assert!(relay.stats().fetch_served_from_cache >= 4);
    assert_eq!(relay.stats().fetch_served_upstream, 0);
    assert!(publisher.next_request(Duration::from_millis(300)).await.is_none(), "the relay sent the publisher no request");
    late.close(0, "done");
    publisher.close(0, "done");
    server.shutdown().await;
}

/// Group 1 reaches the relay whole before group 0 does, as streams finish out of order on a
/// loaded host. The cache still serves a FETCH from group 0: coverage is the run of whole
/// groups, not wherever the first object kept happened to fall.
async fn reordered() {
    let server = relay(|_| {}).await;
    let key = ed25519_dalek::SigningKey::from_bytes(&[0x2b; 32]);
    let t = Tracks::new(ns(&["vod", "reorder"]), "ep", 10, 7, 128);
    let (publisher, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let _held = publish_tracks(&publisher, &t, None).await;
    let relay = server.relay().clone();
    let (mk, rk) = ((t.ns.clone(), t.media.clone()), (t.ns.clone(), t.receipts.clone()));

    // Group 1 first: its receipts stream and its object's stream, each to its FIN.
    publish_group(&publisher, &key, &t, 1, &payload("ep", 1, 1500)).await;
    assert!(until(STEP, || relay.largest(&mk) == Some(Location::new(1, 0))).await, "the relay has group 1");
    // Then group 0, then the rest in order.
    for g in [0, 2, 3, 4, 5] {
        publish_group(&publisher, &key, &t, g, &payload("ep", g, 1500)).await;
    }
    assert!(
        until(STEP, || relay.largest(&mk) == Some(Location::new(5, 0)) && relay.cache_covers(&mk, Location::new(0, 0)) && relay.cache_covers(&rk, Location::new(0, 0))).await,
        "the cache covers the track from group 0, though group 0 came after group 1"
    );

    let (late, _) = connect(&server, Via::WebSocket, ClientOptions::default()).await;
    let (ok, media) = late.fetch(t.ns.clone(), t.media.clone(), Parameters::new()).await.expect("FETCH from group 0");
    assert_eq!(ok.end_location, Location::new(5, 0));
    let (_, rec) = late.fetch(t.ns.clone(), t.receipts.clone(), Parameters::new()).await.expect("FETCH its receipts");
    let (media, rec) = (objects(&media), objects(&rec));
    assert_eq!(media.iter().map(|o| o.group_id).collect::<Vec<_>>(), (0..=5).collect::<Vec<_>>());
    let verifier = ReceiptVerifier::with_anchors(vec![key.verifying_key()]);
    for (m, r) in media.iter().zip(&rec) {
        let b = ReceiptBundle::from_payload(&r.payload).unwrap();
        let vg = verifier.verify_group(&b.group).unwrap();
        verifier.verify_payload(&m.payload, &b.objects[0], &vg).expect("each object verifies");
    }
    assert_eq!(relay.stats().fetch_served_upstream, 0, "nothing went upstream");
    assert!(relay.stats().fetch_served_from_cache >= 2);
    assert!(publisher.next_request(Duration::from_millis(300)).await.is_none(), "the relay sent the publisher no request");
    late.close(0, "done");
    publisher.close(0, "done");
    server.shutdown().await;
}

async fn from_upstream() {
    let server = relay(|_| {}).await;
    let space = ns(&["archive"]);
    let (publisher, _ep) = connect(&server, Via::Quic, ClientOptions::default()).await;
    let _ns = publisher.publish_namespace(space.clone(), Parameters::new()).await.unwrap();
    let (late, _el) = connect(&server, Via::Quic, ClientOptions::default()).await;

    let asking = late.fetch(space.clone(), name("old"), Parameters::new());
    let answering = async {
        let (tx, _rx, first) = publisher.next_request(STEP).await.expect("the relay's upstream FETCH");
        let ControlMessage::Fetch(r) = first else { panic!("expected FETCH, got {}", first.name()) };
        assert_eq!(r.name, name("old"));
        let end = Location::new(2, 0);
        tx.send(&ControlMessage::FetchOk(FetchOk { end_of_track: true, end_location: end, params: Parameters::new(), track_properties: Default::default() })).await.unwrap();
        let mut b = Vec::new();
        FetchHeader { request_id: r.request_id }.encode(&mut b);
        let mut codec = FetchCodec::new(r.params.group_order().unwrap_or(GroupOrder::Ascending));
        for g in 0..=2u64 {
            codec.encode(&FetchEntry::Object(FetchObject { group_id: g, object_id: 0, subgroup_id: Some(0), publisher_priority: 7, properties: None, payload: payload("old", g, 300) }), &mut b).unwrap();
        }
        let mut s = publisher.conn.open_uni().await.unwrap();
        s.write_all(&b).await.unwrap();
        s.finish().unwrap();
        tx
    };
    let (got, _tx) = tokio::join!(asking, answering);
    let (ok, entries) = got.expect("the FETCH answered through the relay");
    assert!(ok.end_of_track);
    let objs = objects(&entries);
    assert_eq!(objs.iter().map(|o| (o.group_id, o.payload.clone())).collect::<Vec<_>>(), (0..=2u64).map(|g| (g, payload("old", g, 300))).collect::<Vec<_>>());
    assert_eq!(server.relay().stats().fetch_served_upstream, 1);
    late.close(0, "done");
    publisher.close(0, "done");
    server.shutdown().await;
}