arcly-stream 0.7.2

//! Multi-node clustering contracts: origin/edge discovery and stream relay.
//!
//! These are **contracts only** — the engine defines the seams an edge tier
//! plugs into, but ships no concrete discovery or transport (those depend on the
//! deployment's service mesh / gossip / control plane). An edge node implements
//! [`ClusterRelay`] to locate a stream's origin, pull it locally, and announce
//! the streams it serves.

use crate::bus::StreamHandle;
use crate::{MediaFrame, Result, StreamKey};
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;

/// Address of a node in the cluster (opaque to the engine).
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct NodeAddr(pub String);

/// A transport-generic source of frames for the **shared mirror loop**
/// ([`mirror`]). Both relays mirror the same way — replay the instant-start GOP,
/// then forward live frames — differing only in *where the frames come from*: an
/// in-process [`StreamHandle`] subscription, or a TCP `read_record` loop in the
/// network relay. Implement this for the transport and call [`mirror`] instead of
/// hand-copying the loop.
#[async_trait]
pub trait FrameSource: Send {
    /// Frames to replay immediately for instant start (the origin's cached GOP),
    /// in order. Defaults to none.
    fn replay(&mut self) -> Vec<MediaFrame> {
        Vec::new()
    }

    /// The next live frame, or `None` when the source has ended.
    async fn next(&mut self) -> Option<MediaFrame>;
}

/// Mirror `src` into `dst`: replay the instant-start buffer, then forward live
/// frames until the source ends, `cancel` fires, or the local mirror is released.
/// This is the one loop both [`InProcessRelay`] and the network relay share, so
/// the federation semantics can't drift between them.
pub async fn mirror<S: FrameSource>(mut src: S, dst: &StreamHandle, cancel: &CancellationToken) {
    for frame in src.replay() {
        if dst.publish_frame(frame).is_err() {
            return; // local mirror released
        }
    }
    loop {
        tokio::select! {
            _ = cancel.cancelled() => return,
            frame = src.next() => match frame {
                Some(frame) => {
                    if dst.publish_frame(frame).is_err() {
                        return;
                    }
                }
                None => return, // source ended
            },
        }
    }
}

impl<T: Into<String>> From<T> for NodeAddr {
    fn from(s: T) -> Self {
        NodeAddr(s.into())
    }
}

/// Origin/edge relay contract. Implement to federate streams across nodes:
/// an edge that lacks a stream locally [`locate`](Self::locate)s its origin and
/// [`pull`](Self::pull)s it; an origin [`announce`](Self::announce)s what it has.
#[async_trait]
pub trait ClusterRelay: Send + Sync + 'static {
    /// Find a node currently serving `key`, if any (e.g. via the control plane).
    async fn locate(&self, key: &StreamKey) -> Result<Option<NodeAddr>>;

    /// Begin relaying `key` from `origin` into the local engine, returning once
    /// the local mirror is publishing.
    async fn pull(&self, key: &StreamKey, origin: &NodeAddr) -> Result<()>;

    /// Advertise that this node serves `key` so edges can discover it.
    async fn announce(&self, key: &StreamKey) -> Result<()>;

    /// Withdraw a previous [`announce`](Self::announce) when the stream ends.
    async fn withdraw(&self, key: &StreamKey) -> Result<()>;

    /// Demand policy on top of the [`locate`](Self::locate)/[`pull`](Self::pull)
    /// mechanism: make `key` available locally, returning `true` if a mirror is
    /// now (or already) present and `false` if no origin serves it.
    ///
    /// A host serving a play/WHEP/HLS request calls this so it doesn't have to
    /// hand-write "locate, then pull if not already local." Implementations
    /// **should de-duplicate concurrent callers** for the same key (single-flight)
    /// so a burst of N simultaneous viewer requests for a not-yet-mirrored stream
    /// triggers **one** pull, not N racing `start_publish` calls.
    ///
    /// The default is a plain `locate` + `pull` with **no** coalescing or
    /// locality check; [`NetworkRelay`](../../arcly_stream_cluster/struct.NetworkRelay.html)
    /// overrides it with a locality fast-path and single-flight.
    async fn ensure_mirrored(&self, key: &StreamKey) -> Result<bool> {
        match self.locate(key).await? {
            Some(origin) => {
                self.pull(key, &origin).await?;
                Ok(true)
            }
            None => Ok(false),
        }
    }
}

#[cfg(feature = "cluster")]
pub use relay::{ClusterDirectory, InProcessRelay};

