Better handling of decoder spawn

Author: JBRS307

smelter-api/src/input/moq_into.rs

@@ -25,7 +25,7 @@ impl TryFrom<MoqInput> for core::RegisterInputOptions {
 
         let input_options = core::MoqServerInputOptions {
             broadcast_path,
-            decoders: core::MoqServerInputDecoders { h264 },
+            decoders: core::MoqServerInputDecoders { h264, aac: None },
             queue_options: core::QueueInputOptions {
                 required: required.unwrap_or(false),
                 video_side_channel: side_channel.video.unwrap_or(false),

smelter-core/src/pipeline/moq/connection.rs

@@ -11,7 +11,9 @@ use tracing::{debug, info, trace, warn};
 use crate::utils::{H264AvcDecoderConfig, H264AvccToAnnexB};
 use crate::{
     MediaKind, PipelineCtx, PipelineEvent,
-    codecs::{AudioCodec, FdkAacDecoderOptions, VideoCodec, VideoDecoderOptions},
+    codecs::{
+        AudioCodec, AudioDecoderOptions, FdkAacDecoderOptions, VideoCodec, VideoDecoderOptions,
+    },
     error::DecoderInitError,
     pipeline::{
         decoder::{
@@ -77,13 +79,6 @@ pub(crate) fn spawn_broadcast_handler(
             }
         };
 
-        if let Some(v) = &discovered.video {
-            info!(input_id = %input_id_str, track = %v.name, "Discovered MoQ video track");
-        }
-        if let Some(a) = &discovered.audio {
-            info!(input_id = %input_id_str, track = %a.name, "Discovered MoQ audio track");
-        }
-
         let has_video = discovered.video.is_some();
         let has_audio = discovered.audio.is_some();
 
@@ -93,20 +88,70 @@ pub(crate) fn spawn_broadcast_handler(
             offset: QueueTrackOffset::Pts(ctx.queue_ctx.effective_last_pts() + MOQ_BUFFER),
         });
 
+        if let Some(v) = &discovered.video {
+            info!(input_id = %input_id_str, track = %v.name, "Discovered MoQ video track");
+        }
+        if let Some(a) = &discovered.audio {
+            info!(input_id = %input_id_str, track = %a.name, "Discovered MoQ audio track");
+        }
+
+        let video_decoder_handle = match &discovered.video {
+            Some(video) => {
+                match process_video_config(
+                    &ctx,
+                    &input_ref,
+                    &decoders,
+                    video,
+                    // TODO: (@jbrs) Handle it nicer in the future
+                    video_sender.expect("MUST be initialized only if video is present"),
+                ) {
+                    Ok(handle) => Some(handle),
+                    Err(err) => {
+                        warn!(
+                            "MoQ video config error: {}",
+                            ErrorStack::new(&err).into_string()
+                        );
+                        None
+                    }
+                }
+            }
+            None => None,
+        };
+
+        let audio_decoder_handle = match &discovered.audio {
+            Some(audio) => {
+                match process_audio_config(
+                    &ctx,
+                    &input_ref,
+                    audio,
+                    audio_sender.expect("MUST be initialized only if audio is present"),
+                ) {
+                    Ok(handle) => Some(handle),
+                    Err(err) => {
+                        warn!(
+                            "MoQ audio config error: {}",
+                            ErrorStack::new(&err).into_string()
+                        );
+                        None
+                    }
+                }
+            }
+            None => None,
+        };
+
         let video_fut = async {
-            if let Some(video) = discovered.video {
-                let avcc = video.description.clone();
+            if let Some(video) = discovered.video
+                && let Some(decoder_handle) = video_decoder_handle
+            {
                 match broadcast.subscribe_track(&Track::new(&video.name)) {
                     Ok(track) => {
                         let consumer =
                             ContainerConsumer::new(track, video.container).with_latency(MOQ_BUFFER);
                         if let Err(err) = read_video_track(
                             ctx.clone(),
                             input_ref.clone(),
-                            decoders.clone(),
                             consumer,
-                            avcc,
-                            video_sender,
+                            decoder_handle,
                         )
                         .await
                         {
@@ -124,8 +169,9 @@ pub(crate) fn spawn_broadcast_handler(
         };
 
         let audio_fut = async {
-            if let Some(audio) = discovered.audio {
-                let asc = audio.description.clone();
+            if let Some(audio) = discovered.audio
+                && let Some(decoder_handle) = audio_decoder_handle
+            {
                 match broadcast.subscribe_track(&Track::new(&audio.name)) {
                     Ok(track) => {
                         let consumer =
@@ -134,8 +180,7 @@ pub(crate) fn spawn_broadcast_handler(
                             ctx.clone(),
                             input_ref.clone(),
                             consumer,
-                            asc,
-                            audio_sender,
+                            decoder_handle,
                         )
                         .await
                         {
@@ -224,6 +269,83 @@ async fn read_catalog(
     Ok(DiscoveredTracks { video, audio })
 }
 
+fn process_video_config(
+    ctx: &Arc<PipelineCtx>,
+    input_ref: &Ref<InputId>,
+    decoders: &MoqServerInputDecoders,
+    video: &DiscoveredVideo,
+    frame_sender: QueueSender<Frame>,
+) -> Result<DecoderThreadHandle, MoqConnectionError> {
+    // TODO: (@jbrs) Proper error
+    let avcc_bytes = video
+        .description
+        .clone()
+        .ok_or(MoqConnectionError::UnsupportedVideoCodec)?;
+    let h264_config = H264AvcDecoderConfig::parse(avcc_bytes).unwrap();
+    let options = VideoDecoderThreadOptions {
+        ctx: ctx.clone(),
+        transformer: Some(H264AvccToAnnexB::new(h264_config)),
+        frame_sender,
+        input_buffer_size: MOQ_MAX_BUFFER,
+    };
+
+    let h264_decoder =
+        decoders
+            .h264
+            .unwrap_or_else(|| match ctx.graphics_context.has_vulkan_decoder_support() {
+                true => VideoDecoderOptions::VulkanH264,
+                false => VideoDecoderOptions::FfmpegH264,
+            });
+
+    match h264_decoder {
+        VideoDecoderOptions::FfmpegH264 => {
+            VideoDecoderThread::<ffmpeg_h264::FfmpegH264Decoder, _>::spawn(
+                input_ref.clone(),
+                options,
+            )
+            .map_err(MoqConnectionError::InitVideoDecoder)
+        }
+        VideoDecoderOptions::VulkanH264 => {
+            VideoDecoderThread::<vulkan_h264::VulkanH264Decoder, _>::spawn(
+                input_ref.clone(),
+                options,
+            )
+            .map_err(MoqConnectionError::InitVideoDecoder)
+        }
+        _ => Err(MoqConnectionError::UnsupportedVideoCodec),
+    }
+}
+
+fn process_audio_config(
+    ctx: &Arc<PipelineCtx>,
+    input_ref: &Ref<InputId>,
+    audio: &DiscoveredAudio,
+    samples_sender: QueueSender<InputAudioSamples>,
+) -> Result<DecoderThreadHandle, MoqConnectionError> {
+    let aac_decoder_options = {
+        let asc = audio
+            .description
+            .as_ref()
+            .expect("process_audio_config called with description present")
+            .clone();
+        AudioDecoderOptions::FdkAac(FdkAacDecoderOptions { asc: Some(asc) })
+    };
+
+    match aac_decoder_options {
+        AudioDecoderOptions::FdkAac(decoder_options) => {
+            let options = AudioDecoderThreadOptions {
+                ctx: ctx.clone(),
+                decoder_options,
+                samples_sender,
+                input_buffer_size: MOQ_MAX_BUFFER,
+            };
+            AudioDecoderThread::<FdkAacDecoder>::spawn(input_ref.clone(), options)
+                .map_err(MoqConnectionError::InitAudioDecoder)
+        }
+        _ => Err(MoqConnectionError::UnsupportedAudioCodec),
+    }
+}
+
 #[derive(thiserror::Error, Debug)]
 enum MoqConnectionError {
     #[error("MoQ track error")]
@@ -235,20 +357,20 @@ enum MoqConnectionError {
     #[error("Catalog track produced no frames")]
     CatalogEmpty,
 
-    #[error("Failed to read catalog")]
-    CatalogError(#[source] moq_mux::Error),
-
     #[error("Catalog contains no recognizable video or audio tracks")]
     CatalogNoTracks,
 
     #[error("Failed to initialize H264 decoder")]
-    InitH264Decoder(#[source] DecoderInitError),
+    InitVideoDecoder(#[source] DecoderInitError),
 
-    #[error("Invalid video decoder provided, expected H264 decoder")]
-    InvalidVideoDecoder,
+    #[error("Unsupported video codec, H264 expected.")]
+    UnsupportedVideoCodec,
 
     #[error("Failed to initialize AAC decoder")]
-    InitAacDecoder(#[source] DecoderInitError),
+    InitAudioDecoder(#[source] DecoderInitError),
+
+    #[error("Unsupported audio codec, AAC expected.")]
+    UnsupportedAudioCodec,
 
     #[error("Decoder channel closed")]
     ChannelClosed,
@@ -257,25 +379,14 @@ enum MoqConnectionError {
     ContainerError(#[source] moq_mux::Error),
 }
 
-// TODO: This is absolutely terrible and cannot be squashed in one function like that
 async fn read_video_track(
     ctx: Arc<PipelineCtx>,
     input_ref: Ref<InputId>,
-    decoders: MoqServerInputDecoders,
     mut consumer: ContainerConsumer<Hang>,
-    avcc_from_catalog: Option<Bytes>,
-    frame_sender: Option<QueueSender<Frame>>,
+    decoder_handle: DecoderThreadHandle,
 ) -> Result<(), MoqConnectionError> {
-    let Some(frame_sender) = frame_sender else {
-        return Ok(());
-    };
-
-    let mut decoder_handle: Option<DecoderThreadHandle> = None;
-    let mut pending_frame_sender = Some(frame_sender);
     let mut first_pts: Option<Duration> = None;
 
-    let h264_config = H264AvcDecoderConfig::parse(avcc_from_catalog.unwrap()).unwrap();
-
     while let Some(frame) = consumer
         .read()
         .await
@@ -287,48 +398,6 @@ async fn read_video_track(
         trace!(?pts, "MoQ video frame");
         let payload = frame.payload;
 
-        if decoder_handle.is_none() {
-            let frame_sender = pending_frame_sender
-                .take()
-                .expect("frame_sender consumed only once");
-            let options = VideoDecoderThreadOptions {
-                ctx: ctx.clone(),
-                transformer: Some(H264AvccToAnnexB::new(h264_config.clone())),
-                frame_sender,
-                input_buffer_size: MOQ_MAX_BUFFER,
-            };
-
-            let h264_decoder = decoders.h264.unwrap_or_else(|| {
-                match ctx.graphics_context.has_vulkan_decoder_support() {
-                    true => VideoDecoderOptions::VulkanH264,
-                    false => VideoDecoderOptions::FfmpegH264,
-                }
-            });
-
-            let handle = match h264_decoder {
-                VideoDecoderOptions::FfmpegH264 => {
-                    VideoDecoderThread::<ffmpeg_h264::FfmpegH264Decoder, _>::spawn(
-                        input_ref.clone(),
-                        options,
-                    )
-                    .map_err(MoqConnectionError::InitH264Decoder)?
-                }
-                VideoDecoderOptions::VulkanH264 => {
-                    VideoDecoderThread::<vulkan_h264::VulkanH264Decoder, _>::spawn(
-                        input_ref.clone(),
-                        options,
-                    )
-                    .map_err(MoqConnectionError::InitH264Decoder)?
-                }
-                _ => {
-                    return Err(MoqConnectionError::InvalidVideoDecoder);
-                }
-            };
-
-            decoder_handle = Some(handle);
-        }
-
-        let handle = decoder_handle.as_ref().unwrap();
         let chunk = EncodedInputChunk {
             data: payload,
             pts,
@@ -342,7 +411,7 @@ async fn read_video_track(
                 .into_event(&input_ref, StatsTrackKind::Video),
         );
 
-        handle
+        decoder_handle
             .chunk_sender
             .send(PipelineEvent::Data(chunk))
             .map_err(|_| MoqConnectionError::ChannelClosed)?;
@@ -355,15 +424,8 @@ async fn read_audio_track(
     ctx: Arc<PipelineCtx>,
     input_ref: Ref<InputId>,
     mut consumer: ContainerConsumer<Hang>,
-    asc_from_catalog: Option<Bytes>,
-    samples_sender: Option<QueueSender<InputAudioSamples>>,
+    decoder_handle: DecoderThreadHandle,
 ) -> Result<(), MoqConnectionError> {
-    let Some(samples_sender) = samples_sender else {
-        return Ok(());
-    };
-
-    let mut decoder_handle: Option<DecoderThreadHandle> = None;
-    let mut pending_samples_sender = Some(samples_sender);
     let mut first_pts: Option<Duration> = None;
 
     while let Some(frame) = consumer
@@ -377,39 +439,6 @@ async fn read_audio_track(
         trace!(?pts, "MoQ audio frame");
         let payload = frame.payload;
 
-        if decoder_handle.is_none() {
-            let samples_sender = pending_samples_sender
-                .take()
-                .expect("samples_sender consumed only once");
-
-            let asc = if let Some(ref asc) = asc_from_catalog {
-                asc.clone()
-            } else {
-                let asc = payload.clone();
-                let options = AudioDecoderThreadOptions {
-                    ctx: ctx.clone(),
-                    decoder_options: FdkAacDecoderOptions { asc: Some(asc) },
-                    samples_sender,
-                    input_buffer_size: MOQ_MAX_BUFFER,
-                };
-                let handle = AudioDecoderThread::<FdkAacDecoder>::spawn(input_ref.clone(), options)
-                    .map_err(MoqConnectionError::InitAacDecoder)?;
-                decoder_handle = Some(handle);
-                continue;
-            };
-
-            let options = AudioDecoderThreadOptions {
-                ctx: ctx.clone(),
-                decoder_options: FdkAacDecoderOptions { asc: Some(asc) },
-                samples_sender,
-                input_buffer_size: MOQ_MAX_BUFFER,
-            };
-            let handle = AudioDecoderThread::<FdkAacDecoder>::spawn(input_ref.clone(), options)
-                .map_err(MoqConnectionError::InitAacDecoder)?;
-            decoder_handle = Some(handle);
-        }
-
-        let handle = decoder_handle.as_ref().unwrap();
         let chunk = EncodedInputChunk {
             data: payload,
             pts,
@@ -423,7 +452,7 @@ async fn read_audio_track(
                 .into_event(&input_ref, StatsTrackKind::Audio),
         );
 
-        handle
+        decoder_handle
             .chunk_sender
             .send(PipelineEvent::Data(chunk))
             .map_err(|_| MoqConnectionError::ChannelClosed)?;