diff --git a/crates/enc-ffmpeg/src/base.rs b/crates/enc-ffmpeg/src/base.rs index 41b34700dce..a108fabfdb4 100644 --- a/crates/enc-ffmpeg/src/base.rs +++ b/crates/enc-ffmpeg/src/base.rs @@ -98,6 +98,9 @@ pub struct EncoderBase { /// capture gap that lands on a fragment cut, playing post-gap content /// during the gap. The bool records whether the duration was synthesized. held_packet: Option<(Packet, bool)>, + /// Pts (encoder ticks) and key flag of each encoded packet, for + /// recording encoders that place keyframes on capture time. + packet_keys: Option>, } impl EncoderBase { @@ -110,9 +113,21 @@ impl EncoderBase { stream_index, last_written_dts: None, held_packet: None, + packet_keys: None, } } + pub(crate) fn record_packet_keys(&mut self) { + self.packet_keys = Some(Vec::new()); + } + + pub(crate) fn take_packet_keys(&mut self) -> Vec<(i64, bool)> { + self.packet_keys + .as_mut() + .map(std::mem::take) + .unwrap_or_default() + } + pub fn set_packet_stats(&mut self, stats: Arc) { self.packet_stats = Some(stats); } @@ -202,6 +217,9 @@ impl EncoderBase { encoder: &mut encoder::encoder::Encoder, ) -> Result<(), ffmpeg::Error> { while packet_available(encoder.receive_packet(&mut self.packet))? { + if let (Some(keys), Some(pts)) = (&mut self.packet_keys, self.packet.pts()) { + keys.push((pts, self.packet.is_key())); + } self.packet.set_stream(self.stream_index); self.packet.rescale_ts( encoder.time_base(), diff --git a/crates/enc-ffmpeg/src/video/h264.rs b/crates/enc-ffmpeg/src/video/h264.rs index eaeac3a9815..e7c4dc720b9 100644 --- a/crates/enc-ffmpeg/src/video/h264.rs +++ b/crates/enc-ffmpeg/src/video/h264.rs @@ -13,6 +13,7 @@ use tracing::{debug, error, info, trace, warn}; use crate::base::EncoderBase; use crate::video::h264_packet::H264PacketEncoder; +use crate::video::keyframe_cadence::KeyframeCadence; #[cfg(target_os = "macos")] use crate::video::videotoolbox_hw::VideoToolboxHwFrames; use std::sync::Arc; @@ -236,7 +237,7 @@ impl H264EncoderBuilder { ); } } - return Ok(encoder); + return Ok(encoder.with_recording_keyframes(self.is_export)); } Err(err) => { debug!("Encoder {} init failed: {:?}", codec_name, err); @@ -298,6 +299,7 @@ impl H264EncoderBuilder { input_width, input_height, converted_frame_pool, + keyframes: None, #[cfg(target_os = "macos")] videotoolbox_hw: None, }) @@ -421,6 +423,7 @@ impl H264EncoderBuilder { "Selected hardware H264 encoder with zero-copy VideoToolbox input" ); + let is_export = self.is_export; Ok(H264Encoder { base: EncoderBase::new(stream_index), encoder, @@ -432,8 +435,10 @@ impl H264EncoderBuilder { input_width: input_config.width, input_height: input_config.height, converted_frame_pool: None, + keyframes: None, videotoolbox_hw: Some(Arc::new(hw_frames)), - }) + } + .with_recording_keyframes(is_export)) } pub fn build_standalone(self) -> Result { @@ -513,7 +518,7 @@ impl H264EncoderBuilder { fps = fps, "Selected standalone H264 encoder" ); - return Ok(encoder); + return Ok(encoder.with_recording_keyframes(self.is_export)); } Err(err) => { debug!("Standalone encoder {} init failed: {:?}", codec_name, err); @@ -821,6 +826,7 @@ pub struct H264Encoder { input_width: u32, input_height: u32, converted_frame_pool: Option, + keyframes: Option, #[cfg(target_os = "macos")] videotoolbox_hw: Option>, } @@ -846,6 +852,30 @@ pub enum QueueFrameError { impl H264Encoder { const TIME_BASE: i32 = 90000; + /// Recording encoders force a keyframe every + /// [`DEFAULT_KEYFRAME_INTERVAL_SECS`] of capture time (see + /// [`KeyframeCadence`]); exports keep the encoder's own cadence. + fn with_recording_keyframes(mut self, is_export: bool) -> Self { + if !is_export { + let gop_size = unsafe { (*self.encoder.as_ptr()).gop_size }; + self.keyframes = Some(KeyframeCadence::new( + Duration::from_secs(u64::from(DEFAULT_KEYFRAME_INTERVAL_SECS)), + self.encoder.time_base(), + gop_size, + )); + self.base.record_packet_keys(); + } + self + } + + fn observe_packets(&mut self) { + if let Some(keyframes) = &mut self.keyframes { + for (pts, is_key) in self.base.take_packet_keys() { + keyframes.observe(pts, is_key); + } + } + } + pub fn builder(input_config: VideoInfo) -> H264EncoderBuilder { H264EncoderBuilder::new(input_config) } @@ -902,6 +932,7 @@ impl H264Encoder { ) -> Result<(), QueueFrameError> { self.base .update_pts(&mut frame, timestamp, &mut self.encoder); + let requested = frame.kind() == ffmpeg::picture::Type::I; let frame_to_send = if let Some(converter) = &mut self.converter { let pts = frame.pts(); @@ -910,14 +941,20 @@ impl H264Encoder { .run(&frame, converted) .map_err(QueueFrameError::Converter)?; converted.set_pts(pts); - converted as &frame::Video + converted } else { - &frame + &mut frame }; - - self.base - .send_frame(frame_to_send, output, &mut self.encoder) - .map_err(QueueFrameError::Encode)?; + match &mut self.keyframes { + Some(keyframes) => keyframes.send(frame_to_send, requested, |frame| { + self.base.send_frame(frame, output, &mut self.encoder) + }), + None => self + .base + .send_frame(frame_to_send, output, &mut self.encoder), + } + .map_err(QueueFrameError::Encode)?; + self.observe_packets(); Ok(()) } @@ -930,6 +967,7 @@ impl H264Encoder { output: &mut format::context::Output, ) -> Result<(), QueueFrameError> { self.base.update_pts(frame, timestamp, &mut self.encoder); + let requested = frame.kind() == ffmpeg::picture::Type::I; let frame_to_send = if let Some(converter) = &mut self.converter { let pts = frame.pts(); @@ -940,14 +978,20 @@ impl H264Encoder { .run(frame, converted) .map_err(QueueFrameError::Converter)?; converted.set_pts(pts); - converted as &frame::Video + converted } else { - frame as &frame::Video + frame }; - - self.base - .send_frame(frame_to_send, output, &mut self.encoder) - .map_err(QueueFrameError::Encode)?; + match &mut self.keyframes { + Some(keyframes) => keyframes.send(frame_to_send, requested, |frame| { + self.base.send_frame(frame, output, &mut self.encoder) + }), + None => self + .base + .send_frame(frame_to_send, output, &mut self.encoder), + } + .map_err(QueueFrameError::Encode)?; + self.observe_packets(); Ok(()) } @@ -970,10 +1014,16 @@ impl H264Encoder { self.base .update_pts(&mut frame, timestamp, &mut self.encoder); + let requested = frame.kind() == ffmpeg::picture::Type::I; - self.base - .send_frame(&frame, output, &mut self.encoder) - .map_err(QueueFrameError::Encode)?; + match &mut self.keyframes { + Some(keyframes) => keyframes.send(&mut frame, requested, |frame| { + self.base.send_frame(frame, output, &mut self.encoder) + }), + None => self.base.send_frame(&frame, output, &mut self.encoder), + } + .map_err(QueueFrameError::Encode)?; + self.observe_packets(); Ok(()) } @@ -1254,6 +1304,7 @@ fn get_codec_and_options( options.set("rc", "vbr"); options.set("spatial-aq", "1"); options.set("temporal-aq", "1"); + options.set("forced-idr", "1"); } options.set("g", &keyframe_interval_str); } @@ -1265,6 +1316,7 @@ fn get_codec_and_options( } else { options.set("preset", "faster"); options.set("look_ahead", "1"); + options.set("forced_idr", "1"); } options.set("g", &keyframe_interval_str); } diff --git a/crates/enc-ffmpeg/src/video/h264_packet.rs b/crates/enc-ffmpeg/src/video/h264_packet.rs index fb7596f7c21..a8c19de0f42 100644 --- a/crates/enc-ffmpeg/src/video/h264_packet.rs +++ b/crates/enc-ffmpeg/src/video/h264_packet.rs @@ -8,6 +8,7 @@ use ffmpeg::{ }; use crate::video::h264::OpenedVideoEncoder; +use crate::video::keyframe_cadence::KeyframeCadence; pub struct H264PacketEncoder { encoder: encoder::Video, @@ -26,6 +27,7 @@ pub struct H264PacketEncoder { encoder_time_base: Rational, frame_rate: Rational, codec_name: String, + keyframes: Option, } #[derive(Debug, Clone)] @@ -69,9 +71,25 @@ impl H264PacketEncoder { encoder_time_base, frame_rate, codec_name, + keyframes: None, } } + /// See `H264Encoder::with_recording_keyframes`. + pub(crate) fn with_recording_keyframes(mut self, is_export: bool) -> Self { + if !is_export { + let gop_size = unsafe { (*self.encoder.as_ptr()).gop_size }; + self.keyframes = Some(KeyframeCadence::new( + Duration::from_secs(u64::from( + crate::video::h264::DEFAULT_KEYFRAME_INTERVAL_SECS, + )), + self.encoder_time_base, + gop_size, + )); + } + self + } + pub fn codec_name(&self) -> &str { &self.codec_name } @@ -134,6 +152,7 @@ impl H264PacketEncoder { F: FnMut(EncodedPacket) -> Result<(), EncodePacketError>, { self.update_pts(&mut frame, timestamp); + let requested = frame.kind() == ffmpeg::picture::Type::I; let frame_to_send = if let Some(converter) = &mut self.converter { let pts = frame.pts(); @@ -144,14 +163,17 @@ impl H264PacketEncoder { .run(&frame, converted) .map_err(EncodePacketError::Converter)?; converted.set_pts(pts); - converted as &frame::Video + converted } else { - &frame + &mut frame }; - - self.encoder - .send_frame(frame_to_send) - .map_err(EncodePacketError::Encode)?; + match &mut self.keyframes { + Some(keyframes) => keyframes.send(frame_to_send, requested, |frame| { + self.encoder.send_frame(frame) + }), + None => self.encoder.send_frame(frame_to_send), + } + .map_err(EncodePacketError::Encode)?; self.drain_packets(on_packet) } @@ -167,6 +189,7 @@ impl H264PacketEncoder { F: FnMut(EncodedPacket) -> Result<(), EncodePacketError>, { self.update_pts(frame, timestamp); + let requested = frame.kind() == ffmpeg::picture::Type::I; let frame_to_send = if let Some(converter) = &mut self.converter { let pts = frame.pts(); @@ -177,14 +200,17 @@ impl H264PacketEncoder { .run(frame, converted) .map_err(EncodePacketError::Converter)?; converted.set_pts(pts); - converted as &frame::Video + converted } else { - frame as &frame::Video + frame }; - - self.encoder - .send_frame(frame_to_send) - .map_err(EncodePacketError::Encode)?; + match &mut self.keyframes { + Some(keyframes) => keyframes.send(frame_to_send, requested, |frame| { + self.encoder.send_frame(frame) + }), + None => self.encoder.send_frame(frame_to_send), + } + .map_err(EncodePacketError::Encode)?; self.drain_packets(on_packet) } @@ -255,6 +281,9 @@ impl H264PacketEncoder { let dts = self.packet.dts().unwrap_or(pts); let duration = self.packet.duration(); let is_keyframe = self.packet.is_key(); + if let Some(keyframes) = &mut self.keyframes { + keyframes.observe(pts, is_keyframe); + } let data = self.packet.data().map(|d| d.to_vec()).unwrap_or_default(); on_packet(EncodedPacket { @@ -330,6 +359,85 @@ mod tests { frame } + /// Keyframe times (seconds) of a recording-encoder run fed one frame + /// every `step` slots of the nominal 30 fps for `seconds`. + fn keyframe_times(codec: &'static [&'static str], step: u64, seconds: u64) -> Option> { + ffmpeg::init().ok(); + let mut encoder = H264EncoderBuilder::new(test_video_info()) + .with_preset(H264Preset::Ultrafast) + .with_encoder_priority_override(codec) + .build_standalone() + .ok()?; + if encoder.codec_name() != codec[0] { + return None; + } + let tb = encoder.time_base(); + let secs = |pts: i64| pts as f64 * f64::from(tb.numerator()) / f64::from(tb.denominator()); + let mut keys = Vec::new(); + for slot in (0..seconds * 30).step_by(step as usize) { + let frame = create_test_frame(320, 240); + encoder + .encode_frame( + frame, + Duration::from_nanos(slot * 1_000_000_000 / 30), + |pkt| { + if pkt.is_keyframe { + keys.push(secs(pkt.pts)); + } + Ok(()) + }, + ) + .expect("encode succeeds"); + } + encoder + .flush(|pkt| { + if pkt.is_keyframe { + keys.push(secs(pkt.pts)); + } + Ok(()) + }) + .expect("flush succeeds"); + Some(keys) + } + + #[test] + fn libx264_keyframes_follow_capture_time_when_frames_arrive_slowly() { + let keys = keyframe_times(&["libx264"], 6, 13).expect("libx264"); + let gaps: Vec = keys.windows(2).map(|w| w[1] - w[0]).collect(); + assert!(keys.len() >= 6, "keyframes at {keys:?}"); + assert!( + gaps.iter().all(|gap| *gap <= 2.0 + 0.2 + 1e-6), + "a 5 fps capture must still get a keyframe every 2 s plus one frame: {keys:?}" + ); + } + + #[test] + fn libx264_keyframes_at_the_nominal_rate_are_unchanged() { + let keys = keyframe_times(&["libx264"], 1, 9).expect("libx264"); + let expected: Vec = (0..5).map(|k| f64::from(k) * 2.0).collect(); + assert_eq!(keys.len(), expected.len(), "keyframes at {keys:?}"); + assert!( + keys.iter() + .zip(&expected) + .all(|(a, b)| (a - b).abs() < 1e-3), + "keyframes at {keys:?}" + ); + } + + #[cfg(target_os = "macos")] + #[test] + fn videotoolbox_keyframes_follow_capture_time_when_frames_arrive_slowly() { + let Some(keys) = keyframe_times(&["h264_videotoolbox"], 30, 13) else { + return; + }; + assert!(keys.len() >= 6, "keyframes at {keys:?}"); + let gaps: Vec = keys.windows(2).map(|w| w[1] - w[0]).collect(); + assert!( + gaps.iter().all(|gap| *gap <= 2.0 + 1.0 + 1e-6), + "a 1 fps capture must still get a keyframe every 2 s plus one frame: {keys:?}" + ); + } + #[test] fn standalone_builder_opens_encoder_and_produces_extradata() { ffmpeg::init().ok(); diff --git a/crates/enc-ffmpeg/src/video/keyframe_cadence.rs b/crates/enc-ffmpeg/src/video/keyframe_cadence.rs new file mode 100644 index 00000000000..3ba2f56f673 --- /dev/null +++ b/crates/enc-ffmpeg/src/video/keyframe_cadence.rs @@ -0,0 +1,180 @@ +//! Keyframes on a capture-time cadence for recording encoders. +//! +//! Every recording encoder is configured with a keyframe interval in frames +//! (`g` = 2 s at the nominal rate), and the fragmented writers can only cut a +//! segment at a keyframe. Screen capture delivers frames only when the screen +//! changes, and loaded machines deliver fewer than nominal, so a 120-frame GOP +//! can span far more than 2 s: a crash then loses the whole open fragment. +//! This forces a keyframe once `interval` of capture time has passed since the +//! last one, and otherwise leaves the encoder's own cadence untouched (frames +//! at the nominal rate keep exactly the keyframes they had). + +use std::time::Duration; + +use ffmpeg::{Rational, frame, picture}; + +pub(crate) struct KeyframeCadence { + /// The interval in pts ticks. + interval: i64, + /// The encoder's own keyframe interval in frames (`gop_size`), if any. + gop: Option, + /// The newest keyframe that was requested, forced or seen in the output. + anchor: Option, + /// A frame the encoder should make a keyframe by its own count, until its + /// packet confirms or refutes it. Encoders may place keyframes on their + /// own (scene cuts), so a prediction never counts once refuted. + predicted: Option, + frames_since_key: u32, +} + +impl KeyframeCadence { + pub(crate) fn new(interval: Duration, time_base: Rational, gop_size: i32) -> Self { + let ticks = interval.as_secs_f64() * f64::from(time_base.denominator()) + / f64::from(time_base.numerator().max(1)); + Self { + interval: (ticks.round() as i64).max(1), + gop: u32::try_from(gop_size).ok().filter(|gop| *gop > 0), + anchor: None, + predicted: None, + frames_since_key: 0, + } + } + + /// Decides whether the frame with `pts` must be sent as a keyframe: + /// when the caller asked for one (`requested`), or when `interval` of + /// capture time passed since the last keyframe. + pub(crate) fn keyframe(&mut self, pts: i64, requested: bool) -> bool { + let reference = match (self.anchor, self.predicted) { + (None, None) => { + self.anchor = Some(pts); + self.frames_since_key = 1; + return requested; + } + (anchor, predicted) => anchor.max(predicted).unwrap_or(pts), + }; + let natural = self.gop.is_some_and(|gop| self.frames_since_key >= gop); + let due = pts - reference >= self.interval; + let keyframe = requested || due; + if keyframe { + self.anchor = Some(pts); + self.predicted = None; + self.frames_since_key = 0; + } else if natural { + self.predicted = Some(pts); + self.frames_since_key = 0; + } + self.frames_since_key += 1; + keyframe + } + + /// Feeds back an encoded packet (pts in the same ticks as the frames). + pub(crate) fn observe(&mut self, pts: i64, is_key: bool) { + if is_key { + self.anchor = Some(self.anchor.map_or(pts, |anchor| anchor.max(pts))); + } + if self.predicted.is_some_and(|predicted| pts >= predicted) { + if is_key && self.predicted == Some(pts) { + self.anchor = Some(self.anchor.map_or(pts, |anchor| anchor.max(pts))); + } + self.predicted = None; + } + } + + /// Sends `frame` as a keyframe when due (see [`Self::keyframe`]) and + /// restores its picture type afterwards, so a pooled or reused frame + /// never carries a forced type into the next send. + pub(crate) fn send( + &mut self, + frame: &mut frame::Video, + requested: bool, + send: impl FnOnce(&frame::Video) -> T, + ) -> T { + let Some(pts) = frame.pts() else { + return send(frame); + }; + if !self.keyframe(pts, requested) || frame.kind() == picture::Type::I { + return send(frame); + } + let kind = frame.kind(); + frame.set_kind(picture::Type::I); + let result = send(frame); + frame.set_kind(kind); + result + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cadence() -> KeyframeCadence { + KeyframeCadence::new(Duration::from_secs(2), Rational(1, 30), 60) + } + + /// Feeds frames and echoes each one back as a packet `lag` frames later, + /// a keyframe when forced or when the encoder's own count says so. + fn run(cadence: &mut KeyframeCadence, pts: &[i64], encoder_keys: &[i64]) -> Vec { + let mut forced = Vec::new(); + for &p in pts { + let key = cadence.keyframe(p, false); + if key { + forced.push(p); + } + cadence.observe(p, key || encoder_keys.contains(&p)); + } + forced + } + + #[test] + fn frames_at_the_nominal_rate_are_keyframes_exactly_where_the_encoder_puts_them() { + let pts: Vec = (0..600).collect(); + let keys: Vec = (0..600).step_by(60).collect(); + let forced = run(&mut cadence(), &pts, &keys); + assert!(forced.iter().all(|p| keys.contains(p)), "{forced:?}"); + } + + #[test] + fn slow_frames_get_a_keyframe_every_interval_of_capture_time() { + let pts: Vec = (0..60).map(|k| k * 10).collect(); + assert_eq!( + run(&mut cadence(), &pts, &[0]), + vec![60, 120, 180, 240, 300, 360, 420, 480, 540] + ); + } + + #[test] + fn a_prediction_the_encoder_did_not_honour_does_not_reset_the_timer() { + let mut cadence = KeyframeCadence::new(Duration::from_secs(2), Rational(1, 300), 60); + let mut pts: Vec = (0..=60).map(|k| k * 9).collect(); + pts.extend((1..10).map(|k| 540 + k * 90)); + let forced = run(&mut cadence, &pts, &[0]); + assert_eq!(forced.first(), Some(&630), "{forced:?}"); + } + + #[test] + fn an_encoder_keyframe_seen_in_the_output_restarts_the_interval() { + let pts: Vec = (0..60).map(|k| k * 10).collect(); + assert_eq!(run(&mut cadence(), &pts, &[0, 30])[0], 90); + } + + #[test] + fn requested_keyframes_are_kept_and_restart_the_interval() { + let mut cadence = cadence(); + assert!(!cadence.keyframe(0, false)); + assert!(cadence.keyframe(6, true)); + assert!(!cadence.keyframe(60, false)); + assert!(cadence.keyframe(66, false)); + } + + #[test] + fn a_forced_picture_type_does_not_outlive_its_send() { + let mut cadence = cadence(); + let mut frame = frame::Video::new(ffmpeg::format::Pixel::YUV420P, 16, 16); + frame.set_pts(Some(0)); + cadence.send(&mut frame, false, |_| ()); + frame.set_pts(Some(90)); + let sent = cadence.send(&mut frame, false, |frame| frame.kind()); + assert_eq!(sent, picture::Type::I); + assert_eq!(frame.kind(), picture::Type::None); + } +} diff --git a/crates/enc-ffmpeg/src/video/mod.rs b/crates/enc-ffmpeg/src/video/mod.rs index 104261fc075..12e4688e7ec 100644 --- a/crates/enc-ffmpeg/src/video/mod.rs +++ b/crates/enc-ffmpeg/src/video/mod.rs @@ -2,6 +2,7 @@ pub mod colour; pub mod h264; pub mod h264_packet; pub mod hevc; +pub(crate) mod keyframe_cadence; pub mod prores; #[cfg(target_os = "macos")] pub mod videotoolbox_hw;