feat(agent): 串行多会话——会话结束不再退出进程

问题(阶段 2 的硬阻塞)
accept() 只接一次,会话结束进程退出。阶段 1 一次验收看不出来,
但阶段 2 要反复测试输入,每测一次都要手动重启被控端。
(即公共阻碍清单里的 P5,此前判断为「只影响调试效率」,
 现在它阻塞迭代,提前做掉。)

改动
1. main:会话体抽成 serve_session(),外层 accept 循环。
   会话失败(含认证失败)只记日志不退出进程。
   明确不做并发多会话:DXGI Desktop Duplication 每个输出仅允许一个
   活动实例,第二个 DuplicateOutput 会让第一个收到 ACCESS_LOST。
2. 认证失败路径补 conn.closed() 等待(agent-core 已有同款修复,
   此处同源补齐 = 手册 §6.1 第 1 条):否则客户端只看到 connection
   lost 而不是 AuthFail。
3. capture 循环加 tx.receiver_count()==0 守卫。
   只靠「send 失败退出」不够——桌面完全静止时那个分支根本不发送,
   线程会永远持有 duplication,下一个会话的 DuplicateOutput 直接失败
   (表现为新会话黑屏)。agent-core 同源修复。
4. 会话收尾:finish 流 + drop 接收端 + 等采集线程退出(2s 超时)后才
   返回,调用方才 accept 下一个连接。

验证(连续两个会话,中间不重启 agent)
  会话1:收到帧 40/40,结论 ✅
  agent 存活:✅
  会话2:收到帧 40/40,结论 ✅
  日志:session closed -> waiting for next connection -> peer connected
        -> auth ok;stderr 有 [capture] no receiver, stop 证明线程真退出
This commit is contained in:
郑博文 2026-10-30 09:31:44 +08:00
parent 58bb66f4c2
commit f119300a57
3 changed files with 60 additions and 8 deletions

View File

