arcly-stream 0.8.1

//! Transcoding & adaptive-bitrate contracts.
//!
//! The actual codec work needs native encoders (FFmpeg/GStreamer, NVENC/QSV via
//! [`HwAccelBackend`](crate::HwAccelBackend)) and is intentionally **not** shipped
//! here. This module defines the seams: a [`Transcoder`] turns one input frame
//! into zero or more output frames, and a [`RenditionSpec`] describes one rung
//! of an ABR ladder. [`drive_abr_ladder`] drives a [`Transcoder`] from a live
//! stream and republishes each output rendition to its layer-key stream, so the
//! WebRTC egress's per-viewer layer selection treats transcoded rungs exactly
//! like simulcast layers.

use crate::bus::{PlaybackRegistry, PublishRegistry};
use crate::frame::CodecId;
use crate::{MediaFrame, Result, StreamKey};
use async_trait::async_trait;
use tokio_util::sync::CancellationToken;

/// One target rendition in an adaptive-bitrate ladder.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RenditionSpec {
    /// Human label / variant id (e.g. `"720p"`).
    pub name: String,
    /// Target encoded width in pixels (0 = keep source).
    pub width: u32,
    /// Target encoded height in pixels (0 = keep source).
    pub height: u32,
    /// Target video bitrate in bits/sec.
    pub video_bitrate_bps: u64,
    /// Output video codec.
    pub codec: CodecId,
}

impl RenditionSpec {
    /// A 16:9 rendition of the given height at `video_bitrate_bps`.
    pub fn new(
        name: impl Into<String>,
        height: u32,
        video_bitrate_bps: u64,
        codec: CodecId,
    ) -> Self {
        Self {
            name: name.into(),
            width: height * 16 / 9,
            height,
            video_bitrate_bps,
            codec,
        }
    }
}

/// Transforms frames: decode/scale/encode, packetize, or fan into renditions.
///
/// Returning a `Vec` lets one input frame yield several outputs (e.g. one per
/// ABR rung) or none (e.g. while priming an encoder).
#[async_trait]
pub trait Transcoder: Send + Sync {
    /// Stable identifier (e.g. `"nvenc-abr"`).
    fn id(&self) -> &str;

    /// The renditions this transcoder emits.
    fn renditions(&self) -> &[RenditionSpec];

    /// Process one input frame into zero or more output frames.
    async fn transcode(&mut self, frame: MediaFrame) -> Result<Vec<MediaFrame>>;

    /// Flush any buffered output at end-of-stream.
    async fn flush(&mut self) -> Result<Vec<MediaFrame>> {
        Ok(Vec::new())
    }
}

/// The [`track_id`](crate::MediaFrame::track_id) a [`Transcoder`] stamps on the
/// **first** rendition's output frames; rendition *n* (its `renditions()[n]`)
/// uses `RENDITION_TRACK_BASE + n`. This is the contract
/// [`drive_abr_ladder`] relies on to route each output to its rung stream.
pub const RENDITION_TRACK_BASE: u32 = 100;

