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

//! A link the relay dials to its upstream (`crate::cascade`) is held to the same receive
//! windows as a session it accepts: `limits.quic_stream_window_bytes` for each stream and
//! `limits.quic_connection_window_bytes` for the connection. An upstream that opens data
//! streams before its SETUP is parked stream by stream (the relay reads each stream's header
//! and holds it until the SETUP arrives); it stalls once its unread bytes reach the
//! connection window, over native QUIC and over WebTransport alike. With quinn's client
//! defaults the connection's window is unbounded, and only the stream windows would stop it.

mod common;

use common::*;
use joule_moq::draft22::data::{SubgroupCodec, SubgroupHeader, SubgroupId, SubgroupObject};
use joule_moq_relay::cert::RelayIdentity;
use joule_moq_relay::config::{UpstreamConfig, UpstreamTransport};
use joule_moq_relay::RelayServer;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::timeout;

const STREAM_WINDOW: u64 = 1 << 20;
const CONNECTION_WINDOW: u64 = 2 << 20;
const STREAMS: u64 = 16;
/// The most the relay reads of a stream before it parks it: one read of the transport.
const READ_CHUNK: u64 = 64 << 10;
/// Each write's size. A write cut off by a stall may have gone in part, so it is not counted.
const WRITE: usize = 4 << 10;
/// A write that makes no progress for this long has stalled on flow control. Every stream
/// past the window costs one or two, and all of them fit in the relay's SETUP timeout (10 s).
const STALL: Duration = Duration::from_millis(200);

/// One data stream's bytes: a group's header, a small first object, then an object twice a
/// stream window long. Never finished.
fn stream_bytes(g: u64) -> Vec<u8> {
    let h = SubgroupHeader { track_alias: 7, group_id: g, subgroup_id: SubgroupId::Zero, publisher_priority: Some(128), properties: false, end_of_group: true, first_object: true };
    let mut b = Vec::new();
    h.encode(&mut b);
    let mut codec = SubgroupCodec::new(h);
    codec.encode_object(&SubgroupObject::normal(0, vec![1; 1000]), &mut b).unwrap();
    codec.encode_object(&SubgroupObject::normal(1, vec![2; 2 * STREAM_WINDOW as usize]), &mut b).unwrap();
    b
}

/// The relay under test, with one upstream at `url` pinned to `pin`, over `transport`.
async fn edge(url: String, transport: UpstreamTransport, pin: [u8; 32]) -> RelayServer {
    relay(|c| {
        c.limits.quic_stream_window_bytes = STREAM_WINDOW;
        c.limits.quic_connection_window_bytes = CONNECTION_WINDOW;
        c.limits.upstream_reconnect_ms = 60_000;
        c.upstream = vec![UpstreamConfig {
            name: "upstream".into(),
            url,
            transport,
            prefixes: vec![vec!["live".into()]],
            ca_path: None,
            cert_sha256: Some(hex::encode(pin)),
            native_roots: false,
            token_sha256: None,
            token_path: None,
            relay_key: None,
        }];
    })
    .await
}

/// What `sent` bytes over `STREAMS` streams say about the windows the relay's link had.
fn check(transport: &str, sent: u64, opened: u64) {
    eprintln!("{transport}: sent {sent} bytes on {opened} of {STREAMS} streams; the connection window {CONNECTION_WINDOW}, a stream window {STREAM_WINDOW}");
    // What the upstream could leave unread is the connection window, plus what the relay read
    // itself, one read a stream.
    let bound = CONNECTION_WINDOW + STREAMS * (READ_CHUNK + WRITE as u64);
    assert!(sent <= bound, "{transport}: the link stalled at its connection window: {sent} bytes sent, the bound {bound}");
    assert!(sent < STREAMS * STREAM_WINDOW / 2, "{transport}: the stream windows did not bound it: {sent} bytes sent");
    // It stalled at the window, not before.
    assert!(sent + STREAMS * WRITE as u64 >= CONNECTION_WINDOW, "{transport}: the upstream filled the connection window: {sent} bytes sent");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_quic_upstream_stalls_at_the_links_connection_window() {
    timeout(Duration::from_secs(120), async {
        let id = RelayIdentity::dev_self_signed(&["localhost".to_string(), "127.0.0.1".to_string()]).expect("an identity");
        let crypto = quinn::crypto::rustls::QuicServerConfig::try_from(id.server_tls().expect("TLS")).expect("QUIC TLS");
        let ep = quinn::Endpoint::server(quinn::ServerConfig::with_crypto(Arc::new(crypto)), loopback()).expect("an endpoint");
        let port = ep.local_addr().unwrap().port();
        let _edge = edge(format!("moqt://127.0.0.1:{port}/moq"), UpstreamTransport::Quic, id.leaf_sha256).await;
        let conn = timeout(STEP, async { ep.accept().await.expect("the link dials").await.expect("the link connects") }).await.expect("the link dials in time");
        let (mut sent, mut opened, mut streams) = (0u64, 0u64, Vec::new());
        for g in 0..STREAMS {
            let Ok(Ok(mut st)) = timeout(STEP, conn.open_uni()).await else { break };
            opened += 1;
            for c in stream_bytes(g).chunks(WRITE) {
                match timeout(STALL, st.write_all(c)).await {
                    Ok(Ok(())) => sent += c.len() as u64,
                    _ => break,
                }
            }
            streams.push(st);
        }
        check("quic", sent, opened);
        drop(streams);
        conn.close(0u32.into(), b"done");
    })
    .await
    .expect("the test finishes in time");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_webtransport_upstream_stalls_at_the_links_connection_window() {
    timeout(Duration::from_secs(120), async {
        let id = RelayIdentity::dev_self_signed(&["localhost".to_string(), "127.0.0.1".to_string()]).expect("an identity");
        let tls = id.server_tls().expect("TLS");
        let cfg = wtransport::ServerConfig::builder().with_bind_address(loopback()).with_custom_tls(tls).build();
        let server = wtransport::Endpoint::server(cfg).expect("an endpoint");
        let port = server.local_addr().unwrap().port();
        let _edge = edge(format!("https://127.0.0.1:{port}/moq"), UpstreamTransport::WebTransport, id.leaf_sha256).await;
        let conn = timeout(STEP, async {
            let request = server.accept().await.await.expect("the link dials");
            request.accept().await.expect("the link's session")
        })
        .await
        .expect("the link dials in time");
        let (mut sent, mut opened, mut streams) = (0u64, 0u64, Vec::new());
        for g in 0..STREAMS {
            // Opening a stream writes its WebTransport header, which flow control can hold too.
            let Ok(Ok(opening)) = timeout(STALL, conn.open_uni()).await else { continue };
            let Ok(Ok(mut st)) = timeout(STALL, opening).await else { continue };
            opened += 1;
            for c in stream_bytes(g).chunks(WRITE) {
                match timeout(STALL, st.write_all(c)).await {
                    Ok(Ok(())) => sent += c.len() as u64,
                    _ => break,
                }
            }
            streams.push(st);
        }
        check("webtransport", sent, opened);
        drop(streams);
    })
    .await
    .expect("the test finishes in time");
}