// 获取 Unix 原始文件描述符所需的 trait use std::os::unix::io::AsRawFd; use anyhow::Result; use clap::Parser; use mio::unix::SourceFd; use mio::{Events, Interest, Poll, Token}; use wayland_client::globals::registry_queue_init; use wayland_client::Connection; // 各功能模块声明 mod args; // 命令行参数解析 mod avhw; // 音视频硬件加速 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 transform; // 图像变换(旋转/翻转) mod webrtc; // WebRTC 传输(str0m Sans-IO) use crate::args::Args; use crate::cap_wlr_screencopy::CapWlrScreencopy; use crate::state::EncConstructionStage; use crate::state::State; // mio 事件循环的 Token 标识:0 = Wayland 合成器事件,1 = 退出信号 const TOKEN_WAYLAND: Token = Token(0); const TOKEN_QUIT: Token = Token(1); /// 程序入口:解析参数 → 初始化日志 → 检测后端 → 启动对应的事件循环 /// /// 整体流程: /// 1. 解析命令行参数(分辨率、编码格式、帧率等) /// 2. 初始化日志系统(verbose 模式输出 DEBUG 级别,否则 INFO) /// 3. 检查编码格式(MVP 阶段仅支持 H.264) /// 4. 自动检测当前桌面环境支持的截屏后端 /// 5. 根据检测结果启动对应的事件循环: /// - wlroots 合成器(Sway/Hyprland)→ run_wlr_screencopy /// - GNOME/KDE 等 → run_portal_pipewire fn main() -> Result<()> { // 解析命令行参数 let args = Args::parse(); // 根据 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") } }); 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); // MVP 阶段仅支持 H.264 编码,不支持 HEVC if args.codec != "h264" { anyhow::bail!("HEVC not supported in MVP. Use --codec h264"); } if args.output.is_none() && args.port == 0 { anyhow::bail!("Either --output or --port is required"); } // 自动检测当前桌面环境可用的截屏后端 // 会尝试列举 Wayland 全局对象,判断合成器是否支持 wlr-screencopy 协议 let backend = crate::backend_detect::detect_backend(&args)?; // 根据检测结果进入对应的事件循环 match backend { crate::backend_detect::CaptureBackend::WlrScreencopy => run_wlr_screencopy(args), crate::backend_detect::CaptureBackend::PortalPipeWire => run_portal_pipewire(args), } } /// 使用 wlroots wlr-screencopy 协议的事件循环(适用于 Sway、Hyprland 等 wlroots 合成器) /// /// 完整工作流程: /// 1. 连接 Wayland 合成器,获取全局注册表 /// 2. 初始化 wlr-screencopy 截屏状态机 /// 3. 获取 Wayland socket 的文件描述符 /// 4. 使用 mio 注册 fd 监听(Wayland 事件 + Unix 信号) /// 5. 进入主事件循环: /// - 监听合成器事件(帧就绪通知、输出信息等) /// - 定时请求截屏帧 /// - 将截取的帧编码为 H.264 并推流 /// 6. 收到退出信号或发生错误时,刷新编码器并退出 fn run_wlr_screencopy(args: Args) -> Result<()> { // Connect to Wayland compositor // 建立 Wayland 连接并初始化全局注册表 // 通过环境变量 $WAYLAND_DISPLAY 找到合成器的 Unix socket let conn = Connection::connect_to_env()?; // registry_queue_init 会绑定全局注册表回调, // 当合成器广播其全局对象(输出、截屏管理器等)时,State 会收到通知 let (gm, mut queue) = registry_queue_init::>(&conn)?; let qhandle = queue.handle(); // State 是 wlr-screencopy 后端的核心状态机, // 内部管理输出探测、截屏请求、编码器构建、帧采集等阶段 let mut state = State::new(gm, args, qhandle)?; // Extract the Wayland fd and consume any immediately-available events. // prepare_read() flushes outgoing requests; read() pulls whatever the // compositor has already sent (may be EAGAIN if nothing yet). // 获取 Wayland socket 的文件描述符,并消费合成器已发送的事件 // 这个 fd 是后续 mio epoll 监听的对象,当合成器写入数据时变为可读 let wayland_fd = { let guard = queue .prepare_read() .ok_or_else(|| anyhow::anyhow!("Failed to prepare Wayland read"))?; // 从 prepare_read 的 guard 中获取底层 socket 的原始文件描述符 let fd = guard.connection_fd().as_raw_fd(); // 尝试非阻塞读取合成器已发送但尚未消费的数据 // 如果没有数据会返回 EAGAIN,这里用 let _ 忽略 let _ = guard.read(); fd }; // 处理队列中的待处理事件 // 在进入主循环前,先处理注册表广播等初始化事件 // 此时状态机应进入 ProbingOutputs 阶段,正在探测可用的显示输出 queue.dispatch_pending(&mut state)?; tracing::info!( "Initial dispatch done, stage is ProbingOutputs: {}", matches!(state.stage, EncConstructionStage::ProbingOutputs { .. }) ); // 调试用:对 Wayland fd 做一次原始 poll,确认 fd 可读性 // 使用 libc 底层 poll 而非 mio,纯粹用于诊断初始化阶段的 fd 状态 { let mut pfd = libc::pollfd { fd: wayland_fd, events: libc::POLLIN, // 监听可读事件 revents: 0, }; // timeout=0 表示非阻塞,立即返回当前 fd 状态 let ret = unsafe { libc::poll(&mut pfd, 1, 0) }; tracing::info!( "Raw poll on wayland fd={wayland_fd}: ret={ret}, revents={}", pfd.revents ); } // Set up mio event loop // 使用 mio 创建事件循环,注册 Wayland fd 和 Unix 信号 // mio 底层在 Linux 上使用 epoll,macOS 上使用 kqueue let mut poll = Poll::new()?; // 事件缓冲区,容量 8 足够(实际只会同时处理 Wayland 事件和信号两种) let mut events = Events::with_capacity(8); // 将 Wayland socket fd 注册为可读监听 // 当合成器发送消息(如帧完成通知、配置变化)时,epoll 会唤醒 poll.registry().register( &mut SourceFd(&wayland_fd), TOKEN_WAYLAND, Interest::READABLE, )?; // 注册 SIGINT / SIGTERM 信号用于优雅退出 // signal_hook_mio 将 Unix 信号转换为 fd 可读事件, // 这样信号也可以通过 epoll 统一监听,不需要单独的信号处理器 let mut signals = signal_hook_mio::v1_0::Signals::new(&[ signal_hook::consts::SIGINT, // Ctrl+C signal_hook::consts::SIGTERM, // kill 命令默认信号 ])?; poll.registry() .register(&mut signals, TOKEN_QUIT, Interest::READABLE)?; tracing::info!("Event loop started"); // Flush outgoing before first poll iteration // 在首次 poll 前刷新所有待发送的 Wayland 请求 // 确保合成器能收到我们的初始化请求(如绑定全局对象、请求截屏等) conn.flush()?; // 主事件循环 // 这是 wlr-screencopy 后端的核心运行循环,负责: // - 接收合成器事件(截屏帧就绪、输出变化) // - 定时触发帧采集和编码 // - 响应退出信号 let mut running = true; while running { // 准备读取 Wayland 事件(非阻塞) // prepare_read() 会先刷出所有待发送的请求, // 然后进入"准备读取"状态,告诉合成器我们已准备好接收数据 let read_guard = queue.prepare_read(); // 如果无法 prepare_read(有待处理数据),先分发 // 返回 None 说明队列中已有待处理的合成器事件, // 需要先 dispatch 掉,否则新事件无法进入队列 if read_guard.is_none() { queue.dispatch_pending(&mut state)?; } // 阻塞等待事件,超时 100ms(用于帧率控制) // poll 会阻塞当前线程,直到以下任一条件满足: // 1. Wayland fd 可读(合成器发来了消息) // 2. 信号 fd 可读(收到了 SIGINT/SIGTERM) // 3. 超过 100ms 没有任何事件(超时返回,触发下一帧采集) // 100ms 超时 ≈ 10 FPS 的帧率上限 poll.poll(&mut events, Some(std::time::Duration::from_millis(100))) .unwrap_or_else(|e| { // EINTR 是信号中断,属于正常情况,继续循环 // 当进程收到信号时,阻塞中的 poll 会被中断并返回 EINTR, // 这不是错误,下一轮循环会继续正常 poll if e.kind() == std::io::ErrorKind::Interrupted { return; } tracing::error!("poll failed: {e}"); running = false; }); // 检查是否收到退出信号 for event in &events { if event.token() == TOKEN_QUIT { tracing::info!("Received quit signal"); running = false; } } // Wayland fd 可读时,读取并分发合成器事件 // 合成器可能发来多种事件:帧数据就绪、输出信息变化、协议错误等 if events.iter().any(|e| e.token() == TOKEN_WAYLAND) { if let Some(guard) = read_guard { match guard.read() { Ok(_) => { // 读取成功后,dispatch_pending 会将合成器事件 // 分发给 State 的对应回调方法处理 queue.dispatch_pending(&mut state)?; } Err(e) => { tracing::error!("Wayland read error: {e}"); running = false; } } } } // 请求状态机分配并编码一帧 // queue_alloc_frame 会根据当前状态机阶段执行不同操作: // - ProbingOutputs: 还在探测输出,跳过 // - AwaitingFrame: 已请求截屏,等待合成器回调 // - FrameReady: 有帧就绪,执行 DMA-BUF → H.264 编码 → 推流 // - Streaming: 正常采集中,请求下一帧 state.queue_alloc_frame(); state.poll_webrtc()?; // 状态机遇到致命错误时退出 if state.errored { tracing::error!("Fatal error in state machine (check preceding error logs), exiting"); running = false; } // 每轮循环结束前刷新 Wayland 发送缓冲区 // 将本轮回合中产生的所有 Wayland 请求(如截屏请求)发送给合成器 conn.flush()?; } // 关闭前刷新编码器,确保所有帧数据已写出 // 先通知帧率限制器停止,再刷新编码器缓冲区中残余的帧数据 tracing::info!("Shutting down, flushing encoder..."); state.fps_limit.flush(); // 仅在编码器已构建完成(Streaming 阶段)时才需要刷新 if let crate::state::EncConstructionStage::Streaming { enc, .. } = &mut state.stage { if let Err(e) = enc.flush() { tracing::error!("Failed to flush encoder: {e}"); } } tracing::info!("Done"); Ok(()) } /// 使用 XDG Portal / PipeWire 后端的事件循环(适用于 KWin、KDE、GNOME 等桌面环境) /// /// 完整工作流程: /// 1. 初始化 Portal 状态机(内部通过 D-Bus 与 XDG Portal 通信) /// 2. 仅注册 Unix 信号监听(不需要监听 Wayland fd,帧数据由 PipeWire 投递) /// 3. 进入主事件循环: /// - 每 10ms 轮询一次 PipeWire 缓冲区 /// - 有帧数据时,取出并编码为 H.264 推流 /// - 收到退出信号时停止 /// 4. 退出时关闭 Portal 连接并释放 PipeWire 资源 fn run_portal_pipewire(args: Args) -> Result<()> { use crate::state_portal::StatePortal; tracing::info!("Using Portal/PipeWire backend (KWin/KDE/GNOME)"); // StatePortal 初始化时会: // 1. 通过 D-Bus 连接到 XDG Portal 的 ScreenCast 接口 // 2. 请求用户授权屏幕录制权限 // 3. 建立 PipeWire 流连接,准备接收帧数据 let mut state = StatePortal::new(args)?; // Set up signal handling only (no Wayland fd needed) // Portal 后端不需要监听 Wayland fd,只需处理 Unix 信号 // 因为帧数据是通过 PipeWire 独立投递的,不走 Wayland 协议 let mut signals = signal_hook_mio::v1_0::Signals::new(&[ signal_hook::consts::SIGINT, signal_hook::consts::SIGTERM, ])?; let mut poll = mio::Poll::new()?; let mut events = mio::Events::with_capacity(8); // 只注册信号 fd,没有 Wayland fd // 所以 poll.poll 在这里只负责检测 SIGINT/SIGTERM // 实际的帧采集完全依赖 poll_and_encode 的轮询 poll.registry() .register(&mut signals, mio::Token(1), mio::Interest::READABLE)?; // 主事件循环(非阻塞信号检测 + recv_timeout 等待帧) // poll 超时为 0ms(非阻塞),实际等待由 poll_and_encode 的 recv_timeout 实现 let mut running = true; while running { // poll 在此循环中只监听信号 fd(非阻塞): // - 收到 SIGINT/SIGTERM → 事件触发,设置 running=false // - 无事件 → 立即返回,继续执行 poll_and_encode(内部 recv_timeout 等待帧) poll.poll(&mut events, Some(std::time::Duration::from_millis(0))) .unwrap_or_else(|e| { if e.kind() == std::io::ErrorKind::Interrupted { return; } tracing::error!("poll failed: {e}"); running = false; }); for event in &events { if event.token() == mio::Token(1) { tracing::info!("Received quit signal"); running = false; } } // Process all available PipeWire frames // 处理所有可用的 PipeWire 帧数据 // poll_and_encode 会从 PipeWire 缓冲区取出帧, // 编码为 H.264 并推送。返回 true 表示还有更多帧待处理, // 返回 false 表示当前没有帧了,while 循环退出等待下一轮 poll if state.poll_and_encode(true)? { while state.poll_and_encode(false)? {} } // Portal 状态机遇到致命错误时退出 if state.is_errored() { tracing::error!("Fatal error in portal state machine (check preceding error logs), exiting"); running = false; } } tracing::info!("Shutting down..."); // 关闭 Portal 连接,释放 PipeWire 流和编码器资源 state.shutdown(); tracing::info!("Done"); Ok(()) }