every.channel: sanitized baseline

third_party/iroh-live/moq-media/src/subscribe.rs

@@ -0,0 +1,712 @@
+use std::{collections::BTreeMap, sync::Arc, time::Duration};
+
+use hang::{
+    Timestamp, TrackConsumer,
+    catalog::{AudioConfig, Catalog, CatalogConsumer, VideoConfig},
+};
+use moq_lite::{BroadcastConsumer, Track};
+use n0_error::{Result, StackResultExt, StdResultExt};
+use n0_future::task::AbortOnDropHandle;
+use n0_watcher::{Watchable, Watcher};
+use tokio::{
+    sync::mpsc::{self, error::TryRecvError},
+    time::Instant,
+};
+use tokio_util::sync::{CancellationToken, DropGuard};
+use tracing::{Span, debug, error, info, info_span, trace, warn};
+
+use crate::{
+    av::{
+        AudioDecoder, AudioSink, AudioSinkHandle, DecodeConfig, DecodedFrame, Decoders,
+        PlaybackConfig, Quality, VideoDecoder, VideoSource,
+    },
+    ffmpeg::util::Rescaler,
+    util::spawn_thread,
+};
+
+const DEFAULT_MAX_LATENCY: Duration = Duration::from_millis(150);
+
+#[derive(derive_more::Debug, Clone)]
+pub struct SubscribeBroadcast {
+    broadcast_name: String,
+    #[debug("BroadcastConsumer")]
+    broadcast: BroadcastConsumer,
+    // catalog_watcher: n0_watcher::Direct<CatalogWrapper>,
+    catalog_watchable: Watchable<CatalogWrapper>,
+    shutdown: CancellationToken,
+    _catalog_task: Arc<AbortOnDropHandle<()>>,
+}
+
+#[derive(Debug, derive_more::PartialEq, derive_more::Eq, Default, Clone, derive_more::Deref)]
+pub struct CatalogWrapper {
+    #[eq(skip)]
+    #[deref]
+    inner: Arc<Catalog>,
+    seq: usize,
+}
+
+impl CatalogWrapper {
+    fn new(inner: Catalog, seq: usize) -> Self {
+        Self {
+            inner: Arc::new(inner),
+            seq,
+        }
+    }
+
+    pub fn video_renditions(&self) -> impl Iterator<Item = &str> {
+        let mut renditions: Vec<_> = self
+            .inner
+            .video
+            .as_ref()
+            .iter()
+            .map(|v| v.renditions.iter())
+            .flatten()
+            .map(|(name, config)| (name.as_str(), config.coded_width))
+            .collect();
+        renditions.sort_by(|a, b| a.1.cmp(&b.1));
+        renditions.into_iter().map(|(name, _w)| name)
+    }
+
+    pub fn audio_renditions(&self) -> impl Iterator<Item = &str> + '_ {
+        self.inner
+            .audio
+            .as_ref()
+            .into_iter()
+            .map(|v| v.renditions.iter())
+            .flatten()
+            .map(|(name, _config)| name.as_str())
+    }
+
+    pub fn select_video_rendition(&self, quality: Quality) -> Result<String> {
+        let video = self.inner.video.as_ref().context("no video published")?;
+        let track_name =
+            select_video_rendition(&video.renditions, quality).context("no video renditions")?;
+        Ok(track_name)
+    }
+
+    pub fn select_audio_rendition(&self, quality: Quality) -> Result<String> {
+        let audio = self.inner.audio.as_ref().context("no video published")?;
+        let track_name =
+            select_audio_rendition(&audio.renditions, quality).context("no video renditions")?;
+        Ok(track_name)
+    }
+}
+
+impl CatalogWrapper {
+    pub fn into_inner(self) -> Arc<Catalog> {
+        self.inner
+    }
+}
+
+impl SubscribeBroadcast {
+    pub async fn new(broadcast_name: String, broadcast: BroadcastConsumer) -> Result<Self> {
+        let shutdown = CancellationToken::new();
+
+        let (catalog_watchable, catalog_task) = {
+            let track = broadcast.subscribe_track(&Catalog::default_track());
+            let mut consumer = CatalogConsumer::new(track);
+            let initial_catalog = consumer
+                .next()
+                .await
+                .std_context("Broadcast closed before receiving catalog")?
+                .context("Catalog track closed before receiving catalog")?;
+            let watchable = Watchable::new(CatalogWrapper::new(initial_catalog, 0));
+
+            let task = tokio::spawn({
+                let shutdown = shutdown.clone();
+                let watchable = watchable.clone();
+                async move {
+                    for seq in 1.. {
+                        match consumer.next().await {
+                            Ok(Some(catalog)) => {
+                                watchable.set(CatalogWrapper::new(catalog, seq)).ok();
+                            }
+                            Ok(None) => {
+                                debug!("subscribed broadcast catalog track ended");
+                                break;
+                            }
+                            Err(err) => {
+                                debug!("subscribed broadcast closed: {err:#}");
+                                break;
+                            }
+                        }
+                    }
+                    shutdown.cancel();
+                }
+            });
+            (watchable, task)
+        };
+        Ok(Self {
+            broadcast_name,
+            broadcast,
+            catalog_watchable,
+            _catalog_task: Arc::new(AbortOnDropHandle::new(catalog_task)),
+            shutdown: CancellationToken::new(),
+        })
+    }
+
+    pub fn broadcast_name(&self) -> &str {
+        &self.broadcast_name
+    }
+
+    pub fn catalog_watcher(&mut self) -> n0_watcher::Direct<CatalogWrapper> {
+        self.catalog_watchable.watch()
+    }
+
+    pub fn catalog(&self) -> CatalogWrapper {
+        self.catalog_watchable.get()
+    }
+
+    pub fn watch_and_listen<D: Decoders>(
+        self,
+        audio_out: impl AudioSink,
+        playback_config: PlaybackConfig,
+    ) -> Result<AvRemoteTrack> {
+        AvRemoteTrack::new::<D>(self, audio_out, playback_config)
+    }
+
+    pub fn watch<D: VideoDecoder>(&self) -> Result<WatchTrack> {
+        self.watch_with::<D>(&Default::default(), Quality::Highest)
+    }
+
+    pub fn watch_with<D: VideoDecoder>(
+        &self,
+        playback_config: &DecodeConfig,
+        quality: Quality,
+    ) -> Result<WatchTrack> {
+        let track_name = self.catalog().select_video_rendition(quality)?;
+        self.watch_rendition::<D>(playback_config, &track_name)
+    }
+
+    pub fn watch_rendition<D: VideoDecoder>(
+        &self,
+        playback_config: &DecodeConfig,
+        track_name: &str,
+    ) -> Result<WatchTrack> {
+        let catalog = self.catalog();
+        let video = catalog.video.as_ref().context("no video published")?;
+        let config = video
+            .renditions
+            .get(track_name)
+            .context("rendition not found")?;
+        let consumer = TrackConsumer::new(
+            self.broadcast.subscribe_track(&Track {
+                name: track_name.to_string(),
+                priority: video.priority,
+            }),
+            DEFAULT_MAX_LATENCY,
+        );
+        let span = info_span!("videodec", %track_name);
+        WatchTrack::from_consumer::<D>(
+            track_name.to_string(),
+            consumer,
+            &config,
+            playback_config,
+            self.shutdown.child_token(),
+            span,
+        )
+    }
+    pub fn listen<D: AudioDecoder>(&self, output: impl AudioSink) -> Result<AudioTrack> {
+        self.listen_with::<D>(Quality::Highest, output)
+    }
+
+    pub fn listen_with<D: AudioDecoder>(
+        &self,
+        quality: Quality,
+        output: impl AudioSink,
+    ) -> Result<AudioTrack> {
+        let track_name = self.catalog().select_audio_rendition(quality)?;
+        self.listen_rendition::<D>(&track_name, output)
+    }
+
+    pub fn listen_rendition<D: AudioDecoder>(
+        &self,
+        name: &str,
+        output: impl AudioSink,
+    ) -> Result<AudioTrack> {
+        let catalog = self.catalog();
+        let audio = catalog.audio.as_ref().context("no audio published")?;
+        let config = audio.renditions.get(name).context("rendition not found")?;
+        let consumer = TrackConsumer::new(
+            self.broadcast.subscribe_track(&Track {
+                name: name.to_string(),
+                priority: audio.priority,
+            }),
+            DEFAULT_MAX_LATENCY,
+        );
+        let span = info_span!("audiodec", %name);
+        AudioTrack::spawn::<D>(
+            name.to_string(),
+            consumer,
+            config.clone(),
+            output,
+            self.shutdown.child_token(),
+            span,
+        )
+    }
+
+    pub fn closed(&self) -> impl Future<Output = ()> + 'static {
+        self.broadcast.closed()
+    }
+
+    pub fn shutdown(&self) {
+        self.shutdown.cancel();
+    }
+}
+
+fn select_rendition<T, P: ToString>(
+    renditions: &BTreeMap<String, T>,
+    order: &[P],
+) -> Option<String> {
+    order
+        .iter()
+        .map(ToString::to_string)
+        .find(|k| renditions.contains_key(k.as_str()))
+        .or_else(|| renditions.keys().next().cloned())
+}
+
+fn select_video_rendition<'a, T>(
+    renditions: &'a BTreeMap<String, T>,
+    q: Quality,
+) -> Option<String> {
+    use crate::av::VideoPreset::*;
+    let order = match q {
+        Quality::Highest => [P1080, P720, P360, P180],
+        Quality::High => [P720, P360, P180, P1080],
+        Quality::Mid => [P360, P180, P720, P1080],
+        Quality::Low => [P180, P360, P720, P1080],
+    };
+
+    select_rendition(renditions, &order)
+}
+
+fn select_audio_rendition<'a, T>(
+    renditions: &'a BTreeMap<String, T>,
+    q: Quality,
+) -> Option<String> {
+    use crate::av::AudioPreset::*;
+    let order = match q {
+        Quality::Highest | Quality::High => [Hq, Lq],
+        Quality::Mid | Quality::Low => [Lq, Hq],
+    };
+    select_rendition(renditions, &order)
+}
+
+pub struct AudioTrack {
+    name: String,
+    handle: Box<dyn AudioSinkHandle>,
+    shutdown_token: CancellationToken,
+    _task_handle: AbortOnDropHandle<()>,
+    _thread_handle: std::thread::JoinHandle<()>,
+}
+
+impl AudioTrack {
+    pub(crate) fn spawn<D: AudioDecoder>(
+        name: String,
+        consumer: TrackConsumer,
+        config: AudioConfig,
+        output: impl AudioSink,
+        shutdown: CancellationToken,
+        span: Span,
+    ) -> Result<Self> {
+        let _guard = span.enter();
+        let (packet_tx, packet_rx) = mpsc::channel(32);
+        let output_format = output.format()?;
+        info!(?config, "audio thread start");
+        let decoder = D::new(&config, output_format)?;
+        let handle = output.handle();
+        let thread_name = format!("adec-{}", name);
+        let thread = spawn_thread(thread_name, {
+            let shutdown = shutdown.clone();
+            let span = span.clone();
+            move || {
+                let _guard = span.enter();
+                if let Err(err) = Self::run_loop(decoder, packet_rx, output, &shutdown) {
+                    error!("audio decoder failed: {err:#}");
+                }
+                info!("audio decoder thread stop");
+            }
+        });
+        let task = tokio::spawn(forward_frames(consumer, packet_tx));
+        Ok(Self {
+            name,
+            handle,
+            shutdown_token: shutdown,
+            _task_handle: AbortOnDropHandle::new(task),
+            _thread_handle: thread,
+        })
+    }
+
+    pub fn stopped(&self) -> impl Future<Output = ()> + 'static {
+        let shutdown_token = self.shutdown_token.clone();
+        async move { shutdown_token.cancelled().await }
+    }
+
+    pub fn rendition(&self) -> &str {
+        &self.name
+    }
+
+    pub fn handle(&self) -> &dyn AudioSinkHandle {
+        self.handle.as_ref()
+    }
+
+    pub(crate) fn run_loop(
+        mut decoder: impl AudioDecoder,
+        mut packet_rx: mpsc::Receiver<hang::Frame>,
+        mut sink: impl AudioSink,
+        shutdown: &CancellationToken,
+    ) -> Result<()> {
+        const INTERVAL: Duration = Duration::from_millis(10);
+        let mut remote_start = None;
+        let loop_start = Instant::now();
+
+        'main: for i in 0.. {
+            let tick = Instant::now();
+
+            if shutdown.is_cancelled() {
+                debug!("stop audio thread: cancelled");
+                break;
+            }
+
+            loop {
+                match packet_rx.try_recv() {
+                    Ok(packet) => {
+                        let remote_start = *remote_start.get_or_insert_with(|| packet.timestamp);
+
+                        if tracing::enabled!(tracing::Level::TRACE) {
+                            let loop_elapsed = tick.duration_since(loop_start);
+                            let remote_elapsed: Duration = packet
+                                .timestamp
+                                .checked_sub(remote_start)
+                                .unwrap_or(Timestamp::ZERO)
+                                .into();
+                            let diff_ms =
+                                (loop_elapsed.as_secs_f32() - remote_elapsed.as_secs_f32()) * 1000.;
+                            trace!(len = packet.payload.num_bytes(), ts=?packet.timestamp, ?loop_elapsed, ?remote_elapsed, ?diff_ms, "recv packet");
+                        }
+
+                        // TODO: Skip outdated packets?
+
+                        if !sink.is_paused() {
+                            decoder.push_packet(packet)?;
+                            if let Some(samples) = decoder.pop_samples()? {
+                                sink.push_samples(samples)?;
+                            }
+                        }
+                    }
+                    Err(TryRecvError::Disconnected) => {
+                        debug!("stop audio thread: packet_rx disconnected");
+                        break 'main;
+                    }
+                    Err(TryRecvError::Empty) => {
+                        trace!("no packet to recv");
+                        break;
+                    }
+                }
+            }
+
+            let target_time = i * INTERVAL;
+            let real_time = Instant::now().duration_since(loop_start);
+            let sleep = target_time.saturating_sub(real_time);
+            if !sleep.is_zero() {
+                std::thread::sleep(sleep);
+            }
+        }
+        shutdown.cancel();
+        Ok(())
+    }
+}
+
+impl Drop for AudioTrack {
+    fn drop(&mut self) {
+        self.shutdown_token.cancel();
+    }
+}
+
+pub struct WatchTrack {
+    video_frames: WatchTrackFrames,
+    handle: WatchTrackHandle,
+}
+
+pub struct WatchTrackHandle {
+    rendition: String,
+    viewport: Watchable<(u32, u32)>,
+    _guard: WatchTrackGuard,
+}
+
+impl WatchTrackHandle {
+    pub fn set_viewport(&self, w: u32, h: u32) {
+        self.viewport.set((w, h)).ok();
+    }
+
+    pub fn rendition(&self) -> &str {
+        &self.rendition
+    }
+}
+
+pub struct WatchTrackFrames {
+    rx: mpsc::Receiver<DecodedFrame>,
+}
+
+impl WatchTrackFrames {
+    pub fn current_frame(&mut self) -> Option<DecodedFrame> {
+        let mut out = None;
+        while let Ok(item) = self.rx.try_recv() {
+            out = Some(item);
+        }
+        out
+    }
+
+    pub async fn next_frame(&mut self) -> Option<DecodedFrame> {
+        if let Some(frame) = self.current_frame() {
+            Some(frame)
+        } else {
+            self.rx.recv().await
+        }
+    }
+}
+
+struct WatchTrackGuard {
+    _shutdown_token_guard: DropGuard,
+    _task_handle: Option<AbortOnDropHandle<()>>,
+    _thread_handle: Option<std::thread::JoinHandle<()>>,
+}
+
+impl WatchTrack {
+    pub fn empty(rendition: impl ToString) -> Self {
+        let (tx, rx) = mpsc::channel(1);
+        let task = tokio::task::spawn(async move {
+            std::future::pending::<()>().await;
+            let _ = tx;
+        });
+        let guard = WatchTrackGuard {
+            _shutdown_token_guard: CancellationToken::new().drop_guard(),
+            _task_handle: Some(AbortOnDropHandle::new(task)),
+            _thread_handle: None,
+        };
+        Self {
+            video_frames: WatchTrackFrames { rx },
+            handle: WatchTrackHandle {
+                rendition: rendition.to_string(),
+                viewport: Default::default(),
+                _guard: guard,
+            },
+        }
+    }
+
+    pub fn from_video_source(
+        rendition: String,
+        shutdown: CancellationToken,
+        mut source: impl VideoSource,
+        decode_config: DecodeConfig,
+    ) -> Self {
+        let viewport = Watchable::new((1u32, 1u32));
+        let (frame_tx, frame_rx) = tokio::sync::mpsc::channel::<DecodedFrame>(2);
+        let thread_name = format!("vpr-{:>4}-{:>4}", source.name(), rendition);
+        let thread = spawn_thread(thread_name, {
+            let mut viewport = viewport.watch();
+            let shutdown = shutdown.clone();
+            move || {
+                // TODO: Make configurable.
+                let fps = 30;
+                let mut rescaler = Rescaler::new(decode_config.pixel_format.to_ffmpeg(), None)
+                    .expect("failed to create rescaler");
+                let frame_duration = Duration::from_secs_f32(1. / fps as f32);
+                if let Err(err) = source.start() {
+                    warn!("Video source failed to start: {err:?}");
+                    return;
+                }
+                let start = Instant::now();
+                for i in 1.. {
+                    // let t = Instant::now();
+                    if shutdown.is_cancelled() {
+                        break;
+                    }
+                    if viewport.update() {
+                        let (w, h) = viewport.peek();
+                        rescaler.set_target_dimensions(*w, *h);
+                    }
+                    match source.pop_frame() {
+                        Ok(Some(frame)) => {
+                            // trace!(t=?t.elapsed(), "pop");
+                            let frame = frame.to_ffmpeg();
+                            let frame = rescaler.process(&frame).expect("rescaler failed");
+                            let frame =
+                                DecodedFrame::from_ffmpeg(frame, frame_duration, start.elapsed());
+                            // trace!(t=?t.elapsed(), "convert");
+                            let _ = frame_tx.blocking_send(frame);
+                            // trace!(t=?t.elapsed(), "send");
+                        }
+                        Ok(None) => {}
+                        Err(_) => break,
+                    }
+                    let expected_time = i * frame_duration;
+                    let actual_time = start.elapsed();
+                    if expected_time > actual_time {
+                        std::thread::sleep(expected_time - actual_time);
+                        // trace!(t=?t.elapsed(), slept=?(actual_time - expected_time), ?expected_time, ?actual_time, "done");
+                    }
+                }
+                if let Err(err) = source.stop() {
+                    warn!("Video source failed to stop: {err:?}");
+                    return;
+                }
+            }
+        });
+        let guard = WatchTrackGuard {
+            _shutdown_token_guard: shutdown.drop_guard(),
+            _task_handle: None,
+            _thread_handle: Some(thread),
+        };
+        WatchTrack {
+            video_frames: WatchTrackFrames { rx: frame_rx },
+            handle: WatchTrackHandle {
+                rendition,
+                viewport,
+                _guard: guard,
+            },
+        }
+    }
+
+    pub(crate) fn from_consumer<D: VideoDecoder>(
+        rendition: String,
+        consumer: TrackConsumer,
+        config: &VideoConfig,
+        playback_config: &DecodeConfig,
+        shutdown: CancellationToken,
+        span: Span,
+    ) -> Result<Self> {
+        let (packet_tx, packet_rx) = mpsc::channel(32);
+        let (frame_tx, frame_rx) = mpsc::channel(32);
+        let viewport = Watchable::new((1u32, 1u32));
+        let viewport_watcher = viewport.watch();
+
+        let _guard = span.enter();
+        debug!(?config, "video decoder start");
+        let decoder = D::new(config, playback_config)?;
+        let thread_name = format!("vdec-{}", rendition);
+        let thread = spawn_thread(thread_name, {
+            let shutdown = shutdown.clone();
+            let span = span.clone();
+            move || {
+                let _guard = span.enter();
+                if let Err(err) =
+                    Self::run_loop(&shutdown, packet_rx, frame_tx, viewport_watcher, decoder)
+                {
+                    error!("video decoder failed: {err:#}");
+                }
+                shutdown.cancel();
+            }
+        });
+        let task = tokio::task::spawn(forward_frames(consumer, packet_tx));
+        let guard = WatchTrackGuard {
+            _shutdown_token_guard: shutdown.drop_guard(),
+            _task_handle: Some(AbortOnDropHandle::new(task)),
+            _thread_handle: Some(thread),
+        };
+        Ok(WatchTrack {
+            video_frames: WatchTrackFrames { rx: frame_rx },
+            handle: WatchTrackHandle {
+                rendition,
+                viewport,
+                _guard: guard,
+            },
+        })
+    }
+
+    pub fn split(self) -> (WatchTrackFrames, WatchTrackHandle) {
+        (self.video_frames, self.handle)
+    }
+
+    pub fn set_viewport(&self, w: u32, h: u32) {
+        self.handle.set_viewport(w, h);
+    }
+
+    pub fn rendition(&self) -> &str {
+        self.handle.rendition()
+    }
+
+    pub fn current_frame(&mut self) -> Option<DecodedFrame> {
+        self.video_frames.current_frame()
+    }
+
+    pub(crate) fn run_loop(
+        shutdown: &CancellationToken,
+        mut input_rx: mpsc::Receiver<hang::Frame>,
+        output_tx: mpsc::Sender<DecodedFrame>,
+        mut viewport_watcher: n0_watcher::Direct<(u32, u32)>,
+        mut decoder: impl VideoDecoder,
+    ) -> Result<(), anyhow::Error> {
+        loop {
+            if shutdown.is_cancelled() {
+                break;
+            }
+            let Some(packet) = input_rx.blocking_recv() else {
+                break;
+            };
+            if viewport_watcher.update() {
+                let (w, h) = viewport_watcher.peek();
+                decoder.set_viewport(*w, *h);
+            }
+            let t = Instant::now();
+            decoder
+                .push_packet(packet)
+                .context("failed to push packet")?;
+            trace!(t=?t.elapsed(), "videodec: push_packet");
+            while let Some(frame) = decoder.pop_frame().context("failed to pop frame")? {
+                trace!(t=?t.elapsed(), "videodec: pop frame");
+                if output_tx.blocking_send(frame).is_err() {
+                    break;
+                }
+                trace!(t=?t.elapsed(), "videodec: tx");
+            }
+        }
+        Ok(())
+    }
+}
+
+async fn forward_frames(mut track: hang::TrackConsumer, sender: mpsc::Sender<hang::Frame>) {
+    loop {
+        let frame = track.read_frame().await;
+        match frame {
+            Ok(Some(frame)) => {
+                if sender.send(frame).await.is_err() {
+                    break;
+                }
+            }
+            Ok(None) => break,
+            Err(err) => {
+                warn!("failed to read frame: {err:?}");
+                break;
+            }
+        }
+    }
+}
+
+pub struct AvRemoteTrack {
+    pub broadcast: SubscribeBroadcast,
+    pub video: Option<WatchTrack>,
+    pub audio: Option<AudioTrack>,
+}
+
+impl AvRemoteTrack {
+    pub fn new<D: Decoders>(
+        broadcast: SubscribeBroadcast,
+        audio_out: impl AudioSink,
+        playback_config: PlaybackConfig,
+    ) -> Result<Self> {
+        let audio = broadcast
+            .listen_with::<D::Audio>(playback_config.quality, audio_out)
+            .inspect_err(|err| tracing::warn!("no audio track: {err}"))
+            .ok();
+        let video = broadcast
+            .watch_with::<D::Video>(&playback_config.decode_config, playback_config.quality)
+            .inspect_err(|err| tracing::warn!("no video track: {err}"))
+            .ok();
+        Ok(Self {
+            broadcast,
+            audio,
+            video,
+        })
+    }
+}