diff --git a/src/avhw.rs b/src/avhw.rs index 284887e..446c175 100644 --- a/src/avhw.rs +++ b/src/avhw.rs @@ -25,6 +25,9 @@ use crate::transform::{transpose_if_transform_transposed, Transform}; pub enum BitrateCommand { UpdateBitrate { target_bps: u64 }, UpdateResolution { width: u32, height: u32 }, + /// Force the next encoded frame to be an IDR. Sent by the WebRTC thread + /// in response to str0m `Event::KeyframeRequest` or a resolution change. + ForceKeyframe, } #[derive(Clone, Copy, Debug)] @@ -856,6 +859,7 @@ impl SwEncImport { requested = Some((width & !1, height & !1)); } BitrateCommand::UpdateBitrate { .. } => {} + BitrateCommand::ForceKeyframe => {} } } @@ -966,6 +970,10 @@ pub struct SwEncEncode { fps: u32, bitrate: u64, gop_size: u32, + /// Set true when WebRTC requests a keyframe. Forces the next frame to + /// `AV_PICTURE_TYPE_I` and bypasses the dedup hash check. Cleared only + /// after `avcodec_send_frame` accepts the forced frame. + force_keyframe_pending: bool, } const FNV1A_OFFSET_BASIS: u64 = 0xcbf29ce484222325; @@ -1028,6 +1036,7 @@ impl SwEncEncode { fps, bitrate, gop_size, + force_keyframe_pending: false, }) } @@ -1066,6 +1075,7 @@ impl SwEncEncode { fps, bitrate, gop_size, + force_keyframe_pending: false, }) } @@ -1103,9 +1113,15 @@ impl SwEncEncode { } } BitrateCommand::UpdateResolution { .. } => {} + BitrateCommand::ForceKeyframe => { + self.force_keyframe_pending = true; + tracing::debug!("encode thread: ForceKeyframe requested"); + } } } + let force_this_frame = self.force_keyframe_pending; + while let Ok(change) = self.resolution_rx.try_recv() { self.recreate_encoder(change.width, change.height)?; } @@ -1130,7 +1146,7 @@ impl SwEncEncode { self.frame_count = self.frame_count.saturating_add(1); let current_hash = hash_sampled_y_plane(&frame.y_data, width, height, frame.y_stride); let force_gop_frame = self.gop_size > 0 && frame_index % u64::from(self.gop_size) == 0; - if frame_index > 0 && !force_gop_frame && current_hash == self.last_frame_hash { + if frame_index > 0 && !force_gop_frame && !force_this_frame && current_hash == self.last_frame_hash { tracing::debug!(frame_index, "skipping duplicate frame"); self.last_frame_hash = current_hash; return Ok(()); @@ -1172,8 +1188,16 @@ impl SwEncEncode { let start_ts = self.starting_timestamp.unwrap_or(0); // SAFETY: yuv_frame is initialized, writable, and matches the opened encoder format. + // pict_type is reset every frame: the AVFrame is reused, so without resetting to NONE + // a previously-forced I-type would leak into subsequent P-frames. With forced-idr=1 + // set on the encoder, AV_PICTURE_TYPE_I produces a true IDR NALU. unsafe { (*self.yuv_frame).pts = pts; + (*self.yuv_frame).pict_type = if force_this_frame { + ffi::AVPictureType::AV_PICTURE_TYPE_I + } else { + ffi::AVPictureType::AV_PICTURE_TYPE_NONE + }; let ret = ffi::avcodec_send_frame(self.enc_video.as_mut_ptr(), self.yuv_frame); if ret < 0 { bail!( @@ -1183,6 +1207,10 @@ impl SwEncEncode { } } + if force_this_frame { + self.force_keyframe_pending = false; + } + self.drain_encoder(start_ts) } @@ -1215,6 +1243,7 @@ impl SwEncEncode { self.enc_height = height; self.last_frame_hash = 0; self.frame_count = 0; + self.force_keyframe_pending = true; Ok(()) } @@ -1741,6 +1770,14 @@ fn create_software_h264_encoder( // avcodec_alloc_context3. Setting level is a simple i32 field // assignment on a properly aligned struct. (*enc.as_mut_ptr()).level = 42; // H.264 Level 4.2 (up to 1440p@30) + // SAFETY: priv_data belongs to the unopened libx264 encoder context. + // `forced-idr` is an FFmpeg-level private option (not x264-native), + // so it must be set via av_opt_set, NOT via the x264opts string. + // With forced-idr=1, setting AV_PICTURE_TYPE_I on an input frame + // produces a true IDR NALU with inline SPS/PPS (repeat_headers=1). + let key = CString::new("forced-idr").unwrap(); + let val = CString::new("1").unwrap(); + ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0); let key = CString::new("x264opts").unwrap(); let vbv_maxrate = bitrate; let vbv_bufsize = bitrate / 4; diff --git a/src/state_portal.rs b/src/state_portal.rs index deefede..4b0a61d 100644 --- a/src/state_portal.rs +++ b/src/state_portal.rs @@ -702,6 +702,10 @@ fn webrtc_thread_loop( break; } + if wrtc.take_force_keyframe() { + let _ = bitrate_tx.try_send(BitrateCommand::ForceKeyframe); + } + let connected = wrtc.is_connected(); let was_paused = paused.load(Ordering::Relaxed); let now_paused = !connected; diff --git a/src/webrtc.rs b/src/webrtc.rs index c072318..d70b516 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -200,6 +200,7 @@ struct WebRtcInner { video_pt: Option, connected: bool, need_keyframe: bool, + force_keyframe_to_encode: bool, current_bwe_estimate: Option, rtp_clock: u32, buf: Vec, @@ -353,6 +354,17 @@ impl WebRtcState { pub fn set_need_keyframe(&mut self) { if let Some(inner) = self.inner.as_mut() { inner.need_keyframe = true; + inner.force_keyframe_to_encode = true; + } + } + + pub fn take_force_keyframe(&mut self) -> bool { + if let Some(inner) = self.inner.as_mut() { + let v = inner.force_keyframe_to_encode; + inner.force_keyframe_to_encode = false; + v + } else { + false } } } @@ -429,6 +441,7 @@ impl WebRtcInner { video_pt: None, connected: false, need_keyframe: false, + force_keyframe_to_encode: false, current_bwe_estimate: None, rtp_clock: 0, buf: vec![0u8; 65535], @@ -446,6 +459,7 @@ impl WebRtcInner { .map_err(|e| anyhow::anyhow!("accept_offer: {e}"))?; self.need_keyframe = true; + self.force_keyframe_to_encode = true; tracing::info!("SDP exchange complete, waiting for ICE/DTLS..."); self.discover_video_params(); @@ -503,6 +517,7 @@ impl WebRtcInner { tracing::info!("WebRTC connected!"); self.connected = true; self.need_keyframe = true; + self.force_keyframe_to_encode = true; self.discover_video_params(); } Event::IceConnectionStateChange(IceConnectionState::Disconnected) => { @@ -525,6 +540,7 @@ impl WebRtcInner { Event::KeyframeRequest(_) => { tracing::info!("received keyframe request from viewer"); self.need_keyframe = true; + self.force_keyframe_to_encode = true; } Event::EgressBitrateEstimate(est) => { let bitrate = match est {