/// A working, in-process reference [`ClusterRelay`] (feature `cluster`).
///
/// It federates streams between nodes that live in the same process — the shape
/// integration tests and single-box multi-engine setups need — over an in-memory
/// [`ClusterDirectory`] control plane. A production edge swaps the directory for
/// its gossip/service-mesh and `pull` for a real transport (RTMP/SRT/QUIC), but
/// the mirror loop (subscribe origin → republish locally) is identical.
#[cfg(feature = "cluster")]
mod relay {
    use super::{ClusterRelay, FrameSource, NodeAddr};
    use crate::bus::{PlaybackRegistry, PublishRegistry, StreamHandle, Subscription};
    use crate::{MediaFrame, Result, StreamError, StreamKey};
    use async_trait::async_trait;
    use std::collections::HashMap;
    use std::time::{Duration, Instant};
    use tokio_util::sync::CancellationToken;

    /// A [`FrameSource`] over an in-process origin [`StreamHandle`]: snapshot the
    /// GOP first (for instant start), then forward the live subscription — so the
    /// in-process relay rides the shared [`mirror`](super::mirror) loop.
    struct HandleSource {
        replay: Vec<MediaFrame>,
        sub: Subscription,
    }

    impl HandleSource {
        fn new(handle: StreamHandle) -> Self {
            // Snapshot the replay buffer *before* subscribing, matching the
            // original ordering so live frames don't duplicate the GOP tail.
            let replay = handle
                .replay_buffer()
                .into_iter()
                .map(|f| (*f).clone())
                .collect();
            let sub = handle.subscribe_resilient();
            Self { replay, sub }
        }
    }

    #[async_trait]
    impl FrameSource for HandleSource {
        fn replay(&mut self) -> Vec<MediaFrame> {
            std::mem::take(&mut self.replay)
        }
        async fn next(&mut self) -> Option<MediaFrame> {
            self.sub.recv().await.map(|f| (*f).clone())
        }
    }
    use std::sync::{Arc, Mutex};

    /// In-memory cluster directory (control plane): which nodes serve which
    /// streams. Shared (`Arc`) by every node's relay in the process.
    ///
    /// By default entries live until an explicit [`withdraw`](Self::withdraw).
    /// Opting into a TTL with [`with_ttl`](Self::with_ttl) gives **lease parity**
    /// with the network [`DirectoryServer`](../../arcly_stream_cluster/struct.DirectoryServer.html):
    /// an [`announce`](Self::announce) is valid only for the TTL, so a peer that
    /// stops a stream without `withdraw` (e.g. it crashed) is dropped on its own,
    /// and `locate` then steers to a live replica instead of a dead node. Re-
    /// announce within the TTL to keep a lease (the engine's `PublishEnded` event
    /// makes a tiny auto-withdraw adapter easy, too).
    #[derive(Debug, Default)]
    pub struct ClusterDirectory {
        serving: Mutex<HashMap<StreamKey, HashMap<NodeAddr, Option<Instant>>>>,
        ttl: Option<Duration>,
    }

    impl ClusterDirectory {
        /// A fresh, empty directory whose entries never expire (withdraw-only).
        pub fn new() -> Self {
            Self::default()
        }

        /// A directory whose announces carry a `ttl` lease (network-relay parity);
        /// an entry not re-announced within `ttl` is reaped on the next access.
        pub fn with_ttl(ttl: Duration) -> Self {
            Self {
                serving: Mutex::new(HashMap::new()),
                ttl: Some(ttl),
            }
        }

        /// Record that `node` serves `key` (refreshing its lease if TTL is set).
        pub fn announce(&self, node: &NodeAddr, key: &StreamKey) {
            let expires = self.ttl.map(|t| Instant::now() + t);
            self.serving
                .lock()
                .unwrap()
                .entry(key.clone())
                .or_default()
                .insert(node.clone(), expires);
        }

        /// Drop `node` from `key`'s server set.
        pub fn withdraw(&self, node: &NodeAddr, key: &StreamKey) {
            if let Some(set) = self.serving.lock().unwrap().get_mut(key) {
                set.remove(node);
            }
        }

        /// A live node serving `key` other than `exclude`, if any. Expired leases
        /// are skipped (and lazily reaped).
        pub fn locate(&self, key: &StreamKey, exclude: &NodeAddr) -> Option<NodeAddr> {
            let now = Instant::now();
            let mut map = self.serving.lock().unwrap();
            let set = map.get_mut(key)?;
            set.retain(|_, exp| exp.is_none_or(|e| e > now)); // reap expired
            set.keys().find(|n| *n != exclude).cloned()
        }
    }

    /// An in-process [`ClusterRelay`] for one node: mirrors a peer's stream into
    /// the local engine.
    pub struct InProcessRelay {
        node: NodeAddr,
        local: Arc<dyn PublishRegistry>,
        directory: Arc<ClusterDirectory>,
        peers: HashMap<NodeAddr, Arc<dyn PlaybackRegistry>>,
        reconnect: bool,
        /// Active mirrors, so [`withdraw`](ClusterRelay::withdraw) can stop a
        /// reconnecting supervisor.
        mirrors: Mutex<HashMap<StreamKey, CancellationToken>>,
    }

