gamescope-hyperhdr/tests/hyperhdr_mock.rs

//! Network module against the in-process HyperHDR mock, without PipeWire (spec 9.2).

use std::io::Write;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use gamescope_hyperhdr::hyperhdr::client::{self, ClientConfig};
use gamescope_hyperhdr::hyperhdr::mock::{
    Behavior, Disconnect, Event, EventKind, ImageInfo, MockOptions, MockServer,
};
use gamescope_hyperhdr::pacing::PacingConfig;
use gamescope_hyperhdr::slot::{Frame, FrameSlot};
use gamescope_hyperhdr::stats::Stats;

const WAIT: Duration = Duration::from_secs(5);
const MS: Duration = Duration::from_millis(1);

// Client logs are captured globally; each client logs inside a `hyperhdr{server=...}` span.
#[derive(Clone)]
struct LogBuffer(Arc<Mutex<Vec<u8>>>);

impl Write for LogBuffer {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.0.lock().unwrap().extend_from_slice(buf);
        Ok(buf.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

fn log_buffer() -> &'static LogBuffer {
    static LOGS: OnceLock<LogBuffer> = OnceLock::new();
    LOGS.get_or_init(|| {
        let logs = LogBuffer(Arc::default());
        let writer = logs.clone();
        tracing_subscriber::fmt()
            .with_ansi(false)
            .with_max_level(tracing::Level::DEBUG)
            .with_writer(move || writer.clone())
            .init();
        logs
    })
}

fn log_lines(server: &MockServer, needle: &str) -> Vec<String> {
    let tag = format!("server={}}}", server.addr());
    let logs = log_buffer().0.lock().unwrap();
    String::from_utf8_lossy(&logs)
        .lines()
        .filter(|l| l.contains(&tag) && l.contains(needle))
        .map(str::to_owned)
        .collect()
}

fn wait_for_log(server: &MockServer, needle: &str) -> bool {
    let deadline = Instant::now() + WAIT;
    while Instant::now() < deadline {
        if !log_lines(server, needle).is_empty() {
            return true;
        }
        thread::sleep(10 * MS);
    }
    false
}

struct Harness {
    server: MockServer,
    opts: MockOptions,
    slot: Arc<FrameSlot>,
    stats: Arc<Stats>,
    client: Option<JoinHandle<()>>,
}

impl Harness {
    fn start(opts: MockOptions, configure: impl FnOnce(&mut ClientConfig)) -> Self {
        Self::start_with(opts, Behavior::Normal, configure)
    }

    fn start_with(
        opts: MockOptions,
        behavior: Behavior,
        configure: impl FnOnce(&mut ClientConfig),
    ) -> Self {
        log_buffer();
        let server = MockServer::start("127.0.0.1:0", opts.clone()).unwrap();
        server.set_behavior(behavior);
        let mut cfg = ClientConfig {
            pacing: PacingConfig {
                min_interval: 50 * MS,
                keepalive: 200 * MS,
                idle_clear: None,
            },
            connect_timeout: 500 * MS,
            write_timeout: 500 * MS,
            backoff_initial: 50 * MS,
            backoff_max: 400 * MS,
            stats_interval: None,
            ..ClientConfig::new("127.0.0.1", server.addr().port())
        };
        configure(&mut cfg);
        let slot = Arc::new(FrameSlot::new());
        let stats = Arc::new(Stats::default());
        let client = client::spawn(cfg, slot.clone(), stats.clone()).unwrap();
        Self {
            server,
            opts,
            slot,
            stats,
            client: Some(client),
        }
    }

    fn publish(&self, width: u32, height: u32, tag: u8) {
        let mut frame = Frame {
            width,
            height,
            rgb: vec![tag; (width * height * 3) as usize],
        };
        self.slot.publish(&mut frame);
    }

    fn wait_for(&self, pred: impl Fn(&[Event]) -> bool) -> bool {
        self.server.wait_for(WAIT, pred)
    }

