arcly-stream 0.7.2

//! **Network `ClusterRelay` over TCP** — cross-process / cross-machine federation.
//!
//! The shipped [`InProcessRelay`](arcly_stream::cluster) mirrors streams between
//! `Engine`s in *one* process. To actually scale viewers across machines you put
//! each edge on its own box and have it pull the origin's stream over the
//! network. This example is that missing piece: a `ClusterRelay` whose `pull`
//! opens a TCP connection to the origin, and an origin-side relay server that
//! streams a stream's frames to whoever connects.
//!
//! ```text
//!                         ┌─ EDGE box A ─ TcpRelay.pull ─▶ viewers
//!   ingest ─▶ ORIGIN box ─┤  (TCP)
//!   (relay server)        └─ EDGE box B ─ TcpRelay.pull ─▶ viewers
//! ```
//!
//! Run it (origin + one edge in this process, but over a *real* TCP socket — set
//! the origin address to another host to span machines):
//!
//! ```text
//! cargo run -p arcly-stream --example network_relay
//! ```
//!
//! **Zero-copy note.** Within a node, one mirrored frame is still fanned out to
//! every local viewer without copying (the headline property). The *wire* hop is
//! necessarily a copy — you cannot share a pointer across machines — so the edge
//! re-allocates each frame's payload once as it arrives, then shares that one
//! allocation with all of its own viewers.

use std::sync::Arc;
use std::time::Duration;

use arcly_stream::prelude::*;
use arcly_stream::{
    ClusterRelay, FrameFlags, FrameType, NodeAddr, PlaybackRegistry, PublishRegistry, StreamError,
    StreamHandle,
};
use bytes::Bytes;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};

// ── Wire format ─────────────────────────────────────────────────────────────
// Per frame: a fixed 37-byte header followed by the payload. All little-endian.
//   i64 pts | i64 dts | u8 has_dur | u64 dur | u8 codec | u8 frame_type
//   | u16 flags | u32 track_id | u32 data_len | [data_len bytes]
const HDR: usize = 37;

fn encode_frame(f: &MediaFrame) -> Vec<u8> {
    let mut b = Vec::with_capacity(HDR + f.data.len());
    b.extend_from_slice(&f.pts.to_le_bytes());
    b.extend_from_slice(&f.dts.to_le_bytes());
    match f.duration {
        Some(d) => {
            b.push(1);
            b.extend_from_slice(&d.to_le_bytes());
        }
        None => {
            b.push(0);
            b.extend_from_slice(&0u64.to_le_bytes());
        }
    }
    b.push(f.codec as u8);
    b.push(match f.frame_type {
        FrameType::Key => 0,
        FrameType::Delta => 1,
        FrameType::Audio => 2,
    });
    b.extend_from_slice(&f.flags.bits().to_le_bytes());
    b.extend_from_slice(&f.track_id.to_le_bytes());
    b.extend_from_slice(&(f.data.len() as u32).to_le_bytes());
    b.extend_from_slice(&f.data);
    b
}

/// Read one frame, or `None` on a clean end-of-stream.
async fn read_frame(r: &mut TcpStream) -> std::io::Result<Option<MediaFrame>> {
    let mut h = [0u8; HDR];
    if let Err(e) = r.read_exact(&mut h).await {
        if e.kind() == std::io::ErrorKind::UnexpectedEof {
            return Ok(None);
        }
        return Err(e);
    }
    let le8 = |i: usize| i64::from_le_bytes(h[i..i + 8].try_into().unwrap());
    let pts = le8(0);
    let dts = le8(8);
    let duration = (h[16] == 1).then(|| u64::from_le_bytes(h[17..25].try_into().unwrap()));
    let codec = CodecId::try_from(h[25]).unwrap_or(CodecId::Unknown);
    let frame_type = match h[26] {
        0 => FrameType::Key,
        1 => FrameType::Delta,
        _ => FrameType::Audio,
    };
    let flags = FrameFlags::from_bits_truncate(u16::from_le_bytes(h[27..29].try_into().unwrap()));
    let track_id = u32::from_le_bytes(h[29..33].try_into().unwrap());
    let len = u32::from_le_bytes(h[33..37].try_into().unwrap()) as usize;
    let mut data = vec![0u8; len];
    r.read_exact(&mut data).await?;
    Ok(Some(MediaFrame {
        pts,
        dts,
        duration,
        data: Bytes::from(data),
        codec,
        frame_type,
        flags,
        track_id,
    }))
}

async fn write_handshake(s: &mut TcpStream, key: &StreamKey) -> std::io::Result<()> {
    let (app, stream) = (
        key.app.as_str().as_bytes(),
        key.stream_id.as_str().as_bytes(),
    );
    s.write_all(&[app.len() as u8]).await?;
    s.write_all(app).await?;
    s.write_all(&[stream.len() as u8]).await?;
    s.write_all(stream).await
}

async fn read_handshake(s: &mut TcpStream) -> std::io::Result<StreamKey> {
    let mut one = [0u8; 1];
    s.read_exact(&mut one).await?;
    let mut app = vec![0u8; one[0] as usize];
    s.read_exact(&mut app).await?;
    s.read_exact(&mut one).await?;
    let mut stream = vec![0u8; one[0] as usize];
    s.read_exact(&mut stream).await?;
    Ok(StreamKey::new(
        String::from_utf8_lossy(&app).into_owned(),
        String::from_utf8_lossy(&stream).into_owned(),
    ))
}

