diff --git a/src/bin/vrec-daemon.rs b/src/bin/vrec-daemon.rs index cfe001c..504aa42 100644 --- a/src/bin/vrec-daemon.rs +++ b/src/bin/vrec-daemon.rs @@ -177,10 +177,15 @@ async fn main() -> Result<(), Box> { }; let audio_codec_ctx_ptr = audio_encoder.as_ref().map(|e| e.codec_ctx() as usize); + let video_time_base = unsafe { (*(codec_ctx_ptr as *mut ffmpeg_next::ffi::AVCodecContext)).time_base }; + let audio_time_base = audio_codec_ctx_ptr.map(|p| unsafe { (*(p as *mut ffmpeg_next::ffi::AVCodecContext)).time_base }); + let mut ring = HeapRb::::new((config.replay_duration_sec * 120).max(120) as usize); let mut normal_muxer: Option = None; let mut normal_recording = false; let mut normal_waiting_keyframe = false; + let mut rec_base_video_pts: i64 = 0; + let mut rec_base_audio_pts: i64 = 0; thread::spawn(move || { while let Ok(drain) = mux_rx.recv() { @@ -188,25 +193,73 @@ async fn main() -> Result<(), Box> { let filename = format!("replay_{}.mp4", ts); let full_path = vrec::config::VrecConfig::load().resolve_save_path(&filename); println!("Saving replay to {}...", full_path); - let mut start_idx = 0; + + // Find the first video IDR keyframe (ignore audio packets which always have keyframe flag set) + let start_video_idx = drain.iter().position(|p| p.stream_index() == 0 && p.is_keyframe()); + let first_v_idx = match start_video_idx { + Some(idx) => idx, + None => { + println!("No video keyframe found in buffer, skipping save."); + continue; + } + }; + + let first_video_pts = drain[first_v_idx].pts(); + let first_audio_pts = if let Some(a_tb) = audio_time_base { + unsafe { ffmpeg_next::ffi::av_rescale_q(first_video_pts, video_time_base, a_tb) } + } else { + 0 + }; + + let mut prepared: Vec<(i64, Packet)> = Vec::with_capacity(drain.len()); for (i, p) in drain.iter().enumerate() { - if p.is_keyframe() { - start_idx = i; - break; + if p.stream_index() == 0 { + if i >= first_v_idx { + let rebased = p.rebased(first_video_pts); + let time_us = unsafe { + ffmpeg_next::ffi::av_rescale_q( + rebased.pts(), + video_time_base, + ffmpeg_next::ffi::AVRational { num: 1, den: 1_000_000 }, + ) + }; + prepared.push((time_us, rebased)); + } + } else if p.stream_index() == 1 { + if p.pts() >= first_audio_pts { + let rebased = p.rebased(first_audio_pts); + let time_us = if let Some(a_tb) = audio_time_base { + unsafe { + ffmpeg_next::ffi::av_rescale_q( + rebased.pts(), + a_tb, + ffmpeg_next::ffi::AVRational { num: 1, den: 1_000_000 }, + ) + } + } else { + 0 + }; + prepared.push((time_us, rebased)); + } } } - - if start_idx < drain.len() { - let codec_ctx = codec_ctx_ptr as *mut ffmpeg_next::ffi::AVCodecContext; - let audio_codec_ctx = audio_codec_ctx_ptr.map(|p| p as *mut ffmpeg_next::ffi::AVCodecContext); - let mut muxer = unsafe { Muxer::new(&full_path, codec_ctx, audio_codec_ctx).unwrap() }; - for p in drain.into_iter().skip(start_idx) { - let _ = muxer.write_packet(&p); + + // Sort packets by presentation timestamp so muxer receives strictly chronological stream + prepared.sort_by_key(|(t, _)| *t); + + let codec_ctx = codec_ctx_ptr as *mut ffmpeg_next::ffi::AVCodecContext; + let audio_codec_ctx = audio_codec_ctx_ptr.map(|p| p as *mut ffmpeg_next::ffi::AVCodecContext); + match unsafe { Muxer::new(&full_path, codec_ctx, audio_codec_ctx) } { + Ok(mut muxer) => { + for (_, p) in prepared { + let _ = muxer.write_packet(&p); + } + let _ = muxer.finalize(); + println!("Replay saved to {}!", full_path); + } + Err(e) => { + eprintln!("Failed to create muxer for {}: {}", full_path, e); } - let _ = muxer.finalize(); - println!("Replay saved to {}!", full_path); - } else { - println!("No keyframe found in buffer, skipping save."); } } }); @@ -251,6 +304,8 @@ async fn main() -> Result<(), Box> { } normal_recording = false; normal_waiting_keyframe = false; + rec_base_video_pts = 0; + rec_base_audio_pts = 0; rec_state_clone.store(false, Ordering::SeqCst); rec_start_clone.store(0, Ordering::SeqCst); }, @@ -262,6 +317,8 @@ async fn main() -> Result<(), Box> { } normal_recording = false; normal_waiting_keyframe = false; + rec_base_video_pts = 0; + rec_base_audio_pts = 0; rec_state_clone.store(false, Ordering::SeqCst); rec_start_clone.store(0, Ordering::SeqCst); } else { @@ -303,7 +360,10 @@ async fn main() -> Result<(), Box> { } if normal_recording && !normal_waiting_keyframe && let Some(muxer) = normal_muxer.as_mut() { - let _ = muxer.write_packet(&pkt); + if pkt.pts() >= rec_base_audio_pts { + let rebased = pkt.rebased(rec_base_audio_pts); + let _ = muxer.write_packet(&rebased); + } } } } @@ -318,6 +378,12 @@ async fn main() -> Result<(), Box> { if normal_recording { if normal_waiting_keyframe && pkt.is_keyframe() { normal_waiting_keyframe = false; + rec_base_video_pts = pkt.pts(); + rec_base_audio_pts = if let Some(a_tb) = audio_time_base { + unsafe { ffmpeg_next::ffi::av_rescale_q(rec_base_video_pts, video_time_base, a_tb) } + } else { + 0 + }; let codec_ctx = codec_ctx_ptr as *mut ffmpeg_next::ffi::AVCodecContext; let ts = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); let filename = format!("record_{}.mp4", ts); @@ -329,13 +395,14 @@ async fn main() -> Result<(), Box> { if !normal_waiting_keyframe && let Some(muxer) = normal_muxer.as_mut() { - let _ = muxer.write_packet(&pkt); + let rebased = pkt.rebased(rec_base_video_pts); + let _ = muxer.write_packet(&rebased); + } } } } } } -} } }); diff --git a/src/encoder.rs b/src/encoder.rs index 79f3bbe..9de26ae 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -62,6 +62,7 @@ impl VaapiEncoder { (*codec_ctx).qmin = 16; (*codec_ctx).qmax = 28; // Prevent compression blockiness (*codec_ctx).profile = FF_PROFILE_H264_HIGH; + (*codec_ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; let mut hw_device_ctx: *mut AVBufferRef = ptr::null_mut(); let ret = av_hwdevice_ctx_create( @@ -434,6 +435,7 @@ impl WindowsHwEncoder { (*codec_ctx).qmin = 16; (*codec_ctx).qmax = 28; (*codec_ctx).profile = FF_PROFILE_H264_HIGH; + (*codec_ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; if selected_desc.contains("NVENC") { (*codec_ctx).pix_fmt = AVPixelFormat::AV_PIX_FMT_NV12; @@ -594,6 +596,7 @@ impl AudioEncoder { (*codec_ctx).time_base = AVRational { num: 1, den: sample_rate }; (*codec_ctx).sample_fmt = AVSampleFormat::AV_SAMPLE_FMT_FLTP; (*codec_ctx).bit_rate = 192_000; + (*codec_ctx).flags |= AV_CODEC_FLAG_GLOBAL_HEADER as libc::c_int; av_channel_layout_default(&mut (*codec_ctx).ch_layout, channels); let ret = avcodec_open2(codec_ctx, codec, ptr::null_mut()); @@ -721,3 +724,40 @@ impl Drop for AudioEncoder { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_audio_encoder_global_header() { + let enc = AudioEncoder::new(48000, 2).expect("AudioEncoder creation"); + unsafe { + let flags = (*enc.codec_ctx).flags; + assert!((flags & (AV_CODEC_FLAG_GLOBAL_HEADER as i32)) != 0, "Audio encoder must have GLOBAL_HEADER flag"); + assert!((*enc.codec_ctx).extradata_size > 0, "Audio encoder extradata must be populated"); + } + } + + #[test] + fn test_vaapi_encoder_global_header() { + if let Ok(mut enc) = VaapiEncoder::new(1280, 720) { + unsafe { + let flags = (*enc.codec_ctx).flags; + assert!((flags & (AV_CODEC_FLAG_GLOBAL_HEADER as i32)) != 0, "Vaapi encoder must have GLOBAL_HEADER flag"); + } + // Encode a blank BGRA frame + let blank_frame = Frame::Raw { + width: 1280, + height: 720, + stride: 1280 * 4, + data: vec![0u8; 1280 * 720 * 4], + }; + let pkts = enc.encode_frame(&blank_frame).expect("Encode frame"); + println!("Encoded {} packets from first frame", pkts.len()); + unsafe { + println!("VAAPI extradata_size after encode: {}", (*enc.codec_ctx).extradata_size); + } + } + } +} diff --git a/src/muxer.rs b/src/muxer.rs index fa3ca38..6171555 100644 --- a/src/muxer.rs +++ b/src/muxer.rs @@ -70,6 +70,8 @@ impl Muxer { let mut opts: *mut AVDictionary = std::ptr::null_mut(); av_dict_set(&mut opts, c"movflags".as_ptr(), c"+faststart".as_ptr(), 0); + (*fmt_ctx).avoid_negative_ts = 2; // AVFMT_AVOID_NEG_TS_MAKE_ZERO: Shift timestamps so that they start at 0 + (*fmt_ctx).max_interleave_delta = 1_000_000; let ret = avformat_write_header(fmt_ctx, &mut opts); av_dict_free(&mut opts); if ret < 0 { @@ -88,6 +90,14 @@ impl Muxer { } } + pub fn video_time_base(&self) -> AVRational { + self.video_time_base + } + + pub fn audio_time_base(&self) -> Option { + self.audio_time_base + } + pub fn write_packet(&mut self, packet: &Packet) -> Result<(), String> { unsafe { if packet.ptr.is_null() { diff --git a/src/ring.rs b/src/ring.rs index 8f96b35..4bfe268 100644 --- a/src/ring.rs +++ b/src/ring.rs @@ -25,6 +25,35 @@ impl Packet { ((*self.ptr).flags & AV_PKT_FLAG_KEY) != 0 } } + + pub fn pts(&self) -> i64 { + unsafe { (*self.ptr).pts } + } + + pub fn dts(&self) -> i64 { + unsafe { (*self.ptr).dts } + } + + pub fn set_pts(&mut self, pts: i64) { + unsafe { (*self.ptr).pts = pts; } + } + + pub fn set_dts(&mut self, dts: i64) { + unsafe { (*self.ptr).dts = dts; } + } + + pub fn rebased(&self, base_pts: i64) -> Self { + let cloned = self.clone(); + unsafe { + if (*cloned.ptr).pts != AV_NOPTS_VALUE { + (*cloned.ptr).pts -= base_pts; + } + if (*cloned.ptr).dts != AV_NOPTS_VALUE { + (*cloned.ptr).dts -= base_pts; + } + } + cloned + } } impl Drop for Packet {