@ -192,6 +192,13 @@ fn capture_session(
let mut seq: u64 = 0;
loop {
// 会话结束守卫:发送端已 drop 接收器时立即退出。
// 与 agent/src/capture.rs 同源——静止桌面分支不发送,只靠 send 失败退出
// 会让线程在画面不动时永不退出(X11 无 DXGI 式独占,危害较小但同样泄漏线程)。
if tx.receiver_count() == 0 {
eprintln!("[capture] no receiver, stop");
return LoopEnd::ChannelClosed;
}
let iter_start = Instant::now();
match cap.acquire(80) {
Ok(true) => match cap.copy_frame() {

View File

@ -387,6 +387,14 @@ pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>, congestion: Arc
}
loop {
// 会话结束守卫:发送端已 drop 接收器时立即退出。
// 只靠 `tx.send` 失败退出是不够的——桌面完全静止时这里根本不发送,
// 线程会永远持有 duplication,导致下一个会话的 DuplicateOutput 失败
// (DXGI 每个输出仅允许一个活动实例)。
if tx.receiver_count() == 0 {
eprintln!("[capture] no receiver, stop");
return;
}
let iter_start = Instant::now();
match dup.acquire(80) {
Ok(true) => {

View File

@ -1,6 +1,8 @@
//! 被控端 Agent:DXGI 桌面复制采集 → LZ4 压缩 → QUIC 下行视频;上行接收输入事件注入。
//! 被控端 Agent:DXGI 桌面复制采集 → H.264 编码 → QUIC 下行视频;上行接收输入事件注入。
//! P0 优化:视频帧经 watch 单槽通道(latest-wins,网络慢时只发最新帧,不积压);
//! 键鼠/控制走独立 QUIC 双向流,不与视频互相阻塞。
//! 会话模型:**串行多会话**——accept 循环一次只服务一个主控,会话结束后回到 accept
//! 等下一个连接(进程不退出)。改成并发多会话需要先解决 DXGI 单实例独占,暂不做。
mod capture;
mod input;
@ -10,6 +12,7 @@ use common::{read_control, write_control, ControlMsg};
use rustls::pki_types::PrivatePkcs8KeyDer;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
use tokio::sync::{mpsc, watch};
@ -59,13 +62,31 @@ async fn main() -> Result<()> {
)?;
println!("[agent] listening 0.0.0.0:{port}, verify code = {code}");
let conn = endpoint
.accept()
.await
.ok_or_else(|| anyhow::anyhow!("endpoint closed"))?
.await?;
println!("[agent] peer connected: {}", conn.remote_address());
// 串行多会话:一次只服务一个主控,结束后回到 accept。
// 不做并发多会话——DXGI Desktop Duplication 每个输出仅允许一个活动实例,
// 第二个 DuplicateOutput 会让第一个收到 ACCESS_LOST。
// 会话失败(含认证失败)只记日志,不能让进程退出——否则每次连接都要重启被控端。
loop {
let Some(incoming) = endpoint.accept().await else {
println!("[agent] endpoint closed, exit");
return Ok(());
};
match incoming.await {
Ok(conn) => {
println!("[agent] peer connected: {}", conn.remote_address());
if let Err(e) = serve_session(conn, &code).await {
println!("[agent] session ended: {e:#}");
}
println!("[agent] waiting for next connection");
}
Err(e) => println!("[agent] handshake failed: {e}"),
}
}
}
/// 单次会话:认证 → 采集/发送/输入 → 会话结束。
/// 返回 Err 表示会话异常结束(含认证失败),调用方继续 accept。
async fn serve_session(conn: quinn::Connection, code: &str) -> Result<()> {
// 流 1:认证 + 视频下行
let (mut send, mut recv) = conn.accept_bi().await?;
@ -93,6 +114,10 @@ async fn main() -> Result<()> {
if let Some(msg) = auth_err {
write_control(&mut send, &ControlMsg::AuthFail(msg.clone())).await?;
send.finish()?;
// 必须等对端读到 AuthFail 再关连接:直接返回会让进程带走 QUIC 缓冲区里
// 尚未发出的数据,客户端只会看到 "connection lost" 而不是 AuthFail。
// (agent-core 已有同款修复,此处同源补齐——手册 §6.1 第 1 条)
let _ = tokio::time::timeout(Duration::from_secs(5), conn.closed()).await;
anyhow::bail!(msg);
}
let (w, h) = capture::desktop_size()?;
@ -117,7 +142,14 @@ async fn main() -> Result<()> {
// 拥塞反馈标志:发送侧 write 超过单帧预算(20ms) → 置位,采集线程跳帧降注入率
let congestion = Arc::new(AtomicBool::new(false));
let congestion_capture = congestion.clone();
std::thread::spawn(move || capture::run(tx, keyreq_capture, congestion_capture));
// 采集线程退出信号:会话结束必须等它真正退出才能 accept 下一个连接——
// DXGI Desktop Duplication 每个输出只允许一个活动实例,旧线程还持有 duplication
// 时新线程的 DuplicateOutput 会让旧实例收到 ACCESS_LOST(表现为下一个会话黑屏)。
let (done_tx, done_rx) = std::sync::mpsc::channel::<()>();
std::thread::spawn(move || {
capture::run(tx, keyreq_capture, congestion_capture);
let _ = done_tx.send(());
});
// 上行控制消息(Pong)经通道交给主循环写流 2,避免多任务争抢
let (ctl_tx, mut ctl_rx) = mpsc::unbounded_channel::<ControlMsg>();
@ -182,6 +214,11 @@ async fn main() -> Result<()> {
}
}
reader.abort();
// 关掉流并丢弃接收端,让采集线程下一次 send 失败而退出
let _ = send.finish();
drop(frames_rx);
// 等采集线程真正退出后再返回,调用方才会 accept 下一个连接(DXGI 独占,见上)
let _ = tokio::task::spawn_blocking(move || done_rx.recv_timeout(Duration::from_secs(2))).await;
println!("[agent] session closed");
Ok(())
}