    fn shutdown(&mut self) -> Duration {
        let started = Instant::now();
        self.slot.shutdown();
        if let Some(client) = self.client.take() {
            client.join().unwrap();
        }
        started.elapsed()
    }
}

impl Drop for Harness {
    fn drop(&mut self) {
        self.slot.shutdown();
        if let Some(client) = self.client.take() {
            let _ = client.join();
        }
    }
}

fn images(events: &[Event]) -> Vec<(Instant, u64, ImageInfo)> {
    events
        .iter()
        .filter_map(|e| match &e.kind {
            EventKind::Image(info) => Some((e.at, e.conn, info.clone())),
            _ => None,
        })
        .collect()
}

fn count(events: &[Event], pred: impl Fn(&EventKind) -> bool) -> usize {
    events.iter().filter(|e| pred(&e.kind)).count()
}

fn is_connected(k: &EventKind) -> bool {
    matches!(k, EventKind::Connected { .. })
}

fn has_image(events: &[Event]) -> bool {
    count(events, |k| matches!(k, EventKind::Image(_))) > 0
}

/// Scenario 1: lazy connect, `Register` first with the configured origin and priority.
#[test]
fn registers_before_sending_images() {
    let h = Harness::start(MockOptions::default(), |cfg| {
        cfg.origin = "test-origin".into();
        cfg.priority = 123;
    });
    thread::sleep(300 * MS);
    assert!(
        h.server.events().is_empty(),
        "must not connect before the first frame"
    );

    h.publish(32, 18, 1);
    assert!(h.wait_for(has_image));
    let events = h.server.events();
    assert!(is_connected(&events[0].kind));
    assert_eq!(
        events[1].kind,
        EventKind::Register {
            origin: "test-origin".into(),
            priority: 123
        }
    );
    let EventKind::Image(image) = &events[2].kind else {
        panic!("expected an image, got {:?}", events[2].kind);
    };
    assert_eq!((image.width, image.height, image.duration), (32, 18, -1));
    assert!(image.size_ok());
    assert_eq!(image.first_pixel, [1, 1, 1]);
    assert!(wait_for_log(&h.server, "registered with priority 123"));
}

/// Scenario 2: correct sizes, rate below max-fps, the newest frame wins in a burst.
#[test]
fn paces_frames_and_sends_the_newest() {
    let h = Harness::start(MockOptions::default(), |cfg| {
        cfg.pacing.min_interval = 100 * MS;
        cfg.pacing.keepalive = Duration::from_secs(4);
    });
    h.publish(64, 36, 1);
    assert!(h.wait_for(has_image));

    for tag in 2..=60 {
        h.publish(64, 36, tag);
        thread::sleep(2 * MS);
    }
    assert!(h.wait_for(|ev| images(ev)
        .last()
        .is_some_and(|(_, _, i)| i.first_pixel == [60; 3])));
    thread::sleep(300 * MS);

    let sent = images(&h.server.events());
    assert!(sent
        .iter()
        .all(|(_, _, i)| i.size_ok() && (i.width, i.height) == (64, 36)));
    for pair in sent.windows(2) {
        let gap = pair[1].0 - pair[0].0;
        assert!(gap >= 80 * MS, "images {gap:?} apart, max-fps is 10");
    }
    // 59 frames in ~130 ms at 10 fps: only a handful go out.
    assert!(sent.len() <= 5, "{} images sent", sent.len());
    assert_eq!(sent.last().unwrap().2.first_pixel, [60; 3]);
    assert!(h.stats.superseded.load(Ordering::Relaxed) > 0);
}

/// Scenario 3: keepalives keep the connection open longer than the server timeout.
#[test]
fn keepalive_outlives_server_timeout() {
    let opts = MockOptions {
        idle_timeout: 500 * MS,
        ..MockOptions::default()
    };
    let h = Harness::start(opts, |cfg| cfg.pacing.keepalive = 150 * MS);
    h.publish(8, 8, 9);
    thread::sleep(2000 * MS);

    let events = h.server.events();
    assert_eq!(count(&events, is_connected), 1);
    assert_eq!(
        count(&events, |k| matches!(k, EventKind::Disconnected(_))),
        0
    );
    assert!(
        images(&events).len() >= 8,
        "{} images",
        images(&events).len()
    );
    assert!(h.stats.keepalives.load(Ordering::Relaxed) >= 7);
}

/// The mock's timeout itself works (the premise of scenario 3).
#[test]
fn mock_drops_silent_clients() {
    let opts = MockOptions {
        idle_timeout: 300 * MS,
        ..MockOptions::default()
    };
    let h = Harness::start(opts, |cfg| cfg.pacing.keepalive = Duration::from_secs(3));
    h.publish(8, 8, 1);
    assert!(h.wait_for(|ev| count(ev, |k| *k == EventKind::Disconnected(Disconnect::Timeout)) > 0));
}

/// Scenario 4: server restart → reconnect with backoff and a new `Register`.
#[test]
fn reconnects_after_server_restart() {
    let mut h = Harness::start(MockOptions::default(), |_| {});
    h.publish(8, 8, 1);
    assert!(h.wait_for(has_image));

    let addr = h.server.addr();
    h.server.stop();
    thread::sleep(300 * MS);
    h.server = MockServer::start(addr, h.opts.clone()).unwrap();

    assert!(h.wait_for(|ev| {
        count(ev, |k| matches!(k, EventKind::Register { .. })) == 1 && has_image(ev)
    }));
    let events = h.server.events();
    assert!(matches!(
        events[1].kind,
        EventKind::Register { priority: 150, .. }
    ));
    assert!(h.stats.reconnects.load(Ordering::Relaxed) >= 1);
    assert!(wait_for_log(&h.server, "disconnected after"));
}

/// Scenario 5: immediate close (not whitelisted) → one hint in the log, growing backoff.
#[test]
fn rejected_connection_hints_and_backs_off() {
    let h = Harness::start_with(MockOptions::default(), Behavior::Reject, |_| {});
    h.publish(8, 8, 1);
    thread::sleep(2500 * MS);

    let rejected: Vec<Instant> = h
        .server
        .events()
        .iter()
        .filter(|e| matches!(e.kind, EventKind::Rejected { .. }))
        .map(|e| e.at)
        .collect();
    assert!(rejected.len() >= 4, "{} attempts", rejected.len());
    assert!(
        rejected.len() <= 12,
        "{} attempts, backoff not applied",
        rejected.len()
    );
    let gaps: Vec<Duration> = rejected.windows(2).map(|w| w[1] - w[0]).collect();
    assert!(gaps[0] < 250 * MS, "first gap {:?}", gaps[0]);
    assert!(
        *gaps.last().unwrap() >= 350 * MS,
        "last gap {:?}",
        gaps.last()
    );

    let hints = log_lines(&h.server, "/24 subnet");
    assert_eq!(hints.len(), 1, "{hints:#?}");
}

/// Scenario 6: a server that does not read → bounded memory, write timeout, reconnect.
#[test]
fn unread_socket_triggers_write_timeout() {
    let mut h = Harness::start_with(MockOptions::default(), Behavior::NoRead, |cfg| {
        cfg.pacing.min_interval = 5 * MS;
        cfg.write_timeout = 300 * MS;
    });
    let rss_before = rss_kib();

    let stop = Arc::new(AtomicBool::new(false));
    let publisher = {
        let (slot, stop) = (h.slot.clone(), stop.clone());
        thread::spawn(move || {
            let mut tag = 0u8;
            while !stop.load(Ordering::Relaxed) {
                tag = tag.wrapping_add(1);
                let mut frame = Frame {
                    width: 640,
                    height: 360,
                    rgb: vec![tag; 640 * 360 * 3],
                };
                slot.publish(&mut frame);
                thread::sleep(5 * MS);
            }
        })
    };

    assert!(h.wait_for(|ev| count(ev, is_connected) >= 1));
    h.server.set_behavior(Behavior::Normal);
    assert!(h.wait_for(|ev| count(ev, is_connected) >= 2 && has_image(ev)));
    assert!(wait_for_log(&h.server, "write timed out"));

    stop.store(true, Ordering::Relaxed);
    publisher.join().unwrap();
    h.shutdown();
    let growth = rss_kib().saturating_sub(rss_before);
    assert!(growth < 64 * 1024, "RSS grew by {growth} KiB");
}

fn rss_kib() -> u64 {
    std::fs::read_to_string("/proc/self/status")
        .ok()
        .and_then(|s| {
            s.lines()
                .find(|l| l.starts_with("VmRSS:"))
                .and_then(|l| l.split_whitespace().nth(1)?.parse().ok())
        })
        .unwrap_or(0)
}

/// Scenario 7: source gone → `Clear` → FIN; the next frame opens a new connection.
#[test]
fn source_gone_clears_and_disconnects() {
    let h = Harness::start(MockOptions::default(), |_| {});
    h.publish(8, 8, 1);
    assert!(h.wait_for(has_image));

    h.slot.source_gone();
    assert!(h.wait_for(|ev| count(ev, |k| *k
        == EventKind::Disconnected(Disconnect::ClosedByClient))
        == 1));
    let events = h.server.events();
    let clear = events
        .iter()
        .position(|e| e.kind == EventKind::Clear { priority: 150 });
    let fin = events
        .iter()
        .position(|e| e.kind == EventKind::Disconnected(Disconnect::ClosedByClient));
    assert!(clear.is_some() && clear < fin, "{events:#?}");

    thread::sleep(300 * MS);
    assert_eq!(
        count(&h.server.events(), is_connected),
        1,
        "no frames, no connection"
    );

    h.publish(8, 8, 2);
    assert!(h.wait_for(|ev| {
        count(ev, is_connected) == 2
            && images(ev)
                .last()
                .is_some_and(|(_, conn, i)| *conn != ev[0].conn && i.first_pixel == [2; 3])
    }));
}

/// Scenario 8: shutdown → `Clear` → FIN within a second.
#[test]
fn shutdown_clears_quickly() {
    let mut h = Harness::start(MockOptions::default(), |_| {});
    h.publish(8, 8, 1);
    assert!(h.wait_for(has_image));

    let took = h.shutdown();
    assert!(took < Duration::from_secs(1), "shutdown took {took:?}");
    assert!(h.server.wait_for(Duration::from_secs(1), |ev| {
        count(ev, |k| *k == EventKind::Clear { priority: 150 }) == 1
            && count(ev, |k| {
                *k == EventKind::Disconnected(Disconnect::ClosedByClient)
            }) == 1
    }));
}

/// Scenario 9: error replies are logged with a rate limit.
#[test]
fn error_replies_are_rate_limited() {
    let h = Harness::start_with(MockOptions::default(), Behavior::ErrorReplies, |cfg| {
        cfg.pacing.min_interval = 10 * MS;
        cfg.pacing.keepalive = 20 * MS;
    });
    h.publish(8, 8, 1);
    thread::sleep(1500 * MS);

    assert!(h.stats.error_replies.load(Ordering::Relaxed) >= 20);
    let logged = log_lines(&h.server, "replied with an error");
    assert_eq!(logged.len(), 1, "{logged:#?}");
}

/// Idle clear: `Clear` once, connection kept with heartbeats, resumed by a new frame.
#[test]
fn idle_clear_and_resume() {
    let opts = MockOptions {
        idle_timeout: Duration::from_secs(1),
        ..MockOptions::default()
    };
    let h = Harness::start(opts, |cfg| {
        cfg.pacing.keepalive = 100 * MS;
        cfg.pacing.idle_clear = Some(300 * MS);
    });
    h.publish(8, 8, 1);
    thread::sleep(1600 * MS);

    let events = h.server.events();
    let clear_at = events
        .iter()
        .position(|e| e.kind == EventKind::Clear { priority: 150 })
        .expect("Clear after idle");
    assert_eq!(count(&events, |k| matches!(k, EventKind::Clear { .. })), 1);
    assert_eq!(
        count(&events[clear_at..], |k| matches!(k, EventKind::Image(_))),
        0
    );
    assert!(
        count(&events[clear_at..], |k| matches!(
            k,
            EventKind::Register { .. }
        )) >= 5
    );
    assert_eq!(
        count(&events, |k| matches!(k, EventKind::Disconnected(_))),
        0
    );

    h.publish(8, 8, 2);
    assert!(h.wait_for(|ev| images(ev)
        .last()
        .is_some_and(|(_, _, i)| i.first_pixel == [2; 3])));
    assert_eq!(count(&h.server.events(), is_connected), 1);
}