every.channel/third_party/iroh-live/moq-media/src/subscribe.rs
2026-02-15 16:17:27 -05:00

712 lines
23 KiB
Rust

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,
})
}
}