    impl InProcessRelay {
        /// A relay for `node`, mirroring into `local`, discovering via `directory`.
        pub fn new(
            node: impl Into<NodeAddr>,
            local: Arc<dyn PublishRegistry>,
            directory: Arc<ClusterDirectory>,
        ) -> Self {
            Self {
                node: node.into(),
                local,
                directory,
                peers: HashMap::new(),
                reconnect: false,
                mirrors: Mutex::new(HashMap::new()),
            }
        }

        /// Register a peer node's playback registry so [`pull`](ClusterRelay::pull)
        /// can subscribe to its streams.
        pub fn with_peer(
            mut self,
            addr: impl Into<NodeAddr>,
            playback: Arc<dyn PlaybackRegistry>,
        ) -> Self {
            self.peers.insert(addr.into(), playback);
            self
        }

        /// Opt into **reconnect/failover parity** with the network relay: when a
        /// mirrored origin's stream ends, re-[`locate`](ClusterRelay::locate) the
        /// stream (excluding the node that just failed) and resume from a live
        /// replica, keeping the same local publish handle — viewers don't see the
        /// stream drop. Stops when the stream is gone everywhere or on
        /// [`withdraw`](ClusterRelay::withdraw). Off by default (one-shot), keeping
        /// the reference minimal. Pair with [`ClusterDirectory::with_ttl`] so a
        /// crashed origin's lease lapses and failover actually has somewhere to go.
        pub fn with_reconnect(mut self, reconnect: bool) -> Self {
            self.reconnect = reconnect;
            self
        }
    }

    #[async_trait]
    impl ClusterRelay for InProcessRelay {
        async fn locate(&self, key: &StreamKey) -> Result<Option<NodeAddr>> {
            Ok(self.directory.locate(key, &self.node))
        }

        async fn pull(&self, key: &StreamKey, origin: &NodeAddr) -> Result<()> {
            let peer = self.peers.get(origin).ok_or_else(|| {
                StreamError::protocol(format!("cluster: unknown origin node {}", origin.0))
            })?;
            // Resolve the origin's live handle and claim the local mirror once.
            let src = peer.get_stream(key)?;
            let dst = self.local.start_publish(key).await?;

            let cancel = CancellationToken::new();
            self.mirrors
                .lock()
                .unwrap()
                .insert(key.clone(), cancel.clone());

            // Mirror via the shared kernel loop (replay GOP → forward live →
            // release), so in-process and network relays stay in lock-step.
            let local = Arc::clone(&self.local);
            let directory = Arc::clone(&self.directory);
            let peers = self.peers.clone();
            let reconnect = self.reconnect;
            let key = key.clone();
            let mut origin = origin.clone();
            tokio::spawn(async move {
                super::mirror(HandleSource::new(src), &dst, &cancel).await;
                // Reconnect/failover supervisor: re-locate (excluding the failed
                // origin) and resume on the **same** local handle, until the
                // stream is gone everywhere or this mirror is withdrawn.
                while reconnect && !cancel.is_cancelled() {
                    let Some(next) = directory.locate(&key, &origin) else {
                        break; // gone everywhere
                    };
                    origin = next;
                    let Some(peer) = peers.get(&origin) else {
                        break;
                    };
                    let Ok(src) = peer.get_stream(&key) else {
                        break;
                    };
                    super::mirror(HandleSource::new(src), &dst, &cancel).await;
                }
                let _ = local.end_publish(&key).await;
            });
            Ok(())
        }

        async fn announce(&self, key: &StreamKey) -> Result<()> {
            self.directory.announce(&self.node, key);
            Ok(())
        }

