shutdown() fired twice on exit (explicit call in main.rs + Drop impl), causing "Total: N frames" and "StatePortal shutdown complete" to log twice 23us apart. Add `shutdown_started: bool` guard at function entry. Plain bool (not AtomicBool) because &mut self already grants exclusive access. Guard is set BEFORE cleanup so Drop re-entry during panic unwinding is suppressed. Includes scripts/test_shutdown_idempotency.sh as a live regression test (requires Wayland session). Verified PASS on KWin: both lines print exactly once.
1136 lines
43 KiB
Rust
1136 lines
43 KiB
Rust
// 采集门户状态模块 —— 通过 PipeWire/DMA-BUF 进行屏幕采集并编码
|
||
use std::os::fd::AsRawFd;
|
||
use std::path::PathBuf;
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
use std::sync::Arc;
|
||
use std::time::{Duration, Instant};
|
||
|
||
use anyhow::{bail, Result}; // 错误处理工具
|
||
|
||
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 信令与媒体传输
|
||
|
||
/// 门户采集的阶段状态
|
||
/// - WaitingForFormat: 等待接收到第一帧 DMA-BUF 以确定视频格式参数
|
||
/// - Streaming: 已完成初始化,正在持续编码流
|
||
enum PortalStage {
|
||
WaitingForFormat,
|
||
Streaming,
|
||
}
|
||
|
||
struct EncodeThreadTiming {
|
||
sws_us: u64,
|
||
encode_us: u64,
|
||
output_bytes: usize,
|
||
}
|
||
|
||
struct EncodeThread {
|
||
handle: Option<std::thread::JoinHandle<()>>,
|
||
input_tx: crossbeam_channel::Sender<CpuNv12Frame>,
|
||
timing_rx: crossbeam_channel::Receiver<EncodeThreadTiming>,
|
||
}
|
||
|
||
struct WebrtcThread {
|
||
handle: Option<std::thread::JoinHandle<()>>,
|
||
sent_gap_rx: crossbeam_channel::Receiver<f64>,
|
||
}
|
||
|
||
/// 门户模式的主状态机
|
||
///
|
||
/// 负责管理从 PipeWire 采集屏幕帧、通过 VAAPI 硬件编码的完整生命周期。
|
||
/// 工作流程:等待第一帧 → 创建编码器 → 持续编码帧数据。
|
||
pub struct StatePortal {
|
||
stage: PortalStage, // 当前采集阶段(等待首帧 / 流式编码中)
|
||
enc: Option<SwEncState>, // 软件编码器,首帧到达后初始化
|
||
enc_import: Option<SwEncImport>,
|
||
enc_thread: Option<EncodeThread>,
|
||
cap: CapPortal, // PipeWire 屏幕采集端点
|
||
args: Args, // 用户命令行参数
|
||
errored: bool, // 是否遇到不可恢复的错误
|
||
drm_device: Option<PathBuf>, // DRM 渲染设备路径(可自动检测)
|
||
frames_encoded: u64, // 已编码帧数(用于 PTS 编号)
|
||
start_time: Option<Instant>, // 编码开始时间
|
||
stats: PipelineStats, // 管道统计(窗口化帧计时 + 每秒快照)
|
||
pw_dropped_prev: u64, // 上一窗口的 PipeWire 丢弃帧数(用于增量计算)
|
||
webrtc: Option<WebRtcState>,
|
||
webrtc_thread: Option<WebrtcThread>,
|
||
webrtc_paused: Option<Arc<AtomicBool>>,
|
||
last_capture_arrival: Option<Instant>, // timestamp of last real frame arrival
|
||
stall_start: Option<Instant>, // when current stall began
|
||
last_stall_log: Option<Instant>, // rate-limiting for stall warnings
|
||
last_fillable_frame: Option<CpuNv12Frame>, // cached last frame for filler duplication
|
||
next_filler_at: Option<Instant>, // when to send next filler frame
|
||
filler_frames_sent: u64,
|
||
shutdown_started: bool, // idempotency guard; plain bool because &mut self is exclusive (not AtomicBool)
|
||
}
|
||
|
||
impl StatePortal {
|
||
/// 创建门户状态实例
|
||
///
|
||
/// 初始化 DRM 设备路径和 PipeWire 采集端点,编码器延迟到第一帧到达时创建。
|
||
pub fn new(args: Args) -> Result<Self> {
|
||
let drm_device = resolve_drm_device(&args)?;
|
||
if let Some(ref drm_device) = drm_device {
|
||
tracing::info!("Using DRM device: {}", drm_device.display());
|
||
} else {
|
||
tracing::info!("DRM device auto-detection enabled");
|
||
}
|
||
|
||
let cap = CapPortal::new(&args)?;
|
||
|
||
let (webrtc, webrtc_paused) = if args.port > 0 {
|
||
let wrtc = WebRtcState::new(args.port, args.fps)?;
|
||
let paused = Arc::new(AtomicBool::new(true));
|
||
(Some(wrtc), Some(paused))
|
||
} else {
|
||
(None, None)
|
||
};
|
||
|
||
Ok(Self {
|
||
stage: PortalStage::WaitingForFormat,
|
||
enc: None,
|
||
enc_import: None,
|
||
enc_thread: None,
|
||
cap,
|
||
args,
|
||
errored: false,
|
||
drm_device,
|
||
frames_encoded: 0,
|
||
start_time: None,
|
||
stats: PipelineStats::new(),
|
||
pw_dropped_prev: 0,
|
||
webrtc,
|
||
webrtc_thread: None,
|
||
webrtc_paused,
|
||
last_capture_arrival: None,
|
||
stall_start: None,
|
||
last_stall_log: None,
|
||
last_fillable_frame: None,
|
||
next_filler_at: None,
|
||
filler_frames_sent: 0,
|
||
shutdown_started: false,
|
||
})
|
||
}
|
||
|
||
/// 轮询 PipeWire 事件并编码帧
|
||
///
|
||
/// `block=true` 时使用 recv_timeout 阻塞等待帧(最多 2ms),
|
||
/// `block=false` 时使用 try_recv 非阻塞检查。
|
||
/// 返回 `Ok(true)` 表示已处理事件,`Ok(false)` 表示暂无数据。
|
||
pub fn poll_and_encode(&mut self, block: bool) -> Result<bool> {
|
||
// 检查 PipeWire 控制事件(流结束 / 错误)
|
||
if let Ok(ctrl) = self.cap.event_receiver().try_recv() {
|
||
match ctrl {
|
||
PwCtrlEvent::StreamEnded => {
|
||
tracing::warn!("PipeWire stream ended");
|
||
self.errored = true;
|
||
return Ok(true);
|
||
}
|
||
PwCtrlEvent::Error(e) => {
|
||
tracing::error!("PipeWire error: {e}");
|
||
self.errored = true;
|
||
return Ok(true);
|
||
}
|
||
PwCtrlEvent::FormatChanged { width, height } => {
|
||
tracing::warn!(
|
||
"PipeWire format renegotiation: new dimensions {}x{} — encoder output remains at original resolution",
|
||
width,
|
||
height
|
||
);
|
||
// No action yet — VAAPI import/scale handles the conversion.
|
||
// Full encoder reinit is a future enhancement.
|
||
}
|
||
}
|
||
}
|
||
|
||
// 根据阻塞模式选择不同的帧接收策略
|
||
let frame = if block {
|
||
// 阻塞模式:最多等待 2ms 接收帧
|
||
match self
|
||
.cap
|
||
.frame_receiver()
|
||
.recv_timeout(std::time::Duration::from_millis(2))
|
||
{
|
||
Ok(frame) => frame,
|
||
Err(_) => {
|
||
self.record_capture_timeout();
|
||
return Ok(false);
|
||
}
|
||
}
|
||
} else {
|
||
// 非阻塞模式:立即尝试接收,无数据则返回
|
||
match self.cap.frame_receiver().try_recv() {
|
||
Ok(frame) => frame,
|
||
Err(_) => {
|
||
self.record_capture_timeout();
|
||
return Ok(false);
|
||
}
|
||
}
|
||
};
|
||
self.record_frame_arrival();
|
||
|
||
match self.stage {
|
||
PortalStage::WaitingForFormat => {
|
||
// 首帧到达,记录 DMA-BUF 格式信息
|
||
tracing::info!(
|
||
"First DMA-BUF frame: {}x{} format=0x{:08X} stride={} modifier=0x{:X}",
|
||
frame.width,
|
||
frame.height,
|
||
frame.format,
|
||
frame.stride,
|
||
frame.modifier
|
||
);
|
||
|
||
// 自动检测或确认 DRM 设备是否支持导入该帧
|
||
let drm_path = self.resolve_drm_device_for_frame(&frame)?;
|
||
// 计算编码目标分辨率(不超过 2560x1440)
|
||
let (enc_width, enc_height) = portal_encode_dimensions(frame.width, frame.height);
|
||
tracing::info!(
|
||
"Portal software encode target: {}x{} -> {}x{} @ {} fps",
|
||
frame.width,
|
||
frame.height,
|
||
enc_width,
|
||
enc_height,
|
||
self.args.fps,
|
||
);
|
||
// 码率:未指定时按分辨率 × 帧率动态计算
|
||
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,最低20),MP4 模式使用 fps
|
||
let actual_gop_size = self.args.gop_size.unwrap_or_else(|| {
|
||
if self.webrtc.is_some() {
|
||
(self.args.fps * 2).max(20)
|
||
} else {
|
||
self.args.fps
|
||
}
|
||
});
|
||
|
||
// 根据是否启用 WebRTC 选择不同的编码器构造方式
|
||
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 (resolution_tx, resolution_rx) =
|
||
crossbeam_channel::bounded::<BitrateCommand>(4);
|
||
let (encoder_resolution_tx, encoder_resolution_rx) =
|
||
crossbeam_channel::bounded::<ResolutionChange>(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::<CpuNv12Frame>(1);
|
||
let (timing_tx, timing_rx) =
|
||
crossbeam_channel::bounded::<EncodeThreadTiming>(32);
|
||
let (bitrate_tx, bitrate_rx) = crossbeam_channel::bounded::<BitrateCommand>(4);
|
||
let encode = SwEncEncode::new_webrtc(
|
||
enc_width,
|
||
enc_height,
|
||
self.args.fps,
|
||
actual_bitrate,
|
||
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,
|
||
});
|
||
|
||
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,
|
||
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()
|
||
.ok_or_else(|| anyhow::anyhow!("--output is required in MP4 file output mode; use --port > 0 for WebRTC mode"))?;
|
||
let enc = avhw::SwEncState::new(
|
||
&drm_path,
|
||
std::path::Path::new(output_path),
|
||
frame.width,
|
||
frame.height,
|
||
enc_width,
|
||
enc_height,
|
||
self.args.fps,
|
||
actual_bitrate,
|
||
actual_gop_size,
|
||
)?;
|
||
self.enc = Some(enc);
|
||
};
|
||
self.stage = PortalStage::Streaming; // 切换到流式编码阶段
|
||
self.start_time = Some(Instant::now());
|
||
tracing::info!(
|
||
"First frame processed, encoder initialized, transitioning to Streaming"
|
||
);
|
||
drop(frame); // 首帧仅用于初始化,不参与编码
|
||
}
|
||
PortalStage::Streaming => {
|
||
// 记录采集帧到达(用于 capture gap 和 capture_fps 统计)
|
||
self.stats.record_capture();
|
||
self.last_capture_arrival = Some(Instant::now());
|
||
// 流式编码阶段:直接处理帧
|
||
self.handle_pw_frame(frame)?;
|
||
}
|
||
}
|
||
|
||
// 每秒输出一次结构化管道统计(仅 --stats 启用时记录日志)
|
||
if self.args.stats && self.stats.should_snapshot() {
|
||
self.stats.set_pipewire_dropped(0, 0);
|
||
self.stats.set_queue_depths(0, 0);
|
||
if let Some(ref enc_thread) = self.enc_thread {
|
||
while let Ok(timing) = enc_thread.timing_rx.try_recv() {
|
||
self.stats.record_encode_thread(
|
||
timing.sws_us,
|
||
timing.encode_us,
|
||
timing.output_bytes,
|
||
);
|
||
}
|
||
}
|
||
if let Some(ref webrtc_thread) = self.webrtc_thread {
|
||
while let Ok(gap_ms) = webrtc_thread.sent_gap_rx.try_recv() {
|
||
self.stats.record_send_from_thread(gap_ms);
|
||
}
|
||
}
|
||
let snap = self.stats.snapshot_and_reset();
|
||
if self.filler_frames_sent > 0 {
|
||
tracing::info!(
|
||
"stats: {snap} filler_frames_sent={}",
|
||
self.filler_frames_sent
|
||
);
|
||
} else {
|
||
tracing::info!("stats: {snap}");
|
||
}
|
||
}
|
||
|
||
Ok(true)
|
||
}
|
||
|
||
fn record_capture_timeout(&mut self) {
|
||
let Some(last_capture_arrival) = self.last_capture_arrival else {
|
||
return;
|
||
};
|
||
|
||
let now = Instant::now();
|
||
let frame_interval = Duration::from_secs_f64(1.0 / f64::from(self.args.fps.max(1)));
|
||
let stall_threshold = Duration::from_millis(100).max(frame_interval * 3);
|
||
if now.duration_since(last_capture_arrival) <= stall_threshold {
|
||
return;
|
||
}
|
||
|
||
if self.stall_start.is_none() {
|
||
self.stall_start = Some(now);
|
||
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)
|
||
});
|
||
if should_log {
|
||
self.last_stall_log = Some(now);
|
||
tracing::warn!("compositor frame delivery stalled");
|
||
}
|
||
}
|
||
|
||
self.maybe_send_filler_frame();
|
||
}
|
||
|
||
fn maybe_send_filler_frame(&mut self) {
|
||
if self.webrtc_thread.is_none() || self.stall_start.is_none() {
|
||
return;
|
||
}
|
||
let Some(cached) = &self.last_fillable_frame else {
|
||
return;
|
||
};
|
||
|
||
const MAX_FILLER_DURATION: Duration = Duration::from_secs(2);
|
||
if let Some(stall_start) = self.stall_start {
|
||
if stall_start.elapsed() > MAX_FILLER_DURATION {
|
||
return;
|
||
}
|
||
}
|
||
|
||
let now = Instant::now();
|
||
let frame_interval = Duration::from_secs_f64(1.0 / f64::from(self.args.fps.max(1)));
|
||
|
||
let Some(next) = self.next_filler_at else {
|
||
self.next_filler_at = Some(now + frame_interval);
|
||
return;
|
||
};
|
||
if now < next {
|
||
return;
|
||
}
|
||
|
||
let filler = CpuNv12Frame {
|
||
y_data: cached.y_data.clone(),
|
||
uv_data: cached.uv_data.clone(),
|
||
y_stride: cached.y_stride,
|
||
uv_stride: cached.uv_stride,
|
||
pts: self.frames_encoded as i64,
|
||
};
|
||
|
||
if let Some(enc_thread) = &self.enc_thread {
|
||
match enc_thread.input_tx.try_send(filler) {
|
||
Ok(()) => {
|
||
self.frames_encoded += 1;
|
||
self.filler_frames_sent += 1;
|
||
self.next_filler_at = Some(next + frame_interval);
|
||
}
|
||
Err(crossbeam_channel::TrySendError::Full(_)) => {}
|
||
Err(crossbeam_channel::TrySendError::Disconnected(_)) => {
|
||
tracing::error!("Encode thread disconnected during filler");
|
||
self.errored = true;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
fn record_frame_arrival(&mut self) {
|
||
if let Some(stall_start) = self.stall_start.take() {
|
||
tracing::info!(
|
||
"compositor frame delivery resumed after {:.0}ms",
|
||
stall_start.elapsed().as_secs_f64() * 1000.0
|
||
);
|
||
self.last_stall_log = None;
|
||
}
|
||
self.last_capture_arrival = Some(Instant::now());
|
||
self.next_filler_at = None;
|
||
}
|
||
|
||
/// 为当前帧解析可用的 DRM 渲染设备
|
||
///
|
||
/// 如果用户已通过 `--drm-device` 指定设备,直接返回;
|
||
/// 否则遍历系统中所有 DRM render node,逐个尝试导入 DMA-BUF 帧来找到兼容设备。
|
||
fn resolve_drm_device_for_frame(&mut self, frame: &PwDmaBufFrame) -> Result<PathBuf> {
|
||
// 用户已显式指定 DRM 设备,直接使用
|
||
if let Some(ref drm) = self.drm_device {
|
||
return Ok(drm.clone());
|
||
}
|
||
|
||
// 查找系统中所有 DRM render node(如 /dev/dri/renderD128)
|
||
let candidates = crate::state::find_drm_render_nodes();
|
||
if candidates.is_empty() {
|
||
bail!("No DRM render device found. Specify --drm-device.");
|
||
}
|
||
|
||
// 逐个尝试导入 DMA-BUF 帧,找到第一个兼容的设备
|
||
let mut failures = Vec::new();
|
||
for candidate in &candidates {
|
||
match crate::avhw::test_dma_buf_import(candidate, frame) {
|
||
Ok(()) => {
|
||
// 成功导入,缓存检测结果并返回
|
||
tracing::info!(
|
||
"Auto-detected DRM device: {} (tested {} candidates)",
|
||
candidate.display(),
|
||
candidates.len(),
|
||
);
|
||
self.drm_device = Some(candidate.clone());
|
||
return Ok(candidate.clone());
|
||
}
|
||
Err(e) => {
|
||
// 导入失败,记录原因,继续尝试下一个设备
|
||
tracing::debug!(
|
||
"DRM device {} cannot import DMA-BUF: {e}",
|
||
candidate.display(),
|
||
);
|
||
failures.push((candidate, e));
|
||
}
|
||
}
|
||
}
|
||
|
||
// 所有候选设备均失败,返回详细错误信息
|
||
bail!(failures
|
||
.into_iter()
|
||
.map(|(p, e)| format!("{} ({e})", p.display()))
|
||
.collect::<Vec<_>>()
|
||
.join(", "));
|
||
}
|
||
|
||
/// 处理单帧 DMA-BUF 数据
|
||
///
|
||
/// 通过 `av_hwframe_map` 零拷贝导入 VAAPI,然后交给 SwEncState 完成:
|
||
/// scale_vaapi GPU 缩放、2K NV12 回读、YUV420P 格式转换、软件 H.264 编码。
|
||
fn handle_pw_frame(&mut self, frame: PwDmaBufFrame) -> Result<()> {
|
||
let t_import_start = Instant::now();
|
||
let pts = self.frames_encoded as i64;
|
||
|
||
if let Some(enc) = self.enc.as_mut() {
|
||
// 将 DMA-BUF 帧零拷贝导入 VAAPI 硬件帧池
|
||
let mut vaapi_frame = unsafe {
|
||
avhw::import_dma_buf_to_vaapi(
|
||
enc.frames_rgb().as_ptr(),
|
||
frame.fd.as_raw_fd(),
|
||
frame.width,
|
||
frame.height,
|
||
frame.format,
|
||
frame.modifier,
|
||
frame.stride,
|
||
frame.offset,
|
||
)
|
||
}?;
|
||
|
||
let import_us = t_import_start.elapsed().as_micros() as u64;
|
||
let t_encode_start = Instant::now();
|
||
|
||
// 设置帧的显示时间戳(PTS),基于已编码帧序号
|
||
unsafe {
|
||
(*vaapi_frame.as_mut_ptr()).pts = pts;
|
||
}
|
||
|
||
// 送入编码器完成:缩放 → 回读 → 格式转换 → H.264 编码
|
||
enc.encode_frame(&vaapi_frame)?;
|
||
let total_us = t_import_start.elapsed().as_micros() as u64;
|
||
let encode_us = t_encode_start.elapsed().as_micros() as u64;
|
||
|
||
self.frames_encoded += 1;
|
||
|
||
// 记录帧计时到管道统计(import + encode 内部各阶段暂不可分离,用 total 覆盖)
|
||
let timings = FrameTimings {
|
||
import_us,
|
||
encode_us,
|
||
total_us,
|
||
..Default::default()
|
||
};
|
||
self.stats.record_encode(&timings);
|
||
} else if let Some(import) = self.enc_import.as_mut() {
|
||
let mut vaapi_frame = unsafe {
|
||
avhw::import_dma_buf_to_vaapi(
|
||
import.frames_rgb().as_ptr(),
|
||
frame.fd.as_raw_fd(),
|
||
frame.width,
|
||
frame.height,
|
||
frame.format,
|
||
frame.modifier,
|
||
frame.stride,
|
||
frame.offset,
|
||
)
|
||
}?;
|
||
unsafe {
|
||
(*vaapi_frame.as_mut_ptr()).pts = pts;
|
||
}
|
||
|
||
let cpu_nv12 = import.import_and_scale(&vaapi_frame)?;
|
||
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 fillable_frame = CpuNv12Frame {
|
||
y_data: cpu_nv12.y_data.clone(),
|
||
uv_data: cpu_nv12.uv_data.clone(),
|
||
y_stride: cpu_nv12.y_stride,
|
||
uv_stride: cpu_nv12.uv_stride,
|
||
pts: 0,
|
||
};
|
||
match enc_thread.input_tx.try_send(cpu_nv12) {
|
||
Ok(()) => {
|
||
self.frames_encoded += 1;
|
||
self.last_fillable_frame = Some(fillable_frame);
|
||
}
|
||
Err(crossbeam_channel::TrySendError::Full(_)) => {
|
||
tracing::debug!("Encode thread input full, dropping portal frame");
|
||
}
|
||
Err(crossbeam_channel::TrySendError::Disconnected(_frame)) => {
|
||
tracing::error!("Encode thread input disconnected");
|
||
self.errored = true;
|
||
}
|
||
}
|
||
} else {
|
||
bail!("encoder not initialized");
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// 关闭状态:刷新编码器并清理资源(幂等)。
|
||
///
|
||
/// `shutdown_started` 守卫在清理之前置位——防止 panic 时 `Drop` 重入 unwinding。
|
||
pub fn shutdown(&mut self) {
|
||
if self.shutdown_started {
|
||
return;
|
||
}
|
||
self.shutdown_started = true;
|
||
|
||
self.last_fillable_frame = None;
|
||
// 1. Stop encode thread (drops webrtc_tx → signals WebRTC thread to exit)
|
||
if let Some(mut enc_thread) = self.enc_thread.take() {
|
||
drop(enc_thread.input_tx);
|
||
if let Some(handle) = enc_thread.handle.take() {
|
||
if handle.join().is_err() {
|
||
tracing::error!("Encode thread panicked during shutdown");
|
||
}
|
||
}
|
||
}
|
||
self.enc_import = None;
|
||
// 2. Wait for WebRTC thread (exits when webrtc_tx is dropped by encode thread)
|
||
if let Some(mut webrtc_thread) = self.webrtc_thread.take() {
|
||
if let Some(handle) = webrtc_thread.handle.take() {
|
||
if handle.join().is_err() {
|
||
tracing::error!("WebRTC thread panicked during shutdown");
|
||
}
|
||
}
|
||
}
|
||
// 3. Flush MP4 encoder if present
|
||
if let Some(mut enc) = self.enc.take() {
|
||
if let Err(e) = enc.flush() {
|
||
tracing::error!("Flush error during shutdown: {e}");
|
||
}
|
||
}
|
||
if let Some(start) = self.start_time {
|
||
if self.frames_encoded > 0 {
|
||
let elapsed = start.elapsed().as_secs_f64();
|
||
let fps = self.frames_encoded as f64 / elapsed;
|
||
tracing::info!(
|
||
"Total: {} frames in {:.1}s, avg {:.1}fps",
|
||
self.frames_encoded,
|
||
elapsed,
|
||
fps,
|
||
);
|
||
}
|
||
}
|
||
tracing::info!("StatePortal shutdown complete");
|
||
}
|
||
|
||
/// 返回是否遇到不可恢复的错误
|
||
pub fn is_errored(&self) -> bool {
|
||
self.errored
|
||
}
|
||
}
|
||
|
||
fn encode_thread_loop(
|
||
mut encode: SwEncEncode,
|
||
input_rx: crossbeam_channel::Receiver<CpuNv12Frame>,
|
||
timing_tx: crossbeam_channel::Sender<EncodeThreadTiming>,
|
||
) {
|
||
loop {
|
||
match input_rx.recv() {
|
||
Ok(frame) => {
|
||
match encode.encode_cpu_frame(&frame) {
|
||
Ok(()) => {
|
||
let t = encode.take_timing();
|
||
let _ = timing_tx.try_send(EncodeThreadTiming {
|
||
sws_us: t.sws_us,
|
||
encode_us: t.encode_us,
|
||
output_bytes: t.output_bytes,
|
||
});
|
||
}
|
||
Err(e) => {
|
||
tracing::error!("Encode thread error: {e}");
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
Err(_) => {
|
||
tracing::info!("Encode thread input closed, flushing encoder");
|
||
if let Err(e) = encode.flush() {
|
||
tracing::error!("Encode thread flush error: {e}");
|
||
}
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
tracing::info!("Encode thread exiting");
|
||
}
|
||
|
||
fn webrtc_thread_loop(
|
||
mut wrtc: WebRtcState,
|
||
webrtc_rx: crossbeam_channel::Receiver<Vec<u8>>,
|
||
fps: u32,
|
||
enc_width: u32,
|
||
enc_height: u32,
|
||
paused: Arc<AtomicBool>,
|
||
sent_gap_tx: crossbeam_channel::Sender<f64>,
|
||
bitrate_tx: crossbeam_channel::Sender<BitrateCommand>,
|
||
resolution_tx: crossbeam_channel::Sender<BitrateCommand>,
|
||
) {
|
||
let mut frames_sent: u64 = 0;
|
||
let mut last_send: Option<std::time::Instant> = None;
|
||
let mut last_sent_bitrate: Option<u64> = 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 {
|
||
if let Err(e) = wrtc.handle_signaling() {
|
||
tracing::error!("WebRTC signaling error: {e}");
|
||
break;
|
||
}
|
||
if let Err(e) = wrtc.poll_and_feed() {
|
||
tracing::error!("WebRTC poll error: {e}");
|
||
break;
|
||
}
|
||
|
||
if wrtc.take_force_keyframe() {
|
||
let _ = bitrate_tx.try_send(BitrateCommand::ForceKeyframe);
|
||
}
|
||
|
||
let connected = wrtc.is_connected();
|
||
let was_paused = paused.load(Ordering::Relaxed);
|
||
let now_paused = !connected;
|
||
if was_paused && !now_paused {
|
||
tracing::info!("WebRTC client connected, resuming encoding");
|
||
} else if !was_paused && now_paused {
|
||
tracing::warn!("WebRTC client disconnected, pausing encoding");
|
||
}
|
||
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) {
|
||
tracing::debug!("WebRTC write frame error: {e}");
|
||
}
|
||
frames_sent = frames_sent.saturating_add(1);
|
||
let gap_ms = last_send
|
||
.map(|l| l.elapsed().as_secs_f64() * 1000.0)
|
||
.unwrap_or(0.0);
|
||
last_send = Some(std::time::Instant::now());
|
||
let _ = sent_gap_tx.try_send(gap_ms);
|
||
}
|
||
} else {
|
||
while webrtc_rx.try_recv().is_ok() {}
|
||
}
|
||
|
||
match webrtc_rx.recv_timeout(timeout) {
|
||
Ok(data) => {
|
||
if wrtc.is_connected() {
|
||
if let Err(e) = wrtc.write_h264_frame(&data, frames_sent, fps) {
|
||
tracing::debug!("WebRTC write frame error: {e}");
|
||
}
|
||
frames_sent = frames_sent.saturating_add(1);
|
||
let gap_ms = last_send
|
||
.map(|l| l.elapsed().as_secs_f64() * 1000.0)
|
||
.unwrap_or(0.0);
|
||
last_send = Some(std::time::Instant::now());
|
||
let _ = sent_gap_tx.try_send(gap_ms);
|
||
}
|
||
}
|
||
Err(crossbeam_channel::RecvTimeoutError::Timeout) => {}
|
||
Err(crossbeam_channel::RecvTimeoutError::Disconnected) => {
|
||
tracing::info!("WebRTC channel disconnected, exiting thread");
|
||
return;
|
||
}
|
||
}
|
||
}
|
||
|
||
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) {
|
||
self.shutdown();
|
||
}
|
||
}
|
||
|
||
/// 计算编码目标分辨率
|
||
///
|
||
/// 将原始分辨率等比缩放至不超过 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; // 目标最大高度
|
||
|
||
// 原始分辨率已在 2K 以内,直接对齐偶数
|
||
if width <= TARGET_W && height <= TARGET_H {
|
||
return (width & !1, height & !1); // & !1 确保为偶数
|
||
}
|
||
|
||
// 按宽度限制等比缩放
|
||
let width_limited_h = ((height as u64) * (TARGET_W as u64) / (width as u64)) as u32;
|
||
if width_limited_h <= TARGET_H {
|
||
(TARGET_W & !1, width_limited_h & !1)
|
||
} else {
|
||
// 按高度限制等比缩放
|
||
let height_limited_w = ((width as u64) * (TARGET_H as u64) / (height as u64)) as u32;
|
||
(height_limited_w & !1, TARGET_H & !1)
|
||
}
|
||
}
|
||
|
||
/// 解析 DRM 渲染设备路径
|
||
///
|
||
/// 仅使用命令行指定的设备路径;未指定则在首帧到达时自动检测。
|
||
fn resolve_drm_device(args: &Args) -> Result<Option<PathBuf>> {
|
||
if let Some(ref drm) = args.drm_device {
|
||
return Ok(Some(PathBuf::from(drm)));
|
||
}
|
||
Ok(None)
|
||
}
|
||
|
||
/// 构建测试用的 AVDRMFrameDescriptor(仅测试用途)
|
||
///
|
||
/// 将 PwDmaBufFrame 转换为 FFmpeg 的 DRM 帧描述符结构体,
|
||
/// 用于验证 DMA-BUF 元数据映射的正确性。
|
||
#[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.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
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::os::fd::{FromRawFd, OwnedFd};
|
||
|
||
/// 创建测试用的 DMA-BUF 帧数据(使用 stderr fd 的副本作为占位)
|
||
fn make_test_frame() -> PwDmaBufFrame {
|
||
// Create a dummy fd from stderr (always valid fd 2)
|
||
// 使用 stderr(fd 2)的副本作为虚拟文件描述符
|
||
let fd = unsafe { OwnedFd::from_raw_fd(libc::dup(2)) };
|
||
PwDmaBufFrame {
|
||
fd,
|
||
offset: 0,
|
||
stride: 1920 * 4, // 每行 1920 像素 × 4 字节(XRGB)
|
||
modifier: 0, // DRM_FORMAT_MOD_LINEAR(线性布局)
|
||
width: 1920,
|
||
height: 1080,
|
||
format: 0x34325258, // XR24 little-endian(XRGB8888)
|
||
pts: 12345,
|
||
}
|
||
}
|
||
|
||
/// 测试 DRM 描述符构建(单平面情况)
|
||
#[test]
|
||
fn build_drm_descriptor_single_plane() {
|
||
let frame = make_test_frame();
|
||
let desc = build_drm_descriptor(&frame);
|
||
|
||
assert_eq!(desc.nb_objects, 1);
|
||
assert_eq!(desc.objects[0].format_modifier, 0);
|
||
assert_eq!(desc.nb_layers, 1);
|
||
assert_eq!(desc.layers[0].format, 0x34325258);
|
||
assert_eq!(desc.layers[0].nb_planes, 1);
|
||
assert_eq!(desc.layers[0].planes[0].object_index, 0);
|
||
assert_eq!(desc.layers[0].planes[0].offset, 0);
|
||
assert_eq!(desc.layers[0].planes[0].pitch, 1920 * 4);
|
||
}
|
||
|
||
/// 测试显式指定 DRM 设备时的解析
|
||
#[test]
|
||
fn resolve_drm_device_explicit() {
|
||
let args = Args {
|
||
output: Some("test.mp4".to_string()),
|
||
output_name: None,
|
||
fps: 30,
|
||
codec: "h264".to_string(),
|
||
hw_accel: "vaapi".to_string(),
|
||
drm_device: Some("/dev/dri/renderD128".to_string()),
|
||
bitrate: None,
|
||
gop_size: None,
|
||
verbose: false,
|
||
backend: None,
|
||
port: 0,
|
||
no_persist: false,
|
||
stats: false,
|
||
};
|
||
let result = resolve_drm_device(&args).unwrap();
|
||
assert_eq!(
|
||
result,
|
||
Some(std::path::PathBuf::from("/dev/dri/renderD128"))
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn resolve_drm_device_none_when_not_specified() {
|
||
let args = Args {
|
||
output: Some("test.mp4".to_string()),
|
||
output_name: None,
|
||
fps: 30,
|
||
codec: "h264".to_string(),
|
||
hw_accel: "vaapi".to_string(),
|
||
drm_device: None,
|
||
bitrate: None,
|
||
gop_size: None,
|
||
verbose: false,
|
||
backend: None,
|
||
port: 0,
|
||
no_persist: false,
|
||
stats: false,
|
||
};
|
||
let result = resolve_drm_device(&args).unwrap();
|
||
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
|
||
width: 3840,
|
||
height: 2160,
|
||
format: 0x34325258,
|
||
pts: 0,
|
||
};
|
||
let desc = build_drm_descriptor(&frame);
|
||
|
||
assert_eq!(desc.nb_objects, 1);
|
||
assert_eq!(desc.objects[0].format_modifier, 0x0100000000000001);
|
||
assert_eq!(desc.layers[0].planes[0].offset, 4096);
|
||
assert_eq!(desc.layers[0].planes[0].pitch, 3840 * 4);
|
||
}
|
||
|
||
// ── issue #8 regression ──
|
||
|
||
#[test]
|
||
fn try_send_full_channel_returns_full_not_block() {
|
||
let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
|
||
tx.send(vec![1]).unwrap();
|
||
tx.send(vec![2]).unwrap();
|
||
|
||
assert!(matches!(
|
||
tx.try_send(vec![3]),
|
||
Err(crossbeam_channel::TrySendError::Full(_))
|
||
));
|
||
assert_eq!(rx.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn try_send_after_rx_dropped_returns_disconnected() {
|
||
let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
|
||
drop(rx);
|
||
|
||
assert!(matches!(
|
||
tx.try_send(vec![1]),
|
||
Err(crossbeam_channel::TrySendError::Disconnected(_))
|
||
));
|
||
}
|
||
|
||
// given: full bounded channel
|
||
// when: rx is dropped, then try_send
|
||
// expect: Disconnected, not blocking
|
||
#[test]
|
||
fn shutdown_rx_drop_prevents_deadlock_on_full_channel() {
|
||
let (tx, rx) = crossbeam_channel::bounded::<Vec<u8>>(2);
|
||
tx.send(vec![1]).unwrap();
|
||
tx.send(vec![2]).unwrap();
|
||
drop(rx);
|
||
|
||
assert!(matches!(
|
||
tx.try_send(vec![3]),
|
||
Err(crossbeam_channel::TrySendError::Disconnected(_))
|
||
));
|
||
}
|
||
|
||
// ── 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);
|
||
}
|
||
}
|