/// Drive an adaptive-bitrate ladder: subscribe to the `base` stream, run every
/// frame through `transcoder`, and **republish each rendition** to its layer
/// stream (`base.layer(&spec.name)`) so the WebRTC egress's per-viewer layer
/// selection treats the transcoded rungs exactly like simulcast layers.
///
/// Codec-agnostic by design — the heavy decode/scale/encode lives in the injected
/// [`Transcoder`] (e.g. the FFmpeg-backed `arcly-stream-transcode`); this driver
/// is the bus plumbing that connects it to the ABR layer-key convention. Each
/// output frame is routed by its [`track_id`](crate::MediaFrame::track_id) offset
/// from [`RENDITION_TRACK_BASE`]. Runs until the base stream ends or `shutdown`
/// fires.
pub async fn drive_abr_ladder(
    playback: &dyn PlaybackRegistry,
    publish: &dyn PublishRegistry,
    base: &StreamKey,
    transcoder: &mut dyn Transcoder,
    shutdown: &CancellationToken,
) -> Result<()> {
    // Open a publish session per rung up front, keyed by the shared layer
    // convention so the egress can discover them.
    let rung_keys: Vec<StreamKey> = transcoder
        .renditions()
        .iter()
        .map(|spec| base.layer(&spec.name))
        .collect();
    let mut rungs = Vec::with_capacity(rung_keys.len());
    for key in &rung_keys {
        rungs.push(publish.start_publish(key).await?);
    }

    // Route one batch of transcoder outputs to their rung streams by track_id.
    let route = |rungs: &[crate::bus::StreamHandle], outputs: Vec<MediaFrame>| -> Result<()> {
        for out in outputs {
            let idx = out.track_id.wrapping_sub(RENDITION_TRACK_BASE) as usize;
            if let Some(handle) = rungs.get(idx) {
                handle.publish_frame(out)?;
            }
        }
        Ok(())
    };

    let mut sub = playback.get_stream(base)?.subscribe_resilient();
    loop {
        tokio::select! {
            _ = shutdown.cancelled() => break,
            next = sub.recv() => match next {
                Some(frame) => route(&rungs, transcoder.transcode((*frame).clone()).await?)?,
                None => break, // base stream ended
            }
        }
    }
    route(&rungs, transcoder.flush().await?)?;
    for key in &rung_keys {
        publish.end_publish(key).await?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::PlaybackRegistry;
    use crate::inbound::IngestContext;
    use crate::{AppSpec, Engine};
    use std::sync::Arc;
    use tokio_util::sync::CancellationToken;

    /// A stand-in transcoder: fans every input frame out to one tagged output per
    /// rung (no real codec work), so the driver's routing can be exercised.
    struct FanoutTranscoder {
        ladder: Vec<RenditionSpec>,
    }

    #[async_trait]
    impl Transcoder for FanoutTranscoder {
        fn id(&self) -> &str {
            "fanout"
        }
        fn renditions(&self) -> &[RenditionSpec] {
            &self.ladder
        }
        async fn transcode(&mut self, frame: MediaFrame) -> Result<Vec<MediaFrame>> {
            Ok((0..self.ladder.len())
                .map(|n| {
                    let mut f = frame.clone();
                    f.track_id = RENDITION_TRACK_BASE + n as u32;
                    f
                })
                .collect())
        }
    }

    #[tokio::test]
    async fn republishes_renditions_to_layer_streams() {
        let engine = Engine::builder()
            .application(AppSpec::new("live").gop_cache(4))
            .build();
        let base = StreamKey::new("live", "cam");

        // Publish a keyframe to the base so the driver's instant-start replay has
        // something to transcode.
        let ctx = IngestContext::new(engine.clone());
        let session = ctx.open_publish(base.clone()).await.unwrap();
        session
            .publish_frame(MediaFrame::new_video(
                0,
                0,
                bytes::Bytes::from_static(&[0, 0, 0, 1, 0x65]),
                CodecId::H264,
                true,
            ))
            .unwrap();

        let shutdown = CancellationToken::new();
        let sd = shutdown.clone();
        let eng = Arc::clone(&engine);
        let driver = tokio::spawn(async move {
            let mut tc = FanoutTranscoder {
                ladder: vec![
                    RenditionSpec::new("720p", 720, 3_000_000, CodecId::H264),
                    RenditionSpec::new("480p", 480, 1_200_000, CodecId::H264),
                ],
            };
            drive_abr_ladder(eng.as_ref(), eng.as_ref(), &base, &mut tc, &sd)
                .await
                .unwrap();
        });

        // Wait for the rung streams to come up.
        let mut up = false;
        for _ in 0..200 {
            if engine
                .get_stream(&StreamKey::new("live", "cam~720p"))
                .is_ok()
                && engine
                    .get_stream(&StreamKey::new("live", "cam~480p"))
                    .is_ok()
            {
                up = true;
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(2)).await;
        }
        shutdown.cancel();
        let _ = driver.await;
        assert!(up, "both ABR rungs were republished as layer streams");
    }
}