diff --git a/src/avhw.rs b/src/avhw.rs index 81a38b4..284887e 100644 --- a/src/avhw.rs +++ b/src/avhw.rs @@ -3,10 +3,10 @@ use std::mem; use std::os::fd::{AsRawFd, RawFd}; use std::os::raw::c_void; use std::path::Path; +use std::ptr; use std::slice; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use std::ptr; use anyhow::{bail, Result}; use ffmpeg_next as ff; @@ -16,6 +16,23 @@ use ffmpeg_next::packet::Mut as _; use crate::cap_portal::PwDmaBufFrame; use crate::transform::{transpose_if_transform_transposed, Transform}; +// --------------------------------------------------------------------------- +// Bitrate feedback command (WebRTC BWE → SW encoder) +// --------------------------------------------------------------------------- + +/// Commands sent from the WebRTC thread to the SW encoder when the +/// bandwidth estimate changes significantly. +pub enum BitrateCommand { + UpdateBitrate { target_bps: u64 }, + UpdateResolution { width: u32, height: u32 }, +} + +#[derive(Clone, Copy, Debug)] +pub struct ResolutionChange { + pub width: u32, + pub height: u32, +} + // --------------------------------------------------------------------------- // AvHwDevCtx // --------------------------------------------------------------------------- @@ -334,7 +351,9 @@ impl EncState { transform, )?; - let mut sink_ctx = video_filter.get("out").ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; + let mut sink_ctx = video_filter + .get("out") + .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; // SAFETY: sink_ctx is a live buffersink; the returned hw_frames_ctx is // borrowed, so av_buffer_ref creates an owned reference. let sink_hw_frames = unsafe { @@ -374,19 +393,19 @@ impl EncState { enc.set_width(enc_width); enc.set_height(enc_height); enc.set_format(ff::format::Pixel::VAAPI); - enc.set_bit_rate(bitrate as usize); - enc.set_gop(gop_size); - enc.set_time_base(ff::Rational::new(1, fps as i32)); - enc.set_max_b_frames(0); + enc.set_bit_rate(bitrate as usize); + enc.set_gop(gop_size); + enc.set_time_base(ff::Rational::new(1, fps as i32)); + enc.set_max_b_frames(0); - // VBV rate limiting: caps IDR burst size for WebRTC. Without this a 4K - // scene change can produce a 256KB keyframe that overflows the UDP send - // buffer. bufsize=bitrate/4 ≈ 250ms of video at the target bitrate. - unsafe { - let ctx_ptr = enc.as_mut_ptr(); - (*ctx_ptr).rc_max_rate = bitrate as i64; - (*ctx_ptr).rc_buffer_size = (bitrate / 4) as i32; - } + // VBV rate limiting: caps IDR burst size for WebRTC. Without this a 4K + // scene change can produce a 256KB keyframe that overflows the UDP send + // buffer. bufsize=bitrate/4 ≈ 250ms of video at the target bitrate. + unsafe { + let ctx_ptr = enc.as_mut_ptr(); + (*ctx_ptr).rc_max_rate = bitrate as i64; + (*ctx_ptr).rc_buffer_size = (bitrate / 4) as i32; + } // SAFETY: AV_CODEC_FLAG_GLOBAL_HEADER must be set BEFORE opening the encoder. // It triggers SPS/PPS extradata generation needed by the muxer for @@ -461,7 +480,10 @@ impl EncState { ffi::avcodec_parameters_from_context((*stream_ptr).codecpar, enc_video.as_ptr()) }; if ret < 0 { - bail!("Failed to copy encoder parameters to stream: {}", ff_err(ret)); + bail!( + "Failed to copy encoder parameters to stream: {}", + ff_err(ret) + ); } // SAFETY: Copy encoder time_base to stream. @@ -510,9 +532,15 @@ impl EncState { } pub fn encode_frame(&mut self, hw_frame: &ff::frame::Video) -> Result<()> { - let mut filter_src_ctx = self.video_filter.get("in").ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; + let mut filter_src_ctx = self + .video_filter + .get("in") + .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; let mut filter_src = filter_src_ctx.source(); - let mut filter_sink_ctx = self.video_filter.get("out").ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; + let mut filter_sink_ctx = self + .video_filter + .get("out") + .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; let mut filter_sink = filter_sink_ctx.sink(); // SAFETY: hw_frame is a valid VAAPI hardware frame from capture. @@ -552,14 +580,20 @@ impl EncState { pub fn flush(&mut self) -> Result<()> { // Flush filter graph - let mut filter_src_ctx = self.video_filter.get("in").ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; + let mut filter_src_ctx = self + .video_filter + .get("in") + .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; let mut filter_src = filter_src_ctx.source(); if let Err(e) = filter_src.flush() { tracing::debug!("filter source flush error: {e}"); } // Drain filter - let mut filter_sink_ctx = self.video_filter.get("out").ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; + let mut filter_sink_ctx = self + .video_filter + .get("out") + .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; let mut filter_sink = filter_sink_ctx.sink(); loop { let mut filtered = ff::frame::Video::empty(); @@ -667,8 +701,13 @@ pub struct SwEncImport { hw_dev: AvHwDevCtx, frames_rgb: AvHwFrameCtx, filter_graph: ff::filter::Graph, + width: u32, + height: u32, enc_width: u32, enc_height: u32, + fps: u32, + resolution_rx: Option>, + encoder_resolution_tx: Option>, } impl SwEncImport { @@ -698,20 +737,50 @@ impl SwEncImport { hw_dev, frames_rgb, filter_graph, + width, + height, enc_width, enc_height, + fps, + resolution_rx: None, + encoder_resolution_tx: None, }) } + #[allow(clippy::too_many_arguments)] + pub fn new_with_resolution_control( + drm_device: &Path, + width: u32, + height: u32, + enc_width: u32, + enc_height: u32, + fps: u32, + resolution_rx: crossbeam_channel::Receiver, + encoder_resolution_tx: crossbeam_channel::Sender, + ) -> Result { + let mut this = Self::new(drm_device, width, height, enc_width, enc_height, fps)?; + this.resolution_rx = Some(resolution_rx); + this.encoder_resolution_tx = Some(encoder_resolution_tx); + Ok(this) + } + pub fn frames_rgb(&self) -> &AvHwFrameCtx { let _ = self.hw_dev.as_ptr(); &self.frames_rgb } pub fn import_and_scale(&mut self, hw_frame: &ff::frame::Video) -> Result { - let mut filter_src_ctx = self.filter_graph.get("in").ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; + self.poll_resolution_commands()?; + + let mut filter_src_ctx = self + .filter_graph + .get("in") + .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; let mut filter_src = filter_src_ctx.source(); - let mut filter_sink_ctx = self.filter_graph.get("out").ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; + let mut filter_sink_ctx = self + .filter_graph + .get("out") + .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; let mut filter_sink = filter_sink_ctx.sink(); filter_src @@ -740,20 +809,28 @@ impl SwEncImport { } if extra_count > 0 { - tracing::warn!("software import filter produced {extra_count} extra frame(s); dropping extras"); + tracing::warn!( + "software import filter produced {extra_count} extra frame(s); dropping extras" + ); } first.ok_or_else(|| anyhow::anyhow!("software pipeline produced no scaled frame")) } pub fn flush_import(&mut self) -> Result> { - let mut filter_src_ctx = self.filter_graph.get("in").ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; + let mut filter_src_ctx = self + .filter_graph + .get("in") + .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?; let mut filter_src = filter_src_ctx.source(); if let Err(e) = filter_src.flush() { tracing::debug!("filter source flush error: {e}"); } - let mut filter_sink_ctx = self.filter_graph.get("out").ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; + let mut filter_sink_ctx = self + .filter_graph + .get("out") + .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?; let mut filter_sink = filter_sink_ctx.sink(); let mut frames = Vec::new(); loop { @@ -767,6 +844,54 @@ impl SwEncImport { Ok(frames) } + fn poll_resolution_commands(&mut self) -> Result<()> { + let Some(rx) = self.resolution_rx.as_ref().cloned() else { + return Ok(()); + }; + + let mut requested = None; + while let Ok(cmd) = rx.try_recv() { + match cmd { + BitrateCommand::UpdateResolution { width, height } => { + requested = Some((width & !1, height & !1)); + } + BitrateCommand::UpdateBitrate { .. } => {} + } + } + + let Some((width, height)) = requested else { + return Ok(()); + }; + if width == self.enc_width && height == self.enc_height { + return Ok(()); + } + + tracing::info!( + from = format_args!("{}x{}", self.enc_width, self.enc_height), + to = format_args!("{}x{}", width, height), + "rebuilding software import filter graph for resolution change" + ); + let _ = self.flush_import(); + self.filter_graph = build_swenc_filter_graph( + &self.hw_dev, + &self.frames_rgb, + self.width, + self.height, + width, + height, + self.fps, + )?; + self.enc_width = width; + self.enc_height = height; + + if let Some(tx) = &self.encoder_resolution_tx { + tx.send(ResolutionChange { width, height }) + .map_err(|_| anyhow::anyhow!("encoder resolution channel disconnected"))?; + } + + Ok(()) + } + fn transfer_filtered_to_cpu(&self, filtered: &ff::frame::Video) -> Result { // SAFETY: av_frame_alloc returns a newly allocated AVFrame or null, // which is checked below. @@ -798,7 +923,9 @@ impl SwEncImport { } let y_stride = (*sw_nv12).linesize[0] as usize; let uv_stride = (*sw_nv12).linesize[1] as usize; - if (*sw_nv12).width != self.enc_width as i32 || (*sw_nv12).height != self.enc_height as i32 { + if (*sw_nv12).width != self.enc_width as i32 + || (*sw_nv12).height != self.enc_height as i32 + { ffi::av_frame_free(&mut sw_nv12); bail!("NV12 transfer frame has unexpected dimensions"); } @@ -826,12 +953,38 @@ pub struct SwEncEncode { enc_video: ff::codec::encoder::video::Video, output: Option, yuv_frame: *mut ffi::AVFrame, + last_frame_hash: u64, + frame_count: u64, starting_timestamp: Option, frames_written: bool, webrtc_disconnected: bool, webrtc_paused: Option>, + bitrate_rx: crossbeam_channel::Receiver, + resolution_rx: crossbeam_channel::Receiver, enc_width: u32, enc_height: u32, + fps: u32, + bitrate: u64, + gop_size: u32, +} + +const FNV1A_OFFSET_BASIS: u64 = 0xcbf29ce484222325; +const FNV1A_PRIME: u64 = 0x100000001b3; +const Y_PLANE_HASH_ROW_STEP: usize = 8; + +fn hash_sampled_y_plane(y_data: &[u8], width: usize, height: usize, stride: usize) -> u64 { + let mut hash = FNV1A_OFFSET_BASIS; + + for row in (0..height).step_by(Y_PLANE_HASH_ROW_STEP) { + let row_start = row * stride; + let row_end = row_start + width; + for &byte in &y_data[row_start..row_end] { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(FNV1A_PRIME); + } + } + + hash } // SAFETY: SwEncEncode owns sws_ctx/yuv_frame/enc_video exclusively after construction. @@ -852,18 +1005,29 @@ impl SwEncEncode { let (enc_video, octx) = create_software_h264_muxer(output_path, enc_width, enc_height, fps, bitrate, gop_size)?; let yuv_frame = alloc_yuv420p_frame(enc_width, enc_height)?; + let (dummy_tx, bitrate_rx) = crossbeam_channel::bounded(1); + drop(dummy_tx); + let (dummy_resolution_tx, resolution_rx) = crossbeam_channel::bounded(1); + drop(dummy_resolution_tx); Ok(Self { sws_ctx, enc_video, output: Some(FrameOutput::Muxer(octx)), yuv_frame, + last_frame_hash: 0, + frame_count: 0, starting_timestamp: None, frames_written: false, webrtc_disconnected: false, webrtc_paused: None, + bitrate_rx, + resolution_rx, enc_width, enc_height, + fps, + bitrate, + gop_size, }) } @@ -876,9 +1040,12 @@ impl SwEncEncode { gop_size: u32, tx: crossbeam_channel::Sender>, webrtc_paused: Arc, + bitrate_rx: crossbeam_channel::Receiver, + resolution_rx: crossbeam_channel::Receiver, ) -> Result { let sws_ctx = create_nv12_to_yuv420p_sws(enc_width, enc_height)?; - let enc_video = create_software_h264_encoder(enc_width, enc_height, fps, bitrate, gop_size)?; + let enc_video = + create_software_h264_encoder(enc_width, enc_height, fps, bitrate, gop_size)?; let yuv_frame = alloc_yuv420p_frame(enc_width, enc_height)?; Ok(Self { @@ -886,12 +1053,19 @@ impl SwEncEncode { enc_video, output: Some(FrameOutput::Channel(tx)), yuv_frame, + last_frame_hash: 0, + frame_count: 0, starting_timestamp: None, frames_written: false, webrtc_disconnected: false, webrtc_paused: Some(webrtc_paused), + bitrate_rx, + resolution_rx, enc_width, enc_height, + fps, + bitrate, + gop_size, }) } @@ -914,6 +1088,28 @@ impl SwEncEncode { if self.webrtc_disconnected { return Ok(()); } + + // Must drain before the stride check: the import thread emits + // ResolutionChange before the new (smaller-stride) frame arrives. + while let Ok(cmd) = self.bitrate_rx.try_recv() { + match cmd { + BitrateCommand::UpdateBitrate { target_bps } => { + tracing::info!(target_bps, "updating encoder bitrate from BWE feedback"); + self.bitrate = target_bps; + // SAFETY: enc_video is an opened AVCodecContext exclusively owned by &mut self. + unsafe { + let ctx = self.enc_video.as_mut_ptr(); + (*ctx).bit_rate = target_bps as i64; + } + } + BitrateCommand::UpdateResolution { .. } => {} + } + } + + while let Ok(change) = self.resolution_rx.try_recv() { + self.recreate_encoder(change.width, change.height)?; + } + if frame.y_stride < self.enc_width as usize || frame.uv_stride < self.enc_width as usize { bail!("CPU NV12 frame stride is smaller than encoder width"); } @@ -923,6 +1119,24 @@ impl SwEncEncode { } } + let width = self.enc_width as usize; + let height = self.enc_height as usize; + let required_y_len = frame.y_stride * height.saturating_sub(1) + width; + if frame.y_data.len() < required_y_len { + bail!("CPU NV12 frame Y plane is smaller than encoder dimensions"); + } + + let frame_index = self.frame_count; + 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 { + tracing::debug!(frame_index, "skipping duplicate frame"); + self.last_frame_hash = current_hash; + return Ok(()); + } + self.last_frame_hash = current_hash; + // SAFETY: yuv_frame is an owned reusable YUV420P frame at the same dimensions as sw_nv12; // sws_ctx was created for NV12 -> YUV420P with no resize, so sws_scale only converts format. unsafe { @@ -930,7 +1144,12 @@ impl SwEncEncode { if ret < 0 { bail!("av_frame_make_writable failed: {}", ff_err(ret)); } - let src_slices = [frame.y_data.as_ptr(), frame.uv_data.as_ptr(), ptr::null(), ptr::null()]; + let src_slices = [ + frame.y_data.as_ptr(), + frame.uv_data.as_ptr(), + ptr::null(), + ptr::null(), + ]; let src_strides = [frame.y_stride as i32, frame.uv_stride as i32, 0, 0]; let scaled = ffi::sws_scale( self.sws_ctx, @@ -957,13 +1176,49 @@ impl SwEncEncode { (*self.yuv_frame).pts = pts; let ret = ffi::avcodec_send_frame(self.enc_video.as_mut_ptr(), self.yuv_frame); if ret < 0 { - bail!("avcodec_send_frame failed for software encoder: {}", ff_err(ret)); + bail!( + "avcodec_send_frame failed for software encoder: {}", + ff_err(ret) + ); } } self.drain_encoder(start_ts) } + fn recreate_encoder(&mut self, width: u32, height: u32) -> Result<()> { + if width == self.enc_width && height == self.enc_height { + return Ok(()); + } + + tracing::info!( + from = format_args!("{}x{}", self.enc_width, self.enc_height), + to = format_args!("{}x{}", width, height), + "recreating WebRTC software encoder for resolution change" + ); + + if !self.sws_ctx.is_null() { + // SAFETY: sws_ctx is owned exclusively by self and will be replaced below. + unsafe { ffi::sws_freeContext(self.sws_ctx) }; + self.sws_ctx = ptr::null_mut(); + } + if !self.yuv_frame.is_null() { + // SAFETY: yuv_frame is owned exclusively by self and will be replaced below. + unsafe { ffi::av_frame_free(&mut self.yuv_frame) }; + } + + self.sws_ctx = create_nv12_to_yuv420p_sws(width, height)?; + self.enc_video = + create_software_h264_encoder(width, height, self.fps, self.bitrate, self.gop_size)?; + self.yuv_frame = alloc_yuv420p_frame(width, height)?; + self.enc_width = width; + self.enc_height = height; + self.last_frame_hash = 0; + self.frame_count = 0; + + Ok(()) + } + fn write_trailer_if_needed(&mut self) -> Result<()> { if self.frames_written { if let Some(FrameOutput::Muxer(ref mut octx)) = self.output { @@ -1026,9 +1281,8 @@ impl SwEncEncode { // `data` is non-null, so `data` points to `size` initialized // bytes owned by the packet. `u8` has alignment 1, and the // slice is copied into a Vec before the packet is unreffed. - let data: &[u8] = unsafe { - std::slice::from_raw_parts(raw.data, raw.size as usize) - }; + let data: &[u8] = + unsafe { std::slice::from_raw_parts(raw.data, raw.size as usize) }; match tx.try_send(data.to_vec()) { Ok(()) => {} Err(crossbeam_channel::TrySendError::Full(frame)) => { @@ -1095,7 +1349,8 @@ impl SwEncState { "SwEncState::new: GPU downscale {width}x{height} BGRA -> {enc_width}x{enc_height} NV12, software H.264" ); let import = SwEncImport::new(drm_device, width, height, enc_width, enc_height, fps)?; - let encode = SwEncEncode::new_muxer(output_path, enc_width, enc_height, fps, bitrate, gop_size)?; + let encode = + SwEncEncode::new_muxer(output_path, enc_width, enc_height, fps, bitrate, gop_size)?; Ok(Self { import, encode }) } @@ -1116,6 +1371,10 @@ impl SwEncState { "SwEncState::new_webrtc: GPU downscale {width}x{height} BGRA -> {enc_width}x{enc_height} NV12, software H.264 -> WebRTC" ); let import = SwEncImport::new(drm_device, width, height, enc_width, enc_height, fps)?; + let (dummy_tx, bitrate_rx) = crossbeam_channel::bounded(1); + drop(dummy_tx); + let (dummy_resolution_tx, resolution_rx) = crossbeam_channel::bounded(1); + drop(dummy_resolution_tx); let encode = SwEncEncode::new_webrtc( enc_width, enc_height, @@ -1124,6 +1383,8 @@ impl SwEncState { gop_size, tx, webrtc_paused, + bitrate_rx, + resolution_rx, )?; Ok(Self { import, encode }) } @@ -1348,6 +1609,10 @@ fn create_software_h264_muxer( let val = CString::new("6").unwrap(); ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0); (*enc.as_mut_ptr()).profile = ffi::AV_PROFILE_H264_HIGH as i32; + // SAFETY: enc is a valid, initialized AVCodecContext from + // avcodec_alloc_context3. Setting level is a simple i32 field + // assignment on a properly aligned struct. + (*enc.as_mut_ptr()).level = 40; // H.264 Level 4.0 (up to 1080p@30) } } @@ -1472,8 +1737,17 @@ fn create_software_h264_encoder( // High profile via AVCodecContext.profile (not x264opts — x264 rejects it there). // High enables CABAC + 8x8dct automatically. (*enc.as_mut_ptr()).profile = ffi::AV_PROFILE_H264_HIGH as i32; + // SAFETY: enc is a valid, initialized AVCodecContext from + // 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) let key = CString::new("x264opts").unwrap(); - let val = CString::new("repeat_headers=1").unwrap(); + let vbv_maxrate = bitrate; + let vbv_bufsize = bitrate / 4; + let val = CString::new(format!( + "repeat_headers=1:vbv-maxrate={vbv_maxrate}:vbv-bufsize={vbv_bufsize}" + )) + .unwrap(); ffi::av_opt_set((*enc.as_mut_ptr()).priv_data, key.as_ptr(), val.as_ptr(), 0); } } @@ -1590,3 +1864,129 @@ fn build_filter_graph( Ok(graph) } + +#[cfg(test)] +mod tests { + use super::*; + + // ── Task 1: VBV x264opts formatting ── + + #[test] + fn vbv_x264opts_format() { + let bitrate: u64 = 5_000_000; + let vbv_maxrate = bitrate; + let vbv_bufsize = bitrate / 4; + let opts = format!("repeat_headers=1:vbv-maxrate={vbv_maxrate}:vbv-bufsize={vbv_bufsize}"); + assert!(opts.contains("vbv-maxrate=5000000")); + assert!(opts.contains("vbv-bufsize=1250000")); + } + + #[test] + fn vbv_bufsize_is_quarter_of_maxrate() { + for bitrate in [1_000_000, 5_000_000, 10_000_000] { + let maxrate = bitrate; + let bufsize = bitrate / 4; + assert_eq!(bufsize * 4, maxrate, "bufsize should be maxrate/4"); + } + } + + // ── Task 3: GOP formula ── + + #[test] + fn webrtc_gop_formula() { + assert_eq!((15u32 * 2).max(20), 30); // 15fps -> 30 + assert_eq!((30u32 * 2).max(20), 60); // 30fps -> 60 + assert_eq!((60u32 * 2).max(20), 120); // 60fps -> 120 + assert_eq!((5u32 * 2).max(20), 20); // 5fps -> 20 (floor) + } + + #[test] + fn h264_level_values() { + // Level 4.0 supports up to 1080p@30fps (used for file muxer) + assert_eq!(40i32, 40); + // Level 4.2 supports up to 1440p@30fps (used for WebRTC low-latency encoder) + assert_eq!(42i32, 42); + } + + // ── Task 4: Duplicate frame hash detection ── + + #[test] + fn hash_sampled_y_plane_first_frame_consistent() { + let width = 64; + let height = 64; + let stride = 64; + let y_data = vec![0u8; stride * height]; + let hash1 = hash_sampled_y_plane(&y_data, width, height, stride); + let hash2 = hash_sampled_y_plane(&y_data, width, height, stride); + assert_eq!(hash1, hash2, "same input should produce same hash"); + } + + #[test] + fn hash_sampled_y_plane_detects_different_frames() { + let width = 64; + let height = 64; + let stride = 64; + let y_data1 = vec![0u8; stride * height]; + let y_data2 = vec![128u8; stride * height]; + let hash1 = hash_sampled_y_plane(&y_data1, width, height, stride); + let hash2 = hash_sampled_y_plane(&y_data2, width, height, stride); + assert_ne!( + hash1, hash2, + "different frame data should produce different hashes" + ); + } + + #[test] + fn hash_sampled_y_plane_samples_every_8th_row() { + // Changing a non-sampled row (e.g., row 1) should NOT change the hash + let width = 64; + let height = 64; + let stride = 64; + let y_data1 = vec![0u8; stride * height]; + let mut y_data2 = vec![0u8; stride * height]; + // Row 1 is NOT sampled (sampling is every 8th row: 0, 8, 16, ...) + y_data2[stride * 1..stride * 1 + width].fill(255); + let hash1 = hash_sampled_y_plane(&y_data1, width, height, stride); + let hash2 = hash_sampled_y_plane(&y_data2, width, height, stride); + assert_eq!( + hash1, hash2, + "non-sampled row change should not affect hash" + ); + } + + #[test] + fn hash_sampled_y_plane_sensitive_to_sampled_row() { + // Changing a sampled row (row 0) SHOULD change the hash + let width = 64; + let height = 64; + let stride = 64; + let y_data1 = vec![0u8; stride * height]; + let mut y_data2 = vec![0u8; stride * height]; + // Row 0 IS sampled (every 8th row starting from 0) + y_data2[stride * 0..stride * 0 + width].fill(255); + let hash1 = hash_sampled_y_plane(&y_data1, width, height, stride); + let hash2 = hash_sampled_y_plane(&y_data2, width, height, stride); + assert_ne!( + hash1, hash2, + "sampled row change should produce different hash" + ); + } + + #[test] + fn hash_sampled_y_plane_handles_stride_greater_than_width() { + // Stride can be larger than width due to alignment; unused padding should not affect hash + let width = 32; + let height = 16; + let stride = 64; // padded stride + let y_data1 = vec![0u8; stride * height]; + let mut y_data2 = vec![0u8; stride * height]; + // Fill the padding area (columns 32..63) of row 0 with garbage + y_data2[width..stride].fill(0xFF); + let hash1 = hash_sampled_y_plane(&y_data1, width, height, stride); + let hash2 = hash_sampled_y_plane(&y_data2, width, height, stride); + assert_eq!( + hash1, hash2, + "padding bytes beyond width should not affect hash" + ); + } +} diff --git a/src/cap_portal.rs b/src/cap_portal.rs index da661d5..5e1b0c0 100644 --- a/src/cap_portal.rs +++ b/src/cap_portal.rs @@ -225,9 +225,7 @@ impl CapPortal { proxy .select_sources(&session, options) .await - .map_err(|e| { - anyhow::anyhow!("Screen sharing permission denied: {e}") - })?; + .map_err(|e| anyhow::anyhow!("Screen sharing permission denied: {e}"))?; let response = proxy .start(&session, None, Default::default()) @@ -272,7 +270,10 @@ fn verify_secure_dir(path: &std::path::Path) -> bool { match std::fs::symlink_metadata(path) { Ok(meta) => { if meta.file_type().is_symlink() { - tracing::warn!("Token parent dir is a symlink, rejecting: {}", path.display()); + tracing::warn!( + "Token parent dir is a symlink, rejecting: {}", + path.display() + ); return false; } // Must be a directory @@ -282,7 +283,10 @@ fn verify_secure_dir(path: &std::path::Path) -> bool { } // Must be owned by current user if meta.uid() != unsafe { libc::getuid() } { - tracing::warn!("Token parent dir not owned by current user: {}", path.display()); + tracing::warn!( + "Token parent dir not owned by current user: {}", + path.display() + ); return false; } // No group or other permissions (mode must be 0o700 exactly within the 0o777 mask) @@ -346,7 +350,10 @@ fn load_restore_token_from(path: PathBuf) -> Option { }; if meta.file_type().is_symlink() { - tracing::warn!("Token file is a symlink, refusing to read: {}", path.display()); + tracing::warn!( + "Token file is a symlink, refusing to read: {}", + path.display() + ); return None; } if !meta.is_file() { @@ -369,7 +376,11 @@ fn load_restore_token_from(path: PathBuf) -> Option { let token = std::fs::read_to_string(&path).ok()?; let trimmed = token.trim().to_string(); - if trimmed.is_empty() { None } else { Some(trimmed) } + if trimmed.is_empty() { + None + } else { + Some(trimmed) + } } fn save_restore_token(token: &str) { @@ -490,7 +501,9 @@ fn pipewire_thread(ctx: PwThreadCtx) { let mainloop = match pw::main_loop::MainLoopBox::new(None) { Ok(ml) => ml, Err(e) => { - if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("MainLoop::new failed: {e}"))) { + if let Err(e) = + event_tx.try_send(PwCtrlEvent::Error(format!("MainLoop::new failed: {e}"))) + { tracing::error!("MainLoop::new failed and error channel also failed: {e}"); } return; @@ -500,7 +513,9 @@ fn pipewire_thread(ctx: PwThreadCtx) { let context = match pw::context::ContextBox::new(mainloop.loop_(), None) { Ok(c) => c, Err(e) => { - if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("Context::new failed: {e}"))) { + if let Err(e) = + event_tx.try_send(PwCtrlEvent::Error(format!("Context::new failed: {e}"))) + { tracing::error!("Context::new failed and error channel also failed: {e}"); } return; @@ -510,7 +525,8 @@ fn pipewire_thread(ctx: PwThreadCtx) { let core = match context.connect_fd(pw_fd, None) { Ok(c) => c, Err(e) => { - if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("connect_fd failed: {e}"))) { + if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("connect_fd failed: {e}"))) + { tracing::error!("connect_fd failed and error channel also failed: {e}"); } return; @@ -534,7 +550,9 @@ fn pipewire_thread(ctx: PwThreadCtx) { ) { Ok(s) => s, Err(e) => { - if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("Stream::new failed: {e}"))) { + if let Err(e) = + event_tx.try_send(PwCtrlEvent::Error(format!("Stream::new failed: {e}"))) + { tracing::error!("Stream::new failed and error channel also failed: {e}"); } return; @@ -608,8 +626,10 @@ fn pipewire_thread(ctx: PwThreadCtx) { "PipeWire format negotiated: {width}x{height}, \ drm_format={drm_format:#010x}, modifier={modifier:#x}, \ framerate={}/{}, max_framerate={}/{}", - framerate.num, framerate.denom, - max_framerate.num, max_framerate.denom, + framerate.num, + framerate.denom, + max_framerate.num, + max_framerate.denom, ); } }) @@ -621,7 +641,6 @@ fn pipewire_thread(ctx: PwThreadCtx) { let frame_tx = frame_tx.clone(); let dropped = dropped; move |stream, _| { - let raw_buf = unsafe { stream.dequeue_raw_buffer() }; if raw_buf.is_null() { tracing::trace!("process: null raw_buf"); @@ -742,7 +761,8 @@ fn pipewire_thread(ctx: PwThreadCtx) { StreamFlags::AUTOCONNECT | StreamFlags::MAP_BUFFERS, &mut params, ) { - if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("stream.connect failed: {e}"))) { + if let Err(e) = event_tx.try_send(PwCtrlEvent::Error(format!("stream.connect failed: {e}"))) + { tracing::error!("stream.connect failed and error channel also failed: {e}"); } return; @@ -918,7 +938,10 @@ mod tests { let meta = std::fs::symlink_metadata(&new_dir).unwrap(); let mode = meta.permissions().mode() & 0o777; - assert_eq!(mode, 0o700, "created directory should be 0700, got {mode:o}"); + assert_eq!( + mode, 0o700, + "created directory should be 0700, got {mode:o}" + ); } #[test] @@ -933,7 +956,10 @@ mod tests { let meta = std::fs::symlink_metadata(path).unwrap(); let mode = meta.permissions().mode() & 0o777; - assert_eq!(mode, 0o700, "tightened directory should be 0700, got {mode:o}"); + assert_eq!( + mode, 0o700, + "tightened directory should be 0700, got {mode:o}" + ); } #[test] @@ -947,7 +973,10 @@ mod tests { let meta = std::fs::symlink_metadata(&token_path).unwrap(); let mode = meta.permissions().mode() & 0o777; assert_eq!(mode, 0o600, "token file should be 0600, got {mode:o}"); - assert_eq!(std::fs::read_to_string(&token_path).unwrap(), "secret-token-123"); + assert_eq!( + std::fs::read_to_string(&token_path).unwrap(), + "secret-token-123" + ); } #[test] diff --git a/src/fps_limit.rs b/src/fps_limit.rs index efdfd98..1fca42d 100644 --- a/src/fps_limit.rs +++ b/src/fps_limit.rs @@ -82,7 +82,12 @@ mod tests { } } - assert!(outputs.len() >= 3, "expected at least 3 outputs, got {} ({:?})", outputs.len(), outputs); + assert!( + outputs.len() >= 3, + "expected at least 3 outputs, got {} ({:?})", + outputs.len(), + outputs + ); assert_eq!(outputs[0], 0); } diff --git a/src/lib.rs b/src/lib.rs index 26513e3..d245351 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,8 +4,8 @@ pub mod backend_detect; pub mod cap_portal; pub mod cap_wlr_screencopy; pub mod fps_limit; -pub mod stats; pub mod state; pub mod state_portal; +pub mod stats; pub mod transform; pub mod webrtc; diff --git a/src/main.rs b/src/main.rs index 6557511..c36e8fd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,9 +15,9 @@ mod backend_detect; // 截屏后端自动检测（wlroots vs Portal/PipeWire） mod cap_portal; // XDG Portal 屏幕捕获 mod cap_wlr_screencopy; // wlroots wlr-screencopy 截屏协议 mod fps_limit; // 帧率限制器 -mod stats; // 管道性能统计（卡顿诊断） mod state; // wlr-screencopy 后端的主状态机 mod state_portal; // Portal/PipeWire 后端的主状态机 +mod stats; // 管道性能统计（卡顿诊断） mod transform; // 图像变换（旋转/翻转） mod webrtc; // WebRTC 传输（str0m Sans-IO） @@ -46,21 +46,27 @@ fn main() -> Result<()> { // 根据 verbose 模式或 RUST_LOG 环境变量设置日志级别 // 支持 RUST_LOG 粒度控制（如 RUST_LOG=wl_webrtc::webrtc=trace） - let env_filter = tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| { - if args.verbose { - tracing_subscriber::EnvFilter::new("debug") - } else { - tracing_subscriber::EnvFilter::new("info") - } - }); + let env_filter = tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| { + if args.verbose { + tracing_subscriber::EnvFilter::new("debug") + } else { + tracing_subscriber::EnvFilter::new("info") + } + }); tracing_subscriber::fmt() .with_env_filter(env_filter) .with_writer(std::io::stderr) .init(); tracing::info!("wl-webrtc starting"); - tracing::debug!("Args: output={:?} fps={} codec={} port={} verbose={}", args.output, args.fps, args.codec, args.port, args.verbose); + tracing::debug!( + "Args: output={:?} fps={} codec={} port={} verbose={}", + args.output, + args.fps, + args.codec, + args.port, + args.verbose + ); // MVP 阶段仅支持 H.264 编码，不支持 HEVC if args.codec != "h264" { @@ -352,7 +358,9 @@ fn run_portal_pipewire(args: Args) -> Result<()> { // Portal 状态机遇到致命错误时退出 if state.is_errored() { - tracing::error!("Fatal error in portal state machine (check preceding error logs), exiting"); + tracing::error!( + "Fatal error in portal state machine (check preceding error logs), exiting" + ); running = false; } } diff --git a/src/state.rs b/src/state.rs index 62de184..4513c71 100644 --- a/src/state.rs +++ b/src/state.rs @@ -539,7 +539,10 @@ impl State { // takes ownership of the fd, and the original fd is owned by map_frame. let fd_dup = unsafe { libc::dup(obj.fd) }; if fd_dup < 0 { - tracing::error!("failed to dup dma-buf fd: {}", std::io::Error::last_os_error()); + tracing::error!( + "failed to dup dma-buf fd: {}", + std::io::Error::last_os_error() + ); // wayland-client does not auto-destroy params on Drop. params.destroy(); self.errored = true; @@ -641,7 +644,11 @@ impl State { if last.elapsed() >= std::time::Duration::from_secs(10) { let delta = self.stats_frames; let fps = delta as f64 / last.elapsed().as_secs_f64(); - tracing::info!(frames = self.stats_frames, fps = format!("{fps:.1}"), "encoding stats"); + tracing::info!( + frames = self.stats_frames, + fps = format!("{fps:.1}"), + "encoding stats" + ); self.stats_last_time = Some(std::time::Instant::now()); self.stats_frames = 0; } @@ -672,7 +679,9 @@ impl State { } pub fn poll_webrtc(&mut self) -> Result<()> { - let Some(ref mut wrtc) = self.webrtc else { return Ok(()) }; + let Some(ref mut wrtc) = self.webrtc else { + return Ok(()); + }; wrtc.handle_signaling()?; wrtc.poll_and_feed()?; @@ -697,7 +706,8 @@ impl State { continue; } count += 1; - if let Err(e) = wrtc.write_h264_frame(&data, self.webrtc_frames_sent, self.args.fps) { + if let Err(e) = wrtc.write_h264_frame(&data, self.webrtc_frames_sent, self.args.fps) + { tracing::debug!("WebRTC write frame error: {e}"); } self.stats.record_send(0.0, None); @@ -709,10 +719,8 @@ impl State { } if self.args.stats && self.stats.should_snapshot() { - self.stats.set_queue_depths( - 0, - self.webrtc_rx.as_ref().map(|r| r.len()).unwrap_or(0), - ); + self.stats + .set_queue_depths(0, self.webrtc_rx.as_ref().map(|r| r.len()).unwrap_or(0)); let snap = self.stats.snapshot_and_reset(); tracing::info!("stats: {snap}"); } @@ -752,9 +760,12 @@ impl State { .unwrap_or_else(|| 2 * (width as u64) * (height as u64) * (fps as u64) / 100); let enc = if let Some(ref tx) = self.webrtc_tx { - let (enc_w, enc_h) = - transpose_if_transform_transposed(output_info.transform, width as i32, height as i32); - let actual_gop_size = self.args.gop_size.unwrap_or((fps / 2).max(10)); + let (enc_w, enc_h) = transpose_if_transform_transposed( + output_info.transform, + width as i32, + height as i32, + ); + let actual_gop_size = self.args.gop_size.unwrap_or((fps * 2).max(20)); match SwEncState::new_webrtc( &drm_path, width, @@ -765,7 +776,10 @@ impl State { bitrate, actual_gop_size, tx.clone(), - self.webrtc_paused.as_ref().expect("webrtc_paused must exist when webrtc_tx exists").clone(), + self.webrtc_paused + .as_ref() + .expect("webrtc_paused must exist when webrtc_tx exists") + .clone(), ) { Ok(enc) => StreamingEncoder::WebRtc(enc), Err(e) => { @@ -775,7 +789,11 @@ impl State { } } } else { - let output_path = self.args.output.as_deref().expect("output required for MP4 mode"); + let output_path = self + .args + .output + .as_deref() + .expect("output required for MP4 mode"); match crate::avhw::create_encoder( &drm_path, Path::new(output_path), @@ -1035,7 +1053,7 @@ impl Dispatch for State { qhandle: &QueueHandle>, ) { use wayland_client::protocol::wl_registry::Event as RegistryEvent; - + match event { RegistryEvent::Global { name, diff --git a/src/state_portal.rs b/src/state_portal.rs index 87147e2..deefede 100644 --- a/src/state_portal.rs +++ b/src/state_portal.rs @@ -5,10 +5,12 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use anyhow::{bail, Result}; // 错误处理工具 +use anyhow::{bail, Result}; // 错误处理工具 -use crate::args::Args; // 命令行参数 -use crate::avhw::{self, CpuNv12Frame, SwEncEncode, SwEncImport, SwEncState}; // 软件编码器状态（VAAPI 导入 + H.264 编码） +use crate::args::Args; // 命令行参数 +use crate::avhw::{ + self, BitrateCommand, CpuNv12Frame, ResolutionChange, SwEncEncode, SwEncImport, SwEncState, +}; // 软件编码器状态（VAAPI 导入 + H.264 编码） use crate::cap_portal::{CapPortal, PwCtrlEvent, PwDmaBufFrame}; // PipeWire 屏幕采集端点 use crate::stats::{FrameTimings, PipelineStats}; // 管道统计（帧计时、每秒快照） use crate::webrtc::WebRtcState; // WebRTC 信令与媒体传输 @@ -43,26 +45,26 @@ struct WebrtcThread { /// 负责管理从 PipeWire 采集屏幕帧、通过 VAAPI 硬件编码的完整生命周期。 /// 工作流程：等待第一帧 → 创建编码器 → 持续编码帧数据。 pub struct StatePortal { - stage: PortalStage, // 当前采集阶段（等待首帧 / 流式编码中） - enc: Option, // 软件编码器，首帧到达后初始化 + stage: PortalStage, // 当前采集阶段（等待首帧 / 流式编码中） + enc: Option, // 软件编码器，首帧到达后初始化 enc_import: Option, enc_thread: Option, - cap: CapPortal, // PipeWire 屏幕采集端点 - args: Args, // 用户命令行参数 - errored: bool, // 是否遇到不可恢复的错误 - drm_device: Option, // DRM 渲染设备路径（可自动检测） - frames_encoded: u64, // 已编码帧数（用于 PTS 编号） - start_time: Option, // 编码开始时间 - stats: PipelineStats, // 管道统计（窗口化帧计时 + 每秒快照） - pw_dropped_prev: u64, // 上一窗口的 PipeWire 丢弃帧数（用于增量计算） + cap: CapPortal, // PipeWire 屏幕采集端点 + args: Args, // 用户命令行参数 + errored: bool, // 是否遇到不可恢复的错误 + drm_device: Option, // DRM 渲染设备路径（可自动检测） + frames_encoded: u64, // 已编码帧数（用于 PTS 编号） + start_time: Option, // 编码开始时间 + stats: PipelineStats, // 管道统计（窗口化帧计时 + 每秒快照） + pw_dropped_prev: u64, // 上一窗口的 PipeWire 丢弃帧数（用于增量计算） webrtc: Option, webrtc_thread: Option, webrtc_paused: Option>, - last_capture_arrival: Option, // timestamp of last real frame arrival - stall_start: Option, // when current stall began - last_stall_log: Option, // rate-limiting for stall warnings - last_fillable_frame: Option, // cached last frame for filler duplication - next_filler_at: Option, // when to send next filler frame + last_capture_arrival: Option, // timestamp of last real frame arrival + stall_start: Option, // when current stall began + last_stall_log: Option, // rate-limiting for stall warnings + last_fillable_frame: Option, // cached last frame for filler duplication + next_filler_at: Option, // when to send next filler frame filler_frames_sent: u64, } @@ -147,7 +149,11 @@ impl StatePortal { // 根据阻塞模式选择不同的帧接收策略 let frame = if block { // 阻塞模式：最多等待 2ms 接收帧 - match self.cap.frame_receiver().recv_timeout(std::time::Duration::from_millis(2)) { + match self + .cap + .frame_receiver() + .recv_timeout(std::time::Duration::from_millis(2)) + { Ok(frame) => frame, Err(_) => { self.record_capture_timeout(); @@ -194,10 +200,10 @@ impl StatePortal { let actual_bitrate = self.args.bitrate.unwrap_or_else(|| { 5 * (enc_width as u64) * (enc_height as u64) * (self.args.fps as u64) / 100 }); - // GOP 大小：WebRTC 模式使用更小的 GOP（fps/2，最低10），MP4 模式使用 fps + // GOP 大小：WebRTC 模式使用较大的 GOP（fps*2，最低20），MP4 模式使用 fps let actual_gop_size = self.args.gop_size.unwrap_or_else(|| { if self.webrtc.is_some() { - (self.args.fps / 2).max(10) + (self.args.fps * 2).max(20) } else { self.args.fps } @@ -207,17 +213,25 @@ impl StatePortal { if self.webrtc.is_some() { let paused = self.webrtc_paused.as_ref() .ok_or_else(|| anyhow::anyhow!("internal invariant broken: webrtc_paused missing while WebRTC mode is active"))?; - let import = SwEncImport::new( + let (resolution_tx, resolution_rx) = + crossbeam_channel::bounded::(4); + let (encoder_resolution_tx, encoder_resolution_rx) = + crossbeam_channel::bounded::(4); + let import = SwEncImport::new_with_resolution_control( &drm_path, frame.width, frame.height, enc_width, enc_height, self.args.fps, + resolution_rx, + encoder_resolution_tx, )?; let (webrtc_tx, webrtc_rx) = crossbeam_channel::bounded(2); let (input_tx, input_rx) = crossbeam_channel::bounded::(1); - let (timing_tx, timing_rx) = crossbeam_channel::bounded::(32); + let (timing_tx, timing_rx) = + crossbeam_channel::bounded::(32); + let (bitrate_tx, bitrate_rx) = crossbeam_channel::bounded::(4); let encode = SwEncEncode::new_webrtc( enc_width, enc_height, @@ -226,24 +240,48 @@ impl StatePortal { actual_gop_size, webrtc_tx, paused.clone(), + bitrate_rx, + encoder_resolution_rx, )?; let handle = std::thread::Builder::new() .name("wl-webrtc-encode".into()) .spawn(move || encode_thread_loop(encode, input_rx, timing_tx))?; self.enc_import = Some(import); - self.enc_thread = Some(EncodeThread { handle: Some(handle), input_tx, timing_rx }); + self.enc_thread = Some(EncodeThread { + handle: Some(handle), + input_tx, + timing_rx, + }); - let wrtc = self.webrtc.take() - .ok_or_else(|| anyhow::anyhow!("internal: WebRtcState missing during init"))?; - let paused = self.webrtc_paused.as_ref() + let wrtc = self.webrtc.take().ok_or_else(|| { + anyhow::anyhow!("internal: WebRtcState missing during init") + })?; + let paused = self + .webrtc_paused + .as_ref() .ok_or_else(|| anyhow::anyhow!("internal: webrtc_paused missing"))? .clone(); let fps = self.args.fps; let (sent_gap_tx, sent_gap_rx) = crossbeam_channel::bounded(64); let webrtc_handle = std::thread::Builder::new() .name("wl-webrtc-webrtc".into()) - .spawn(move || webrtc_thread_loop(wrtc, webrtc_rx, fps, paused, sent_gap_tx))?; - self.webrtc_thread = Some(WebrtcThread { handle: Some(webrtc_handle), sent_gap_rx }); + .spawn(move || { + webrtc_thread_loop( + wrtc, + webrtc_rx, + fps, + enc_width, + enc_height, + paused, + sent_gap_tx, + bitrate_tx, + resolution_tx, + ) + })?; + self.webrtc_thread = Some(WebrtcThread { + handle: Some(webrtc_handle), + sent_gap_rx, + }); } else { // MP4 模式：编码输出写入文件 let output_path = self.args.output.as_deref() @@ -263,7 +301,9 @@ impl StatePortal { }; self.stage = PortalStage::Streaming; // 切换到流式编码阶段 self.start_time = Some(Instant::now()); - tracing::info!("First frame processed, encoder initialized, transitioning to Streaming"); + tracing::info!( + "First frame processed, encoder initialized, transitioning to Streaming" + ); drop(frame); // 首帧仅用于初始化，不参与编码 } PortalStage::Streaming => { @@ -295,7 +335,10 @@ impl StatePortal { } let snap = self.stats.snapshot_and_reset(); if self.filler_frames_sent > 0 { - tracing::info!("stats: {snap} filler_frames_sent={}", self.filler_frames_sent); + tracing::info!( + "stats: {snap} filler_frames_sent={}", + self.filler_frames_sent + ); } else { tracing::info!("stats: {snap}"); } @@ -321,9 +364,9 @@ impl StatePortal { self.last_stall_log = Some(now); tracing::warn!("compositor frame delivery stalled"); } else { - let should_log = self - .last_stall_log - .map_or(true, |last_log| now.duration_since(last_log) >= Duration::from_secs(1)); + let should_log = self.last_stall_log.map_or(true, |last_log| { + now.duration_since(last_log) >= Duration::from_secs(1) + }); if should_log { self.last_stall_log = Some(now); tracing::warn!("compositor frame delivery stalled"); @@ -437,13 +480,11 @@ impl StatePortal { } // 所有候选设备均失败，返回详细错误信息 - bail!( - failures - .into_iter() - .map(|(p, e)| format!("{} ({e})", p.display())) - .collect::>() - .join(", ") - ); + bail!(failures + .into_iter() + .map(|(p, e)| format!("{} ({e})", p.display())) + .collect::>() + .join(", ")); } /// 处理单帧 DMA-BUF 数据 @@ -513,8 +554,11 @@ impl StatePortal { let import_us = t_import_start.elapsed().as_micros() as u64; self.stats.record_import(import_us); - let enc_thread = self.enc_thread.as_ref() - .ok_or_else(|| anyhow::anyhow!("internal invariant broken: encode thread missing while async import is active"))?; + let enc_thread = self.enc_thread.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "internal invariant broken: encode thread missing while async import is active" + ) + })?; let fillable_frame = CpuNv12Frame { y_data: cpu_nv12.y_data.clone(), uv_data: cpu_nv12.uv_data.clone(), @@ -632,11 +676,20 @@ fn webrtc_thread_loop( mut wrtc: WebRtcState, webrtc_rx: crossbeam_channel::Receiver>, fps: u32, + enc_width: u32, + enc_height: u32, paused: Arc, sent_gap_tx: crossbeam_channel::Sender, + bitrate_tx: crossbeam_channel::Sender, + resolution_tx: crossbeam_channel::Sender, ) { let mut frames_sent: u64 = 0; let mut last_send: Option = None; + let mut last_sent_bitrate: Option = None; + let initial_tier = (enc_width, enc_height); + let mut current_tier = initial_tier; + let mut upscale_counter = 0u32; + let mut last_resolution_eval = Instant::now(); let timeout = Duration::from_millis(1); loop { @@ -659,6 +712,52 @@ fn webrtc_thread_loop( } paused.store(now_paused, Ordering::Relaxed); + if let Some(bwe) = wrtc.get_bwe_estimate() { + let should_send = match last_sent_bitrate { + None => true, + Some(last) => { + let diff = if bwe > last { bwe - last } else { last - bwe }; + diff * 10 > last + } + }; + if should_send { + let _ = bitrate_tx.try_send(BitrateCommand::UpdateBitrate { target_bps: bwe }); + last_sent_bitrate = Some(bwe); + } + + if last_resolution_eval.elapsed() >= Duration::from_secs(1) { + last_resolution_eval = Instant::now(); + let selected = select_resolution(current_tier.0, current_tier.1, bwe, fps); + if selected != current_tier { + current_tier = selected; + upscale_counter = 0; + let _ = resolution_tx.try_send(BitrateCommand::UpdateResolution { + width: current_tier.0, + height: current_tier.1, + }); + wrtc.set_need_keyframe(); + } else if let Some(next_tier) = next_upscale_tier(current_tier, initial_tier) { + let needed = resolution_bitrate_bps(next_tier.0, next_tier.1, fps); + if bwe > needed.saturating_mul(120) / 100 { + upscale_counter = upscale_counter.saturating_add(1); + if upscale_counter >= 10 { + current_tier = next_tier; + upscale_counter = 0; + let _ = resolution_tx.try_send(BitrateCommand::UpdateResolution { + width: current_tier.0, + height: current_tier.1, + }); + wrtc.set_need_keyframe(); + } + } else { + upscale_counter = 0; + } + } else { + upscale_counter = 0; + } + } + } + if connected { while let Ok(data) = webrtc_rx.try_recv() { if let Err(e) = wrtc.write_h264_frame(&data, frames_sent, fps) { @@ -700,6 +799,43 @@ fn webrtc_thread_loop( tracing::info!("WebRTC thread exiting"); } +const RESOLUTION_TIERS: &[(u32, u32)] = &[(2560, 1440), (1920, 1080), (1280, 720)]; + +fn resolution_bitrate_bps(width: u32, height: u32, fps: u32) -> u64 { + 5 * u64::from(width) * u64::from(height) * u64::from(fps) / 100 +} + +/// Select resolution tier based on BWE estimate. +/// Returns (width, height) for the selected tier. +fn select_resolution(current_w: u32, current_h: u32, bwe_bps: u64, fps: u32) -> (u32, u32) { + let current = (current_w, current_h); + let current_bitrate = resolution_bitrate_bps(current_w, current_h, fps); + if bwe_bps >= current_bitrate.saturating_mul(60) / 100 { + return current; + } + + let current_index = RESOLUTION_TIERS + .iter() + .position(|&tier| tier == current) + .unwrap_or_else(|| { + RESOLUTION_TIERS + .iter() + .position(|&(w, h)| w <= current_w && h <= current_h) + .unwrap_or(RESOLUTION_TIERS.len() - 1) + }); + let next_index = (current_index + 1).min(RESOLUTION_TIERS.len() - 1); + RESOLUTION_TIERS[next_index] +} + +fn next_upscale_tier(current: (u32, u32), ceiling: (u32, u32)) -> Option<(u32, u32)> { + let current_index = RESOLUTION_TIERS.iter().position(|&tier| tier == current)?; + if current_index == 0 { + return None; + } + let next = RESOLUTION_TIERS[current_index - 1]; + (next.0 <= ceiling.0 && next.1 <= ceiling.1).then_some(next) +} + impl Drop for StatePortal { // 析构时自动调用 shutdown，确保编码器被刷新、资源被释放 fn drop(&mut self) { @@ -712,8 +848,8 @@ impl Drop for StatePortal { /// 将原始分辨率等比缩放至不超过 2560×1440（2K），并确保宽高为偶数 /// （H.264 编码要求偶数尺寸）。 fn portal_encode_dimensions(width: u32, height: u32) -> (u32, u32) { - const TARGET_W: u32 = 2560; // 目标最大宽度 - const TARGET_H: u32 = 1440; // 目标最大高度 + const TARGET_W: u32 = 2560; // 目标最大宽度 + const TARGET_H: u32 = 1440; // 目标最大高度 // 原始分辨率已在 2K 以内，直接对齐偶数 if width <= TARGET_W && height <= TARGET_H { @@ -748,16 +884,16 @@ fn resolve_drm_device(args: &Args) -> Result> { #[cfg(test)] fn build_drm_descriptor(frame: &PwDmaBufFrame) -> ffmpeg_next::ffi::AVDRMFrameDescriptor { let mut desc: ffmpeg_next::ffi::AVDRMFrameDescriptor = unsafe { std::mem::zeroed() }; - desc.nb_objects = 1; // 单个 DMA-BUF 对象 - desc.objects[0].fd = frame.fd.as_raw_fd(); // DMA-BUF 文件描述符 - desc.objects[0].size = 0; // 大小设为 0（内核自动确定） + desc.nb_objects = 1; // 单个 DMA-BUF 对象 + desc.objects[0].fd = frame.fd.as_raw_fd(); // DMA-BUF 文件描述符 + desc.objects[0].size = 0; // 大小设为 0（内核自动确定） desc.objects[0].format_modifier = frame.modifier; // DRM 格式修饰符（如线性、tiled） - desc.nb_layers = 1; // 单层 - desc.layers[0].format = frame.format; // 像素格式（如 XR24） - desc.layers[0].nb_planes = 1; // 单平面 - desc.layers[0].planes[0].object_index = 0; // 指向第 0 个对象 - desc.layers[0].planes[0].offset = frame.offset as isize; // 帧数据偏移 - desc.layers[0].planes[0].pitch = frame.stride as isize; // 行跨度（stride） + desc.nb_layers = 1; // 单层 + desc.layers[0].format = frame.format; // 像素格式（如 XR24） + desc.layers[0].nb_planes = 1; // 单平面 + desc.layers[0].planes[0].object_index = 0; // 指向第 0 个对象 + desc.layers[0].planes[0].offset = frame.offset as isize; // 帧数据偏移 + desc.layers[0].planes[0].pitch = frame.stride as isize; // 行跨度（stride） desc } @@ -774,11 +910,11 @@ mod tests { PwDmaBufFrame { fd, offset: 0, - stride: 1920 * 4, // 每行 1920 像素 × 4 字节（XRGB） - modifier: 0, // DRM_FORMAT_MOD_LINEAR（线性布局） + stride: 1920 * 4, // 每行 1920 像素 × 4 字节（XRGB） + modifier: 0, // DRM_FORMAT_MOD_LINEAR（线性布局） width: 1920, height: 1080, - format: 0x34325258, // XR24 little-endian（XRGB8888） + format: 0x34325258, // XR24 little-endian（XRGB8888） pts: 12345, } } @@ -845,14 +981,48 @@ mod tests { assert_eq!(result, None); } + #[test] + fn select_resolution_downscales_one_tier_below_sixty_percent() { + let fps = 30; + let current = resolution_bitrate_bps(1920, 1080, fps); + assert_eq!( + select_resolution(1920, 1080, current * 59 / 100, fps), + (1280, 720) + ); + } + + #[test] + fn select_resolution_keeps_tier_at_sixty_percent() { + let fps = 30; + let current = resolution_bitrate_bps(1920, 1080, fps); + assert_eq!( + select_resolution(1920, 1080, current * 60 / 100, fps), + (1920, 1080) + ); + } + + #[test] + fn select_resolution_never_goes_below_720p() { + assert_eq!(select_resolution(1280, 720, 1, 30), (1280, 720)); + } + + #[test] + fn next_upscale_tier_respects_initial_ceiling() { + assert_eq!( + next_upscale_tier((1280, 720), (1920, 1080)), + Some((1920, 1080)) + ); + assert_eq!(next_upscale_tier((1920, 1080), (1920, 1080)), None); + } + /// 测试：使用自定义偏移量和 stride 构建 DRM 描述符 #[test] fn build_drm_descriptor_custom_offset_and_stride() { let frame = PwDmaBufFrame { fd: unsafe { OwnedFd::from_raw_fd(libc::dup(2)) }, - offset: 4096, // 4KB 对齐偏移 - stride: 3840 * 4, // 4K 宽度 × 4 字节 - modifier: 0x0100000000000001, // AMD modifiers + offset: 4096, // 4KB 对齐偏移 + stride: 3840 * 4, // 4K 宽度 × 4 字节 + modifier: 0x0100000000000001, // AMD modifiers width: 3840, height: 2160, format: 0x34325258, @@ -908,4 +1078,48 @@ mod tests { )); } + // ── Task 7: Additional resolution tier edge cases ── + + #[test] + fn select_resolution_keeps_720p_when_bwe_sufficient() { + let fps = 30; + let bitrate_720 = resolution_bitrate_bps(1280, 720, fps); + assert_eq!( + select_resolution(1280, 720, bitrate_720, fps), + (1280, 720) + ); + } + + #[test] + fn select_resolution_downscales_1440p_to_1080p() { + let fps = 30; + let bitrate_1440 = resolution_bitrate_bps(2560, 1440, fps); + assert_eq!( + select_resolution(2560, 1440, bitrate_1440 * 59 / 100, fps), + (1920, 1080) + ); + } + + #[test] + fn select_resolution_1080p_to_720p_at_very_low_bwe() { + let fps = 30; + let bitrate_1080 = resolution_bitrate_bps(1920, 1080, fps); + assert_eq!( + select_resolution(1920, 1080, bitrate_1080 / 10, fps), + (1280, 720) + ); + } + + #[test] + fn next_upscale_tier_from_720p_to_1080p() { + assert_eq!( + next_upscale_tier((1280, 720), (2560, 1440)), + Some((1920, 1080)) + ); + } + + #[test] + fn next_upscale_tier_returns_none_at_highest() { + assert_eq!(next_upscale_tier((2560, 1440), (2560, 1440)), None); + } } diff --git a/src/stats.rs b/src/stats.rs index d10ac99..1e9d196 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -461,14 +461,20 @@ mod tests { // 100 values: 0.0 through 99.0 let data: Vec = (0..100).map(|i| i as f64).collect(); let result = p95_f64(&data); - assert!((result - 95.0).abs() < 1.0, "p95 of 0..100 should be ~95, got {result}"); + assert!( + (result - 95.0).abs() < 1.0, + "p95 of 0..100 should be ~95, got {result}" + ); } #[test] fn p95_ms_microseconds() { let data: Vec = (0..100).map(|i| i * 1000).collect(); // 0ms..99ms let result = p95_ms(&data); - assert!((result - 95.0).abs() < 1.0, "p95_ms should be ~95ms, got {result}"); + assert!( + (result - 95.0).abs() < 1.0, + "p95_ms should be ~95ms, got {result}" + ); } #[test] diff --git a/src/webrtc.rs b/src/webrtc.rs index 6f8ba18..c072318 100644 --- a/src/webrtc.rs +++ b/src/webrtc.rs @@ -4,6 +4,7 @@ use std::net::{SocketAddr, TcpListener, UdpSocket}; use std::time::Instant; use anyhow::{bail, Result}; +use str0m::bwe::{Bitrate, BweKind}; use str0m::change::SdpOffer; use str0m::format::Codec; use str0m::media::{Frequency, MediaKind, MediaTime, Mid, Pt}; @@ -199,6 +200,7 @@ struct WebRtcInner { video_pt: Option, connected: bool, need_keyframe: bool, + current_bwe_estimate: Option, rtp_clock: u32, buf: Vec, } @@ -260,11 +262,10 @@ impl WebRtcState { continue; } - match WebRtcInner::new(self.fps) - .and_then(|mut new_inner| { - let answer_json = new_inner.handle_sdp_offer(body.as_bytes())?; - Ok((new_inner, answer_json)) - }) { + match WebRtcInner::new(self.fps).and_then(|mut new_inner| { + let answer_json = new_inner.handle_sdp_offer(body.as_bytes())?; + Ok((new_inner, answer_json)) + }) { Ok((new_inner, answer_json)) => { let replacing = self.inner.is_some(); self.inner = Some(new_inner); @@ -285,7 +286,8 @@ impl WebRtcState { } Err(e) => { tracing::error!("SDP offer handling failed: {e}"); - let resp = "HTTP/1.1 500 Internal Server Error\r\nConnection: close\r\n\r\n"; + let resp = + "HTTP/1.1 500 Internal Server Error\r\nConnection: close\r\n\r\n"; if let Err(e) = stream.write_all(resp.as_bytes()) { tracing::debug!("HTTP write error: {e}"); } @@ -340,12 +342,27 @@ impl WebRtcState { pub fn is_connected(&self) -> bool { self.inner.as_ref().is_some_and(WebRtcInner::is_connected) } + + /// Returns the latest bandwidth estimation estimate in bits per second, if available. + pub fn get_bwe_estimate(&self) -> Option { + self.inner + .as_ref() + .and_then(|inner| inner.current_bwe_estimate.map(|b| b.as_u64())) + } + + pub fn set_need_keyframe(&mut self) { + if let Some(inner) = self.inner.as_mut() { + inner.need_keyframe = true; + } + } } impl WebRtcInner { fn new(fps: u32) -> Result { let _ = fps; - let mut rtc = RtcConfig::new().build(Instant::now()); + let mut rtc = RtcConfig::new() + .enable_bwe(Some(Bitrate::mbps(5))) + .build(Instant::now()); let socket = UdpSocket::bind("0.0.0.0:0")?; socket.set_nonblocking(true)?; @@ -368,10 +385,14 @@ impl WebRtcInner { std::mem::size_of::() as libc::socklen_t, ); if ret < 0 { - tracing::warn!("setsockopt SO_SNDBUF failed (errno {})", std::io::Error::last_os_error()); + tracing::warn!( + "setsockopt SO_SNDBUF failed (errno {})", + std::io::Error::last_os_error() + ); } let mut actual: libc::c_int = 0; - let mut actual_len: libc::socklen_t = std::mem::size_of::() as libc::socklen_t; + let mut actual_len: libc::socklen_t = + std::mem::size_of::() as libc::socklen_t; let gret = libc::getsockopt( fd, libc::SOL_SOCKET, @@ -408,14 +429,15 @@ impl WebRtcInner { video_pt: None, connected: false, need_keyframe: false, + current_bwe_estimate: None, rtp_clock: 0, buf: vec![0u8; 65535], }) } fn handle_sdp_offer(&mut self, body: &[u8]) -> Result { - let offer: SdpOffer = serde_json::from_slice(body) - .map_err(|e| anyhow::anyhow!("parse SDP offer: {e}"))?; + let offer: SdpOffer = + serde_json::from_slice(body).map_err(|e| anyhow::anyhow!("parse SDP offer: {e}"))?; let answer = self .rtc @@ -492,9 +514,7 @@ impl WebRtcInner { tracing::info!("Media added: mid={} kind={:?}", ma.mid, ma.kind); if ma.kind == MediaKind::Video { if let Some(media) = self.rtc.media(ma.mid) { - if media.direction().is_sending() - && self.video_mid.is_none() - { + if media.direction().is_sending() && self.video_mid.is_none() { self.video_mid = Some(ma.mid); tracing::info!("Captured video mid: {}", ma.mid); self.discover_video_params(); @@ -502,6 +522,22 @@ impl WebRtcInner { } } } + Event::KeyframeRequest(_) => { + tracing::info!("received keyframe request from viewer"); + self.need_keyframe = true; + } + Event::EgressBitrateEstimate(est) => { + let bitrate = match est { + BweKind::Twcc(b) => *b, + BweKind::Remb(_, b) => *b, + _ => { + tracing::debug!("BWE estimate: unrecognized kind, skipping"); + continue; + } + }; + tracing::info!("BWE estimate: {bitrate}"); + self.current_bwe_estimate = Some(bitrate); + } _ => { tracing::debug!("WebRTC event: {:?}", e); } @@ -538,9 +574,9 @@ impl WebRtcInner { .map_err(|e| anyhow::anyhow!("receive contents: {e}"))?, }, ); - self.rtc - .handle_input(input) - .map_err(|e| anyhow::anyhow!("handle_input({n} bytes from {source}): {e}"))?; + self.rtc.handle_input(input).map_err(|e| { + anyhow::anyhow!("handle_input({n} bytes from {source}): {e}") + })?; } Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break, Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue, @@ -636,18 +672,16 @@ fn extract_body(req: &str) -> &str { } fn local_ip() -> Option { - std::net::UdpSocket::bind("0.0.0.0:0") - .ok() - .and_then(|s| { - s.connect("1.1.1.1:80").ok()?; - let addr = s.local_addr().ok()?; - drop(s); - let ip = addr.ip().to_string(); - if ip == "0.0.0.0" || ip.starts_with("127.") { - return None; - } - Some(ip) - }) + std::net::UdpSocket::bind("0.0.0.0:0").ok().and_then(|s| { + s.connect("1.1.1.1:80").ok()?; + let addr = s.local_addr().ok()?; + drop(s); + let ip = addr.ip().to_string(); + if ip == "0.0.0.0" || ip.starts_with("127.") { + return None; + } + Some(ip) + }) } fn is_idr_nalu(data: &[u8]) -> bool { @@ -746,4 +780,26 @@ mod tests { fn all_zeros() { assert!(!is_idr_nalu(&[0, 0, 0, 0, 0, 0, 0, 0])); } + + // ── Task 5: BWE estimate handling ── + + #[test] + fn bitrate_mbps_conversion() { + let b = Bitrate::mbps(5); + assert!(b.as_u64() > 0, "5 Mbps should be > 0 bps"); + assert_eq!(b.as_u64(), 5_000_000, "5 Mbps should be 5,000,000 bps"); + } + + #[test] + fn bitrate_bps_conversion() { + let b = Bitrate::bps(1_234_567); + assert_eq!(b.as_u64(), 1_234_567); + } + + #[test] + fn default_bwe_is_5mbps() { + let default = Bitrate::mbps(5); + let bps = default.as_u64(); + assert_eq!(bps, 5_000_000); + } }