video: first frame within ~150 ms after connect or takeover

app/src/main.rs

@@ -40,7 +40,6 @@ fn init_tracing(cli: &Cli) {
 
     let fmt_layer = tracing_subscriber::fmt::layer()
         .with_ansi(!cli.logging.log_no_color)
-        .without_time()
         .with_level(true)
         .with_target(true)
         .with_thread_ids(true)

app/src/pipeline.rs

@@ -90,6 +90,9 @@ pub struct VideoPipeline {
     running: Arc<AtomicBool>,
     raw_frame_waiters: Arc<RawFrameWaiters>,
     raw_frame_requested: Arc<AtomicBool>,
+    /// When the oldest unanswered keyframe request was made; the encoder
+    /// output probe logs how long it took to produce the keyframe.
+    keyframe_requested_at: Arc<parking_lot::Mutex<Option<std::time::Instant>>>,
 }
 
 impl VideoPipeline {
@@ -167,6 +170,34 @@ impl VideoPipeline {
             encoder.set_property("gop", gop);
         }
 
+        let keyframe_requested_at: Arc<parking_lot::Mutex<Option<std::time::Instant>>> =
+            Arc::default();
+        let probe_requested_at = keyframe_requested_at.clone();
+        let is_h264 = matches!(config.encoder, crate::cli::VideoEncoder::H264);
+        encoder.static_pad("src").context("Encoder has no src pad")?.add_probe(
+            gst::PadProbeType::BUFFER,
+            move |_pad, info| {
+                // Decide by the bitstream, not the buffer flags: the MPP encoder
+                // does not reliably clear DELTA_UNIT on IDR frames.
+                if is_h264
+                    && let Some(buffer) = info.buffer()
+                    && let Ok(map) = buffer.map_readable()
+                    && contains_h264_idr(map.as_slice())
+                {
+                    let delta = buffer.flags().contains(gst::BufferFlags::DELTA_UNIT);
+                    match probe_requested_at.lock().take() {
+                        Some(at) => info!(
+                            elapsed_ms = at.elapsed().as_millis() as u64,
+                            delta_flag = delta,
+                            "encoder emitted IDR after request"
+                        ),
+                        None => debug!(delta_flag = delta, "encoder emitted periodic IDR"),
+                    }
+                }
+                gst::PadProbeReturn::Ok
+            },
+        );
+
         let parser_name = config.encoder.parser_element_name();
         let parser = gst::ElementFactory::make(parser_name)
             .property("config-interval", -1i32)
@@ -271,6 +302,7 @@ impl VideoPipeline {
             running,
             raw_frame_waiters,
             raw_frame_requested,
+            keyframe_requested_at,
         })
     }
 
@@ -780,14 +812,61 @@ impl PipelineManager {
         self.video.request_raw_frame()
     }
 
+    /// Requests an IDR from the encoder.
+    ///
+    /// The request carries the current running time. An "as soon as possible"
+    /// request (no running time) only works once per pipeline: GstVideoEncoder
+    /// keeps it pending until the encoder outputs a frame flagged as a sync
+    /// point, which the MPP encoder never does, and while one is pending it
+    /// ignores every later ASAP request. Timed requests are only compared with
+    /// the last forced request, so each one goes through.
+    ///
+    /// Requests are skipped while an earlier one is still waiting for its IDR
+    /// (up to a GOP), so a burst of browser PLIs costs one keyframe, and so
+    /// the encoder's never-cleared request queue grows by a few entries per
+    /// session rather than per PLI.
     pub fn force_keyframe(&self) {
-        let event = gst_video::UpstreamForceKeyUnitEvent::builder().all_headers(true).build();
-        if !self.video.encoder.send_event(event) {
+        {
+            let requested_at = self.video.keyframe_requested_at.lock();
+            if requested_at.is_some_and(|at| at.elapsed() < std::time::Duration::from_secs(1)) {
+                debug!("keyframe already requested, waiting for it");
+                return;
+            }
+        }
+        self.force_keyframe_now();
+    }
+
+    /// Like [`Self::force_keyframe`] but never coalesced: a new viewer needs
+    /// an IDR encoded after its connection is up, even if one is in flight.
+    pub fn force_keyframe_now(&self) {
+        self.video.keyframe_requested_at.lock().get_or_insert_with(std::time::Instant::now);
+        let mut builder = gst_video::UpstreamForceKeyUnitEvent::builder().all_headers(true);
+        if let Some(running_time) = self.video.encoder.current_running_time() {
+            builder = builder.running_time(running_time);
+        }
+        let handled = self.video.encoder.send_event(builder.build());
+        info!(handled, "keyframe requested");
+        if !handled {
             warn!("force-keyframe event was not handled by encoder");
         }
     }
 }
 