        async fn withdraw(&self, key: &StreamKey) -> Result<()> {
            if let Some(cancel) = self.mirrors.lock().unwrap().remove(key) {
                cancel.cancel(); // stop a reconnecting supervisor
            }
            self.directory.withdraw(&self.node, key);
            Ok(())
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use crate::{AppSpec, CodecId, Engine, MediaFrame};

        fn frame(pts: i64) -> MediaFrame {
            MediaFrame::new_video(
                pts,
                pts,
                bytes::Bytes::from_static(b"x"),
                CodecId::H264,
                true,
            )
        }

        #[tokio::test]
        async fn edge_locates_and_mirrors_origin_stream() {
            let directory = Arc::new(ClusterDirectory::new());
            let key = StreamKey::new("live", "cam");

            // Origin node publishes the stream and announces it.
            let origin = Engine::builder()
                .application(AppSpec::new("live").gop_cache(8))
                .build();
            let origin_relay = InProcessRelay::new("origin", origin.clone(), directory.clone());
            let src_handle = origin.start_publish(&key).await.unwrap();
            origin_relay.announce(&key).await.unwrap();
            src_handle.publish_frame(frame(0)).unwrap();

            // Edge node has the stream locally? No — it locates + pulls the origin.
            let edge = Engine::builder()
                .application(AppSpec::new("live").gop_cache(8))
                .build();
            let edge_relay = InProcessRelay::new("edge", edge.clone(), directory.clone())
                .with_peer("origin", origin.clone());

            assert!(edge.get_stream(&key).is_err(), "not local yet");
            let found = edge_relay.locate(&key).await.unwrap();
            assert_eq!(found, Some(NodeAddr::from("origin")));
            edge_relay.pull(&key, &found.unwrap()).await.unwrap();

            // The local mirror is now publishing and receives forwarded frames.
            let mirror = edge.get_stream(&key).expect("local mirror exists");
            let mut sub = mirror.subscribe_resilient();
            src_handle.publish_frame(frame(1)).unwrap();
            let got = tokio::time::timeout(std::time::Duration::from_secs(5), sub.recv())
                .await
                .expect("a frame was mirrored")
                .expect("frame");
            assert!(got.is_video());
        }

        #[tokio::test]
        async fn locate_excludes_self_and_withdraw_clears() {
            let directory = Arc::new(ClusterDirectory::new());
            let key = StreamKey::new("live", "s");
            let engine = Engine::builder().application(AppSpec::new("live")).build();
            let relay = InProcessRelay::new("only", engine, directory.clone());

            relay.announce(&key).await.unwrap();
            // The only server is ourselves → locate returns None.
            assert_eq!(relay.locate(&key).await.unwrap(), None);
            relay.withdraw(&key).await.unwrap();
            assert!(directory.locate(&key, &NodeAddr::from("other")).is_none());
        }

        #[tokio::test]
        async fn ttl_directory_reaps_an_unrenewed_lease() {
            // Real (short) TTL: the directory uses std::Instant, not the mock clock.
            let dir = ClusterDirectory::with_ttl(Duration::from_millis(80));
            let key = StreamKey::new("live", "cam");
            dir.announce(&NodeAddr::from("origin"), &key);
            // Still leased now.
            assert_eq!(
                dir.locate(&key, &NodeAddr::from("edge")),
                Some(NodeAddr::from("origin"))
            );
            // Past the TTL without renewal → reaped, locate finds nothing.
            tokio::time::sleep(Duration::from_millis(160)).await;
            assert_eq!(dir.locate(&key, &NodeAddr::from("edge")), None);
            // Re-announce restores the lease (renewal parity).
            dir.announce(&NodeAddr::from("origin"), &key);
            assert!(dir.locate(&key, &NodeAddr::from("edge")).is_some());
        }

        #[tokio::test]
        async fn reconnect_fails_over_to_a_second_origin() {
            let directory = Arc::new(ClusterDirectory::new());
            let key = StreamKey::new("live", "cam");

            // Two origins (A, B) both publish + announce the same stream.
            let origin_a = Engine::builder()
                .application(AppSpec::new("live").gop_cache(8))
                .build();
            let origin_b = Engine::builder()
                .application(AppSpec::new("live").gop_cache(8))
                .build();
            let a = origin_a.start_publish(&key).await.unwrap();
            let b = origin_b.start_publish(&key).await.unwrap();
            directory.announce(&NodeAddr::from("A"), &key);
            directory.announce(&NodeAddr::from("B"), &key);
            a.publish_frame(frame(0)).unwrap();
            b.publish_frame(frame(0)).unwrap();

            // Edge mirrors with reconnect on, starting from A.
            let edge = Engine::builder()
                .application(AppSpec::new("live").gop_cache(8))
                .build();
            let relay = InProcessRelay::new("edge", edge.clone(), directory.clone())
                .with_peer("A", origin_a.clone())
                .with_peer("B", origin_b.clone())
                .with_reconnect(true);
            relay.pull(&key, &NodeAddr::from("A")).await.unwrap();

            let mirror = edge.get_stream(&key).expect("local mirror");
            let mut sub = mirror.subscribe_resilient();

            // A drops out of the directory and ends its stream → supervisor must
            // re-locate to B on the *same* local handle.
            directory.withdraw(&NodeAddr::from("A"), &key);
            origin_a.end_publish(&key).await.unwrap();

            // Frames from B now reach the unchanged local mirror.
            for n in 1..20 {
                b.publish_frame(frame(n)).unwrap();
            }
            let got = tokio::time::timeout(std::time::Duration::from_secs(5), sub.recv())
                .await
                .expect("a frame arrived after failover")
                .expect("frame");
            assert!(got.is_video());
            assert!(edge.get_stream(&key).is_ok(), "local handle stable");
        }
    }
}