diff --git a/src/bin/vrec-daemon.rs b/src/bin/vrec-daemon.rs index 504aa42..31546c2 100644 --- a/src/bin/vrec-daemon.rs +++ b/src/bin/vrec-daemon.rs @@ -46,6 +46,7 @@ fn ensure_wayland_env() { } #[tokio::main] +#[allow(unused_assignments)] async fn main() -> Result<(), Box> { ensure_wayland_env(); @@ -121,10 +122,10 @@ async fn main() -> Result<(), Box> { let (frame_tx, frame_rx) = bounded::(5); let (cmd_tx, cmd_rx) = bounded::(10); let (mux_tx, mux_rx) = bounded::>(1); - let (audio_tx, audio_rx) = bounded::>(100); + let (audio_tx, audio_rx) = bounded::>(500); let initial_config = vrec::config::VrecConfig::load(); - let audio_info = vrec::capture::audio::AudioCapture::new_with_device_and_mode( - audio_tx, + let audio_capture_initial = vrec::capture::audio::AudioCapture::new_with_device_and_mode( + audio_tx.clone(), Some(&initial_config.audio_device), &initial_config.audio_mode, ).ok(); @@ -145,9 +146,7 @@ async fn main() -> Result<(), Box> { loop { match source.next_frame() { Ok(frame) => { - if capture_tx.send(frame).is_err() { - break; - } + let _ = capture_tx.try_send(frame); } Err(e) => { eprintln!("Capture error: {}", e); @@ -161,6 +160,10 @@ async fn main() -> Result<(), Box> { let rec_start_clone = Arc::clone(&record_start_state); let replay_state_clone = Arc::clone(&replay_enabled_state); let audio_muted_clone = Arc::clone(&audio_muted_state); + let audio_tx_clone = audio_tx.clone(); + let audio_info = audio_capture_initial.as_ref().map(|(_, sr, ch)| (*sr, *ch)); + #[allow(unused_variables, unused_assignments)] + let mut audio_capture = audio_capture_initial.map(|(c, _, _)| c); thread::spawn(move || { let mut config = vrec::config::VrecConfig::load(); @@ -170,10 +173,10 @@ async fn main() -> Result<(), Box> { .expect("Failed to init encoder"); let codec_ctx_ptr = encoder.codec_ctx() as usize; - let mut audio_encoder = if let Some(ref info) = audio_info { - vrec::encoder::AudioEncoder::new(info.1 as i32, info.2 as i32).ok() + let mut audio_encoder = if let Some((sr, ch)) = audio_info { + vrec::encoder::AudioEncoder::new(sr as i32, ch as i32).ok() } else { - None + vrec::encoder::AudioEncoder::new(48000, 2).ok() }; let audio_codec_ctx_ptr = audio_encoder.as_ref().map(|e| e.codec_ctx() as usize); @@ -185,7 +188,7 @@ async fn main() -> Result<(), Box> { 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; + let mut rec_base_audio_pts: i64 = -1; thread::spawn(move || { while let Ok(drain) = mux_rx.recv() { @@ -205,11 +208,15 @@ async fn main() -> Result<(), Box> { }; 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 - }; + + // Find the audio packet in the ring buffer closest in arrival to the first video keyframe + let first_audio_pts = drain + .iter() + .enumerate() + .filter(|(_, p)| p.stream_index() == 1) + .min_by_key(|(i, _)| i.abs_diff(first_v_idx)) + .map(|(_, p)| p.pts()) + .unwrap_or(0); let mut prepared: Vec<(i64, Packet)> = Vec::with_capacity(drain.len()); for (i, p) in drain.iter().enumerate() { @@ -264,6 +271,10 @@ async fn main() -> Result<(), Box> { } }); + let mut ticker = crossbeam_channel::tick(std::time::Duration::from_nanos(1_000_000_000 / config.fps.max(1) as u64)); + let mut latest_frame: Option = None; + let mut has_new_frame = false; + loop { crossbeam_channel::select! { recv(cmd_rx) -> cmd_res => { @@ -272,13 +283,29 @@ async fn main() -> Result<(), Box> { Command::ReloadConfig => { config = vrec::config::VrecConfig::load(); replay_state_clone.store(config.replay_enabled, Ordering::SeqCst); - println!("Daemon config reloaded!"); + audio_muted_clone.store(config.audio_mode == "muted", Ordering::SeqCst); + println!("Daemon config reloaded! Audio mode: {}", config.audio_mode); let new_capacity = (config.replay_duration_sec * 120).max(120) as usize; if ring.capacity().get() != new_capacity { ring = HeapRb::::new(new_capacity); println!("Replay buffer resized to {} packets.", new_capacity); } vrec::hyprland_binds::register_hyprland_binds(&config); + audio_capture = match vrec::capture::audio::AudioCapture::new_with_device_and_mode( + audio_tx_clone.clone(), + Some(&config.audio_device), + &config.audio_mode, + ) { + Ok((c, _, _)) => { + println!("Audio capture reloaded for mode: {}", config.audio_mode); + Some(c) + } + Err(e) => { + eprintln!("Failed to reload audio capture: {}", e); + None + } + }; + ticker = crossbeam_channel::tick(std::time::Duration::from_nanos(1_000_000_000 / config.fps.max(1) as u64)); }, Command::SaveReplay => { if config.replay_enabled { @@ -292,6 +319,7 @@ async fn main() -> Result<(), Box> { normal_recording = true; normal_waiting_keyframe = true; rec_state_clone.store(true, Ordering::SeqCst); + rec_base_audio_pts = -1; let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); rec_start_clone.store(now, Ordering::SeqCst); println!("StartRecording requested, waiting for keyframe..."); @@ -305,7 +333,7 @@ async fn main() -> Result<(), Box> { normal_recording = false; normal_waiting_keyframe = false; rec_base_video_pts = 0; - rec_base_audio_pts = 0; + rec_base_audio_pts = -1; rec_state_clone.store(false, Ordering::SeqCst); rec_start_clone.store(0, Ordering::SeqCst); }, @@ -318,13 +346,14 @@ async fn main() -> Result<(), Box> { normal_recording = false; normal_waiting_keyframe = false; rec_base_video_pts = 0; - rec_base_audio_pts = 0; + rec_base_audio_pts = -1; rec_state_clone.store(false, Ordering::SeqCst); rec_start_clone.store(0, Ordering::SeqCst); } else { normal_recording = true; normal_waiting_keyframe = true; rec_state_clone.store(true, Ordering::SeqCst); + rec_base_audio_pts = -1; let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs(); rec_start_clone.store(now, Ordering::SeqCst); println!("ToggleRecording: StartRecording requested, waiting for keyframe..."); @@ -339,69 +368,81 @@ async fn main() -> Result<(), Box> { } } }, + recv(audio_rx) -> audio_res => { + if let Ok(mut audio_chunk) = audio_res { + if audio_muted_clone.load(Ordering::Relaxed) { + for sample in audio_chunk.iter_mut() { + *sample = 0.0; + } + } + if let Some(enc) = audio_encoder.as_mut() + && let Ok(audio_packets) = enc.encode_pcm(&audio_chunk) { + for mut pkt in audio_packets { + pkt.set_stream_index(1); + if config.replay_enabled { + ring.push_overwrite(pkt.clone()); + } + if normal_recording && !normal_waiting_keyframe + && let Some(muxer) = normal_muxer.as_mut() { + if rec_base_audio_pts < 0 { + rec_base_audio_pts = pkt.pts(); + } + if pkt.pts() >= rec_base_audio_pts { + let rebased = pkt.rebased(rec_base_audio_pts); + let _ = muxer.write_packet(&rebased); + } + } + } + } + } + }, recv(frame_rx) -> frame_res => { - let frame = match frame_res { - Ok(f) => f, - Err(_) => break, + if let Ok(f) = frame_res { + latest_frame = Some(f); + has_new_frame = true; + } + }, + recv(ticker) -> _ => { + let packets_res = if has_new_frame { + has_new_frame = false; + if let Some(ref f) = latest_frame { + encoder.encode_frame(f) + } else { + Ok(Vec::new()) + } + } else { + encoder.encode_cached_frame() }; - while let Ok(mut audio_chunk) = audio_rx.try_recv() { - if audio_muted_clone.load(Ordering::Relaxed) { - for sample in audio_chunk.iter_mut() { - *sample = 0.0; - } - } - if let Some(enc) = audio_encoder.as_mut() - && let Ok(audio_packets) = enc.encode_pcm(&audio_chunk) { - for mut pkt in audio_packets { - pkt.set_stream_index(1); + if let Ok(packets) = packets_res { + for mut pkt in packets { + pkt.set_stream_index(0); if config.replay_enabled { ring.push_overwrite(pkt.clone()); } - if normal_recording && !normal_waiting_keyframe - && let Some(muxer) = normal_muxer.as_mut() { - if pkt.pts() >= rec_base_audio_pts { - let rebased = pkt.rebased(rec_base_audio_pts); - let _ = muxer.write_packet(&rebased); - } + if normal_recording { + if normal_waiting_keyframe && pkt.is_keyframe() { + normal_waiting_keyframe = false; + rec_base_video_pts = pkt.pts(); + rec_base_audio_pts = -1; + 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); + let full_path = config.resolve_save_path(&filename); + let audio_codec_ctx = audio_codec_ctx_ptr.map(|p| p as *mut ffmpeg_next::ffi::AVCodecContext); + normal_muxer = unsafe { Muxer::new(&full_path, codec_ctx, audio_codec_ctx).ok() }; + println!("Started normal recording to {}", full_path); } - } - } - } - - if let Ok(packets) = encoder.encode_frame(&frame) { - for mut pkt in packets { - pkt.set_stream_index(0); - if config.replay_enabled { - ring.push_overwrite(pkt.clone()); - } - 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); - let full_path = config.resolve_save_path(&filename); - let audio_codec_ctx = audio_codec_ctx_ptr.map(|p| p as *mut ffmpeg_next::ffi::AVCodecContext); - normal_muxer = unsafe { Muxer::new(&full_path, codec_ctx, audio_codec_ctx).ok() }; - println!("Started normal recording to {}", full_path); - } - - if !normal_waiting_keyframe - && let Some(muxer) = normal_muxer.as_mut() { - let rebased = pkt.rebased(rec_base_video_pts); - let _ = muxer.write_packet(&rebased); + + if !normal_waiting_keyframe + && let Some(muxer) = normal_muxer.as_mut() { + let rebased = pkt.rebased(rec_base_video_pts); + let _ = muxer.write_packet(&rebased); + } + } } } } - } - } } } }); @@ -477,6 +518,7 @@ async fn main() -> Result<(), Box> { }.to_string(); let _ = cfg.save(); println!("Audio mode cycled to: {}", cfg.audio_mode); + let _ = cmd_tx.try_send(Command::ReloadConfig); }, Command::StopDaemon => { println!("StopDaemon requested: Finalizing active recordings..."); diff --git a/src/capture/audio.rs b/src/capture/audio.rs index b534dbb..67490e5 100644 --- a/src/capture/audio.rs +++ b/src/capture/audio.rs @@ -87,17 +87,38 @@ fn spawn_parec_stream( let mut stdout = child.stdout.take().ok_or("Failed to open parec stdout")?; std::thread::spawn(move || { - let mut buf = [0u8; 4096]; + let mut remainder = Vec::with_capacity(4); + let mut read_buf = [0u8; 4096]; loop { - match stdout.read(&mut buf) { + match stdout.read(&mut read_buf) { Ok(0) => break, Ok(n) => { - let floats_count = n / 4; + let total_len = remainder.len() + n; + let full_bytes = total_len - (total_len % 4); + if full_bytes == 0 { + remainder.extend_from_slice(&read_buf[..n]); + continue; + } + + let mut combined = Vec::with_capacity(total_len); + combined.extend_from_slice(&remainder); + combined.extend_from_slice(&read_buf[..n]); + + let floats_count = full_bytes / 4; let mut floats = Vec::with_capacity(floats_count); for i in 0..floats_count { - let b = [buf[i * 4], buf[i * 4 + 1], buf[i * 4 + 2], buf[i * 4 + 3]]; + let b = [ + combined[i * 4], + combined[i * 4 + 1], + combined[i * 4 + 2], + combined[i * 4 + 3], + ]; floats.push(f32::from_le_bytes(b)); } + + remainder.clear(); + remainder.extend_from_slice(&combined[full_bytes..]); + if sender.send(floats).is_err() { break; } diff --git a/src/encoder.rs b/src/encoder.rs index 9de26ae..3b19b62 100644 --- a/src/encoder.rs +++ b/src/encoder.rs @@ -7,6 +7,7 @@ pub struct VaapiEncoder { hw_device_ctx: *mut AVBufferRef, hw_frames_ctx: *mut AVBufferRef, sws_ctx: *mut SwsContext, + cached_nv12: *mut AVFrame, next_pts: i64, } @@ -113,6 +114,7 @@ impl VaapiEncoder { hw_device_ctx, hw_frames_ctx: hw_frames_ref, sws_ctx: ptr::null_mut(), + cached_nv12: ptr::null_mut(), next_pts: 0, }) } @@ -143,7 +145,7 @@ impl VaapiEncoder { dst_row.copy_from_slice(src_row); } - let mut nv12_frame = av_frame_alloc(); + let nv12_frame = av_frame_alloc(); (*nv12_frame).format = AVPixelFormat::AV_PIX_FMT_NV12 as i32; (*nv12_frame).width = *width as i32; (*nv12_frame).height = *height as i32; @@ -184,15 +186,19 @@ impl VaapiEncoder { av_packet_free(&mut pkt); } av_frame_free(&mut hw_frame); - av_frame_free(&mut nv12_frame); + if !self.cached_nv12.is_null() { + av_frame_free(&mut self.cached_nv12); + } + self.cached_nv12 = nv12_frame; av_frame_free(&mut bgra_frame); } Frame::DmaBuf { width: dma_width, height: dma_height, format, modifier, fd, stride, offset } => { + let mmap_size = (*stride as usize) * (*dma_height as usize); let mut drm_desc = AVDRMFrameDescriptor { nb_objects: 1, objects: [AVDRMObjectDescriptor { fd: *fd, - size: 0, + size: mmap_size, format_modifier: *modifier, }, std::mem::zeroed(), std::mem::zeroed(), std::mem::zeroed()], nb_layers: 1, @@ -265,7 +271,7 @@ impl VaapiEncoder { } libc::munmap(ptr, mmap_size); - let mut nv12_frame = av_frame_alloc(); + let nv12_frame = av_frame_alloc(); (*nv12_frame).format = AVPixelFormat::AV_PIX_FMT_NV12 as i32; (*nv12_frame).width = *dma_width as i32; (*nv12_frame).height = *dma_height as i32; @@ -310,7 +316,10 @@ impl VaapiEncoder { } av_frame_free(&mut fresh_hw_frame); - av_frame_free(&mut nv12_frame); + if !self.cached_nv12.is_null() { + av_frame_free(&mut self.cached_nv12); + } + self.cached_nv12 = nv12_frame; av_frame_free(&mut bgra_frame); } else { eprintln!("mmap failed on dma-buf fd {}: {}", *fd, std::io::Error::last_os_error()); @@ -334,11 +343,54 @@ impl VaapiEncoder { } Ok(packets) } + + pub fn encode_cached_frame(&mut self) -> Result, String> { + let mut packets = Vec::new(); + unsafe { + if self.cached_nv12.is_null() { + return Ok(packets); + } + + let mut fresh_hw_frame = av_frame_alloc(); + (*fresh_hw_frame).format = AVPixelFormat::AV_PIX_FMT_VAAPI as i32; + let get_buf_ret = av_hwframe_get_buffer(self.hw_frames_ctx, fresh_hw_frame, 0); + if get_buf_ret < 0 { + av_frame_free(&mut fresh_hw_frame); + return Err("Failed to get hw frame buffer".into()); + } + + let transfer_ret = av_hwframe_transfer_data(fresh_hw_frame, self.cached_nv12, 0); + if transfer_ret < 0 { + av_frame_free(&mut fresh_hw_frame); + return Err("Failed to transfer cached data to hw frame".into()); + } + + (*fresh_hw_frame).pts = self.next_pts; + self.next_pts += 1; + + let send_ret = avcodec_send_frame(self.codec_ctx, fresh_hw_frame); + if send_ret >= 0 { + let mut pkt = av_packet_alloc(); + while avcodec_receive_packet(self.codec_ctx, pkt) >= 0 { + let new_pkt = av_packet_alloc(); + av_packet_move_ref(new_pkt, pkt); + packets.push(crate::ring::Packet::new(new_pkt)); + } + av_packet_free(&mut pkt); + } + av_frame_free(&mut fresh_hw_frame); + } + Ok(packets) + } } impl Drop for VaapiEncoder { fn drop(&mut self) { unsafe { + if !self.cached_nv12.is_null() { + av_frame_free(&mut self.cached_nv12); + self.cached_nv12 = ptr::null_mut(); + } if !self.sws_ctx.is_null() { sws_freeContext(self.sws_ctx); self.sws_ctx = ptr::null_mut(); @@ -544,6 +596,29 @@ impl WindowsHwEncoder { } Ok(packets) } + + #[cfg(target_os = "windows")] + pub fn encode_cached_frame(&mut self) -> Result, String> { + let mut packets = Vec::new(); + unsafe { + if self.sw_frame.is_null() { + return Ok(packets); + } + (*self.sw_frame).pts = self.next_pts; + self.next_pts += 1; + + if avcodec_send_frame(self.codec_ctx, self.sw_frame) >= 0 { + let mut pkt = av_packet_alloc(); + while avcodec_receive_packet(self.codec_ctx, pkt) >= 0 { + let new_pkt = av_packet_alloc(); + av_packet_move_ref(new_pkt, pkt); + packets.push(crate::ring::Packet::new(new_pkt)); + } + av_packet_free(&mut pkt); + } + } + Ok(packets) + } } #[cfg(target_os = "windows")] @@ -739,6 +814,25 @@ mod tests { } } + #[test] + fn test_audio_encode_sine_wave() { + let mut enc = AudioEncoder::new(48000, 2).expect("AudioEncoder creation"); + // Generate 48000 stereo samples of 440Hz sine wave (1.0s) + let mut pcm = Vec::with_capacity(48000 * 2); + for i in 0..48000 { + let val = ((i as f32 * 440.0 * 2.0 * std::f32::consts::PI) / 48000.0).sin() * 0.8; + pcm.push(val); // L + pcm.push(val); // R + } + + let packets = enc.encode_pcm(&pcm).expect("encode_pcm"); + println!("Encoded {} AAC packets from sine wave", packets.len()); + assert!(!packets.is_empty(), "Must produce AAC packets"); + let total_size: usize = packets.iter().map(|p| unsafe { (*p.ptr).size as usize }).sum(); + println!("Total AAC packet bytes: {}", total_size); + assert!(total_size > 1000, "AAC packets must not be empty or silence"); + } + #[test] fn test_vaapi_encoder_global_header() { if let Ok(mut enc) = VaapiEncoder::new(1280, 720) {