+/// Whether an Annex B H.264 access unit is an IDR: the first slice NAL found
+/// is type 5. Stops at the first slice, so a P frame costs only its headers.
+fn contains_h264_idr(data: &[u8]) -> bool {
+    for w in data.windows(4) {
+        if w[0] == 0 && w[1] == 0 && w[2] == 1 {
+            match w[3] & 0x1f {
+                5 => return true,
+                1 => return false,
+                _ => {}
+            }
+        }
+    }
+    false
+}
+
 pub fn encode_jpeg(sample: &gst::Sample) -> Result<bytes::Bytes> {
     let caps = sample.caps_owned().context("Raw frame has no caps")?;
     let encoder = gst::ElementFactory::make("mppjpegenc")

app/src/video.rs

@@ -69,42 +69,23 @@ type AudioFrame = bytes::Bytes;
 static VIDEO_FRAME_TX: tokio::sync::OnceCell<mpsc::Sender<VideoFrame>> =
     tokio::sync::OnceCell::const_new();
 
+/// Writes encoded frames to the current session's track. The sink is read per
+/// frame rather than cached: a cached track kept feeding a taken-over session
+/// for up to a second after the new one attached, delaying its first frame.
 fn video_writer_loop(mut rx: mpsc::Receiver<VideoFrame>) {
     tokio::spawn(async move {
         let mut last_pts: Option<u64> = None;
-        let mut current_track: Option<Arc<TrackLocalStaticSample>> = None;
-        let mut track_refresh = interval(Duration::from_secs(1));
-        track_refresh.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
-
-        loop {
-            tokio::select! {
-                _ = track_refresh.tick() => {
-                    current_track = VIDEO_SINK.read().await.clone();
-                }
-                maybe_frame = rx.recv() => {
-                    let Some((data, pts)) = maybe_frame else { break; };
-
-                    let dur = match last_pts {
-                        Some(prev) if pts > prev => Duration::from_micros(pts - prev),
-                        _ => Duration::from_millis(16),
-                    };
-                    last_pts = Some(pts);
-
-                    if current_track.is_none() {
-                        current_track = VIDEO_SINK.read().await.clone();
-                    }
-                    if let Some(track) = &current_track {
-                        let sample = Sample {
-                            data,
-                            duration: dur,
-                            ..Default::default()
-                        };
-                        if let Err(e) = track.write_sample(&sample).await {
-                            warn!("Failed to write video sample: {}", e);
-                            current_track = None;
-                        }
-                    }
-                }
+        while let Some((data, pts)) = rx.recv().await {
+            let dur = match last_pts {
+                Some(prev) if pts > prev => Duration::from_micros(pts - prev),
+                _ => Duration::from_millis(16),
+            };
+            last_pts = Some(pts);
+
+            let Some(track) = VIDEO_SINK.read().await.clone() else { continue };
+            let sample = Sample { data, duration: dur, ..Default::default() };
+            if let Err(e) = track.write_sample(&sample).await {
+                warn!("Failed to write video sample: {}", e);
             }
         }
     });
@@ -147,6 +128,8 @@ pub async fn ensure_video_pipeline_started() -> anyhow::Result<()> {
         .map(|_| ())
 }
 
+/// Switches video output to a new session's track. The keyframe for the new
+/// viewer is requested once its connection is fully up (see webrtc.rs).
 pub async fn attach_webrtc_sink(track: Arc<TrackLocalStaticSample>) {
     *VIDEO_SINK.write().await = Some(track);
 }
@@ -163,42 +146,15 @@ pub async fn detach_webrtc_sink() {
 async fn audio_frame_writer(mut rx: mpsc::Receiver<AudioFrame>) {
     const AUDIO_FRAME_DURATION: Duration = Duration::from_millis(10);
 
-    let mut current_track: Option<Arc<TrackLocalStaticSample>> = None;
-    let mut track_check_interval = interval(Duration::from_secs(1));
-    track_check_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
-
-    loop {
-        tokio::select! {
-            _ = track_check_interval.tick() => {
-                current_track = AUDIO_SINK.read().await.clone();
-            }
-            frame = rx.recv() => {
-                match frame {
-                    Some(data) => {
-                        if !is_audio_enabled() {
-                            continue;
-                        }
-                        if let Some(track) = &current_track {
-                            let sample = Sample {
-                                data,
-                                duration: AUDIO_FRAME_DURATION,
-                                ..Default::default()
-                            };
-                            if let Err(e) = track.write_sample(&sample).await {
-                                warn!("error writing audio sample: {}", e);
-                                current_track = AUDIO_SINK.read().await.clone();
-                            }
-                        } else {
-                            current_track = AUDIO_SINK.read().await.clone();
-                        }
-                    }
-                    None => {
-                        tokio::time::sleep(Duration::from_millis(500)).await;
-                        current_track = AUDIO_SINK.read().await.clone();
-                        continue;
-                    }
-                }
-            }
+    // Read the sink per frame for the same reason as the video writer.
+    while let Some(data) = rx.recv().await {
+        if !is_audio_enabled() {
+            continue;
+        }
+        let Some(track) = AUDIO_SINK.read().await.clone() else { continue };
+        let sample = Sample { data, duration: AUDIO_FRAME_DURATION, ..Default::default() };
+        if let Err(e) = track.write_sample(&sample).await {
+            warn!("error writing audio sample: {}", e);
         }
     }
 }
@@ -361,6 +317,13 @@ pub async fn force_video_keyframe() {
     }
 }
 
+/// Uncoalesced keyframe request for a viewer whose connection just came up.
+pub async fn force_new_viewer_keyframe() {
+    if let Some(manager) = PIPELINE_MANAGER.lock().await.as_ref() {
+        manager.force_keyframe_now();
+    }
+}
+
 pub async fn update_video_quality(quality: f32) -> anyhow::Result<()> {
     if !(0.0..=1.0).contains(&quality) {
         anyhow::bail!("Quality must be between 0.0 and 1.0");

app/src/webrtc.rs

@@ -13,6 +13,7 @@ use webrtc::ice_transport::ice_candidate::RTCIceCandidate;
 use webrtc::ice_transport::ice_connection_state::RTCIceConnectionState;
 use webrtc::interceptor::registry::Registry;
 use webrtc::peer_connection::configuration::RTCConfiguration;
+use webrtc::peer_connection::peer_connection_state::RTCPeerConnectionState;
 use webrtc::rtp_transceiver::rtp_codec::{
     RTCRtpCodecCapability, RTCRtpCodecParameters, RTPCodecType,
 };
@@ -309,6 +310,23 @@ impl WebRTCApi {
             },
         ));
 
+        // A keyframe requested at ICE "connected" is encoded before DTLS has
+        // finished, so it never reaches the browser and the viewer waits for
+        // the next periodic IDR. Ask again once the whole connection is up.
+        let session_id_clone = session.id.clone();
+        peer_connection.on_peer_connection_state_change(Box::new(
+            move |state: RTCPeerConnectionState| {
+                let session_id = session_id_clone.clone();
+                Box::pin(async move {
+                    if state != RTCPeerConnectionState::Connected {
+                        return;
+                    }
+                    video::force_new_viewer_keyframe().await;
+                    info!("forced keyframe for new session {}", session_id);
+                })
+            },
+        ));
+
         let session_id_clone = session.id.clone();
         let app_state_clone = config.app_state.clone();
         peer_connection.on_ice_candidate(Box::new(move |candidate: Option<RTCIceCandidate>| {