257 lines
10 KiB
Rust
257 lines
10 KiB
Rust
//! DXGI Desktop Duplication 采集（Windows 8+）。
|
||
//! 帧处理：BGRA → openh264 H.264 编码（屏幕内容实时档，5Mbps）→ 线格式。
|
||
//! 支持主控端 KeyFrameReq 强制关键帧（解码异常恢复）。
|
||
|
||
use anyhow::{anyhow, Result};
|
||
use common::{encode_video_frame, VideoFrameHeader};
|
||
use openh264::encoder::{
|
||
BitRate, Encoder, EncoderConfig, FrameRate, IntraFramePeriod, RateControlMode, UsageType,
|
||
};
|
||
use openh264::formats::{BgraSliceU8, YUVBuffer};
|
||
use openh264::OpenH264API;
|
||
use std::sync::atomic::{AtomicBool, Ordering};
|
||
use std::sync::Arc;
|
||
use tokio::sync::watch::Sender;
|
||
use windows::core::Interface;
|
||
use windows::Win32::Foundation::HMODULE;
|
||
use windows::Win32::Graphics::Direct3D::{
|
||
D3D_DRIVER_TYPE_HARDWARE, D3D_FEATURE_LEVEL_11_0,
|
||
};
|
||
use windows::Win32::Graphics::Direct3D11::{
|
||
D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Resource, ID3D11Texture2D,
|
||
D3D11_CREATE_DEVICE_FLAG, D3D11_CPU_ACCESS_READ, D3D11_MAP, D3D11_MAP_READ,
|
||
D3D11_MAPPED_SUBRESOURCE, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
|
||
};
|
||
use windows::Win32::Graphics::Dxgi::Common::{DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_SAMPLE_DESC};
|
||
use windows::Win32::Graphics::Dxgi::{
|
||
IDXGIDevice, IDXGIOutput1, IDXGIOutputDuplication, IDXGIResource, DXGI_ERROR_WAIT_TIMEOUT,
|
||
DXGI_OUTDUPL_DESC, DXGI_OUTDUPL_FRAME_INFO,
|
||
};
|
||
|
||
use std::time::{Duration, Instant};
|
||
|
||
pub struct Duplication {
|
||
_device: ID3D11Device,
|
||
context: ID3D11DeviceContext,
|
||
staging_res: ID3D11Resource,
|
||
dup: IDXGIOutputDuplication,
|
||
pub width: u32,
|
||
pub height: u32,
|
||
}
|
||
|
||
impl Duplication {
|
||
pub fn new() -> Result<Self> {
|
||
unsafe {
|
||
let mut device: Option<ID3D11Device> = None;
|
||
let mut context: Option<ID3D11DeviceContext> = None;
|
||
D3D11CreateDevice(
|
||
None,
|
||
D3D_DRIVER_TYPE_HARDWARE,
|
||
HMODULE::default(),
|
||
D3D11_CREATE_DEVICE_FLAG(0),
|
||
Some(&[D3D_FEATURE_LEVEL_11_0]),
|
||
D3D11_SDK_VERSION,
|
||
Some(&mut device),
|
||
None,
|
||
Some(&mut context),
|
||
)
|
||
.map_err(|e| anyhow!("D3D11CreateDevice failed: {e}"))?;
|
||
let device = device.ok_or_else(|| anyhow!("no d3d device"))?;
|
||
let context = context.ok_or_else(|| anyhow!("no d3d context"))?;
|
||
|
||
let dxgi: IDXGIDevice = device.cast()?;
|
||
let adapter = dxgi.GetAdapter()?;
|
||
let output = adapter
|
||
.EnumOutputs(0)
|
||
.map_err(|e| anyhow!("no display output: {e}"))?;
|
||
let output1: IDXGIOutput1 = output.cast()?;
|
||
let dup: IDXGIOutputDuplication = output1
|
||
.DuplicateOutput(&device)
|
||
.map_err(|e| anyhow!("DuplicateOutput failed: {e}"))?;
|
||
|
||
let desc: DXGI_OUTDUPL_DESC = dup.GetDesc();
|
||
let (width, height) = (desc.ModeDesc.Width, desc.ModeDesc.Height);
|
||
anyhow::ensure!(width > 0 && height > 0, "bad desktop size");
|
||
|
||
let sdesc = D3D11_TEXTURE2D_DESC {
|
||
Width: width,
|
||
Height: height,
|
||
MipLevels: 1,
|
||
ArraySize: 1,
|
||
Format: DXGI_FORMAT_B8G8R8A8_UNORM,
|
||
SampleDesc: DXGI_SAMPLE_DESC {
|
||
Count: 1,
|
||
Quality: 0,
|
||
},
|
||
Usage: D3D11_USAGE_STAGING,
|
||
BindFlags: 0,
|
||
CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32,
|
||
MiscFlags: 0,
|
||
};
|
||
let mut staging_opt: Option<ID3D11Texture2D> = None;
|
||
device
|
||
.CreateTexture2D(&sdesc, None, Some(&mut staging_opt))
|
||
.map_err(|e| anyhow!("CreateTexture2D failed: {e}"))?;
|
||
let staging = staging_opt.ok_or_else(|| anyhow!("no staging texture"))?;
|
||
let staging_res: ID3D11Resource = staging.cast()?;
|
||
|
||
Ok(Self {
|
||
_device: device,
|
||
context,
|
||
staging_res,
|
||
dup,
|
||
width,
|
||
height,
|
||
})
|
||
}
|
||
}
|
||
|
||
/// true=拿到新帧并已复制到 staging；false=超时无更新
|
||
pub fn acquire(&mut self, timeout_ms: u32) -> Result<bool> {
|
||
unsafe {
|
||
let mut info = DXGI_OUTDUPL_FRAME_INFO::default();
|
||
let mut resource: Option<IDXGIResource> = None;
|
||
match self.dup.AcquireNextFrame(timeout_ms, &mut info, &mut resource) {
|
||
Ok(()) => {
|
||
let res = resource.ok_or_else(|| anyhow!("no desktop resource"))?;
|
||
let src: ID3D11Resource = res.cast()?;
|
||
self.context.CopyResource(&self.staging_res, &src);
|
||
let _ = self.dup.ReleaseFrame();
|
||
Ok(true)
|
||
}
|
||
Err(e) if e.code() == DXGI_ERROR_WAIT_TIMEOUT => Ok(false),
|
||
Err(e) => Err(e.into()),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 把 staging 纹理映射到内存并去掉行间距
|
||
pub fn copy_frame(&mut self) -> Result<Vec<u8>> {
|
||
unsafe {
|
||
let mut mapped = D3D11_MAPPED_SUBRESOURCE::default();
|
||
self.context
|
||
.Map(&self.staging_res, 0, D3D11_MAP_READ, 0, Some(&mut mapped))
|
||
.map_err(|e| anyhow!("Map staging failed: {e}"))?;
|
||
let (w, h) = (self.width as usize, self.height as usize);
|
||
let pitch = mapped.RowPitch as usize;
|
||
let src = std::slice::from_raw_parts(mapped.pData as *const u8, pitch * h);
|
||
let mut out = Vec::with_capacity(w * h * 4);
|
||
for row in 0..h {
|
||
out.extend_from_slice(&src[row * pitch..row * pitch + w * 4]);
|
||
}
|
||
self.context.Unmap(&self.staging_res, 0);
|
||
Ok(out)
|
||
}
|
||
}
|
||
}
|
||
|
||
#[allow(dead_code)]
|
||
pub fn desktop_size() -> Result<(u32, u32)> {
|
||
let d = Duplication::new()?;
|
||
Ok((d.width, d.height))
|
||
}
|
||
|
||
/// 采集主循环：独立线程运行；连接断开（通道关闭）时退出。
|
||
/// 通道为 watch 单槽（latest-wins）：发送即覆盖，网络慢时不产生积压。
|
||
pub fn run(tx: Sender<Option<Vec<u8>>>, keyreq: Arc<AtomicBool>) {
|
||
let start = Instant::now();
|
||
loop {
|
||
let mut dup = match Duplication::new() {
|
||
Ok(d) => d,
|
||
Err(e) => {
|
||
eprintln!("[capture] init failed: {e}, retry in 2s");
|
||
std::thread::sleep(Duration::from_secs(2));
|
||
continue;
|
||
}
|
||
};
|
||
eprintln!(
|
||
"[capture] duplication started {}x{}",
|
||
dup.width, dup.height
|
||
);
|
||
|
||
// H.264 编码器：屏幕内容实时档 + 码率控制（对标 UU：3.2Mbps@60fps）
|
||
let mut encoder = match Encoder::with_api_config(
|
||
OpenH264API::from_source(),
|
||
EncoderConfig::new()
|
||
.bitrate(BitRate::from_bps(5_000_000))
|
||
.max_frame_rate(FrameRate::from_hz(30.0))
|
||
.rate_control_mode(RateControlMode::Bitrate)
|
||
.usage_type(UsageType::ScreenContentRealTime)
|
||
.intra_frame_period(IntraFramePeriod::from_num_frames(300)),
|
||
) {
|
||
Ok(e) => e,
|
||
Err(e) => {
|
||
eprintln!("[capture] encoder init failed: {e}");
|
||
return;
|
||
}
|
||
};
|
||
|
||
let mut last_frame: Option<Vec<u8>> = None;
|
||
let mut last_sent = Instant::now() - Duration::from_secs(10);
|
||
let min_frame = Duration::from_millis(33); // 30fps 上限
|
||
loop {
|
||
let iter_start = Instant::now();
|
||
match dup.acquire(80) {
|
||
Ok(true) => match dup.copy_frame() {
|
||
Ok(raw) => {
|
||
// 主控端请求关键帧（解码异常恢复）
|
||
if keyreq.swap(false, Ordering::Relaxed) {
|
||
encoder.force_intra_frame();
|
||
}
|
||
let header = VideoFrameHeader {
|
||
width: dup.width,
|
||
height: dup.height,
|
||
ts_ms: start.elapsed().as_millis() as u64,
|
||
};
|
||
let yuv = YUVBuffer::from_bgra8_source(BgraSliceU8::new(
|
||
&raw,
|
||
(dup.width as usize, dup.height as usize),
|
||
));
|
||
match encoder.encode(&yuv) {
|
||
Ok(ebs) => {
|
||
let payload = ebs.to_vec();
|
||
let frame = encode_video_frame(&header, &payload);
|
||
if tx.send(Some(frame.clone())).is_err() {
|
||
return; // 会话结束
|
||
}
|
||
last_frame = Some(frame);
|
||
last_sent = Instant::now();
|
||
}
|
||
Err(e) => {
|
||
eprintln!("[capture] encode failed: {e}");
|
||
}
|
||
}
|
||
// 帧间隔节流到 ~30fps
|
||
let el = iter_start.elapsed();
|
||
if el < min_frame {
|
||
std::thread::sleep(min_frame - el);
|
||
}
|
||
}
|
||
Err(e) => {
|
||
eprintln!("[capture] frame copy failed: {e}");
|
||
break; // 重建 duplication
|
||
}
|
||
},
|
||
Ok(false) => {
|
||
// 桌面静止：低频重发最后一帧（重复 P 帧，解码端画面不变）
|
||
if last_sent.elapsed() > Duration::from_millis(500) {
|
||
if let Some(f) = last_frame.as_ref() {
|
||
if tx.send(Some(f.clone())).is_err() {
|
||
return;
|
||
}
|
||
last_sent = Instant::now();
|
||
}
|
||
} else {
|
||
std::thread::sleep(Duration::from_millis(8));
|
||
}
|
||
}
|
||
Err(e) => {
|
||
eprintln!("[capture] acquire failed: {e}");
|
||
break; // ACCESS_LOST 等，重建
|
||
}
|
||
}
|
||
}
|
||
std::thread::sleep(Duration::from_millis(300));
|
||
}
|
||
}
|