// ── Origin: relay server ────────────────────────────────────────────────────
/// Accept connections; for each, read the requested key, replay the instant-start
/// buffer, then forward live frames until the viewer disconnects or the stream
/// ends. One connection = one edge mirroring one stream.
async fn serve_relay(listener: TcpListener, origin: Arc<Engine>) {
    loop {
        let Ok((mut sock, _)) = listener.accept().await else {
            break;
        };
        let origin = origin.clone();
        tokio::spawn(async move {
            let Ok(key) = read_handshake(&mut sock).await else {
                return;
            };
            let Ok(handle) = origin.get_stream(&key) else {
                return; // unknown stream — drop (a real server would signal 404)
            };
            // Instant-start: replay the cached GOP before going live.
            for fr in handle.replay_buffer() {
                if sock.write_all(&encode_frame(&fr)).await.is_err() {
                    return;
                }
            }
            let mut sub = handle.subscribe_resilient();
            while let Some(fr) = sub.recv().await {
                if sock.write_all(&encode_frame(&fr)).await.is_err() {
                    break; // edge went away
                }
            }
        });
    }
}

// ── Edge: the network ClusterRelay ──────────────────────────────────────────
/// A `ClusterRelay` that pulls streams from a single origin over TCP. A larger
/// deployment would resolve `locate` through a shared directory (etcd/DNS/…); here
/// the edge simply knows its origin's address.
struct TcpRelay {
    local: Arc<Engine>,
    origin_addr: String,
}

#[async_trait::async_trait]
impl ClusterRelay for TcpRelay {
    async fn locate(&self, _key: &StreamKey) -> arcly_stream::Result<Option<NodeAddr>> {
        Ok(Some(NodeAddr(self.origin_addr.clone())))
    }

    async fn pull(&self, key: &StreamKey, origin: &NodeAddr) -> arcly_stream::Result<()> {
        let mut sock = TcpStream::connect(&origin.0)
            .await
            .map_err(|e| StreamError::cluster(format!("connect {}: {e}", origin.0)))?;
        write_handshake(&mut sock, key)
            .await
            .map_err(|e| StreamError::cluster(format!("handshake: {e}")))?;

        // Claim the local mirror, then forward frames off the socket into it.
        let dst = self.local.start_publish(key).await?;
        let local = self.local.clone();
        let key = key.clone();
        tokio::spawn(async move {
            while let Ok(Some(frame)) = read_frame(&mut sock).await {
                if dst.publish_frame(frame).is_err() {
                    break;
                }
            }
            let _ = local.end_publish(&key).await;
        });
        Ok(())
    }

    async fn announce(&self, _key: &StreamKey) -> arcly_stream::Result<()> {
        Ok(()) // the origin's relay server is the directory in this reference
    }

    async fn withdraw(&self, _key: &StreamKey) -> arcly_stream::Result<()> {
        Ok(())
    }
}

/// Publish a synthetic stream on the origin: a CONFIG frame then ~25 fps of
/// keyframes/deltas. The relay forwards opaque bytes, so these need not decode.
fn publish_synthetic(handle: StreamHandle) {
    tokio::spawn(async move {
        let mut cfg =
            MediaFrame::new_video(0, 0, Bytes::from_static(b"\x00cfg"), CodecId::H264, true);
        cfg.flags |= FrameFlags::CONFIG;
        let _ = handle.publish_frame(cfg);

        let mut tick = tokio::time::interval(Duration::from_millis(40));
        let mut pts = 0i64;
        for n in 0..u64::MAX {
            tick.tick().await;
            let key = n % 25 == 0;
            let data = Bytes::from(vec![(n & 0xFF) as u8; 256]);
            let mut f = MediaFrame::new_video(pts, pts, data, CodecId::H264, key);
            f.duration = Some(40);
            if handle.publish_frame(f).is_err() {
                break;
            }
            pts += 40;
        }
    });
}

#[tokio::main]
async fn main() -> arcly_stream::Result<()> {
    let key = StreamKey::new("live", "cam");

    // ── ORIGIN: engine + relay server + a live stream ────────────────────────
    let origin = Engine::builder()
        .application(AppSpec::new("live").gop_cache(30))
        .build();
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let origin_addr = listener.local_addr().expect("addr").to_string();
    tokio::spawn(serve_relay(listener, origin.clone()));

    let pub_handle = origin.start_publish(&key).await?;
    publish_synthetic(pub_handle);
    println!("[origin] serving relay at {origin_addr}, publishing {key}");
    tokio::time::sleep(Duration::from_millis(300)).await; // let the GOP cache fill

    // ── EDGE: a different engine that mirrors over TCP ───────────────────────
    let edge = Engine::builder()
        .application(AppSpec::new("live").gop_cache(30))
        .build();
    let relay = TcpRelay {
        local: edge.clone(),
        origin_addr,
    };

    // Edge lacks the stream → locate the origin → pull it over the network.
    let node = relay.locate(&key).await?.expect("origin located");
    relay.pull(&key, &node).await?;
    println!("[edge]   pulling {key} from {} over TCP", node.0);

    // A viewer on the edge sees the mirrored stream (served locally, zero-copy).
    let edge_handle = loop {
        if let Ok(h) = edge.get_stream(&key) {
            break h;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    };
    let mut viewer = edge_handle.subscribe_resilient();
    let (mut frames, mut keyframes) = (0u32, 0u32);
    for _ in 0..60 {
        match tokio::time::timeout(Duration::from_secs(2), viewer.recv()).await {
            Ok(Some(fr)) => {
                frames += 1;
                if fr.is_keyframe() {
                    keyframes += 1;
                }
            }
            _ => break,
        }
    }

    println!("[edge]   viewer received {frames} mirrored frames ({keyframes} keyframes) over TCP");
    assert!(
        frames > 0 && keyframes > 0,
        "edge should mirror live frames"
    );
    println!("\nOK: network ClusterRelay mirrored {key} origin → edge across a TCP socket.");
    println!("    Point `origin_addr` at another host to federate across machines.");
    Ok(())
}