1use std::mem;
2use std::os::windows::prelude::AsRawHandle;
3use std::sync::atomic::{self, AtomicBool};
4use std::sync::{Arc, mpsc};
5use std::thread::{self, JoinHandle};
6
7use parking_lot::Mutex;
8use windows::Win32::Foundation::{HANDLE, LPARAM, WPARAM};
9use windows::Win32::Graphics::Direct3D11::{ID3D11Device, ID3D11DeviceContext};
10use windows::Win32::System::Threading::{GetCurrentThreadId, GetThreadId};
11use windows::Win32::System::WinRT::{
12 CreateDispatcherQueueController, DQTAT_COM_NONE, DQTYPE_THREAD_CURRENT, DispatcherQueueOptions,
13};
14use windows::Win32::UI::WindowsAndMessaging::{
15 DispatchMessageW, GetMessageW, MSG, PostQuitMessage, PostThreadMessageW, TranslateMessage, WM_QUIT,
16};
17use windows::core::Result as WindowsResult;
18use windows_future::AsyncActionCompletedHandler;
19
20use crate::d3d11::{self, create_d3d_device};
21use crate::frame::Frame;
22use crate::graphics_capture_api::{self, GraphicsCaptureApi, InternalCaptureControl};
23use crate::settings::{GraphicsCaptureItemType, Settings};
24use crate::winrt::WinRT;
25
26#[derive(thiserror::Error, Debug)]
27pub enum CaptureControlError<E> {
32 #[error("Failed to join thread")]
37 FailedToJoinThread,
38 #[error("Thread handle is taken out of the struct")]
41 ThreadHandleIsTaken,
42 #[error("Failed to post thread message")]
46 FailedToPostThreadMessage,
47 #[error("Stopped handler error: {0}")]
51 StoppedHandlerError(E),
52 #[error("Windows capture error: {0}")]
56 GraphicsCaptureApiError(#[from] GraphicsCaptureApiError<E>),
57}
58
59pub struct CaptureControl<T: GraphicsCaptureApiHandler + Send + 'static, E> {
61 thread_handle: Option<JoinHandle<Result<(), GraphicsCaptureApiError<E>>>>,
62 halt_handle: Arc<AtomicBool>,
63 callback: Arc<Mutex<T>>,
64}
65
66impl<T: GraphicsCaptureApiHandler + Send + 'static, E> CaptureControl<T, E> {
67 #[inline]
69 #[must_use]
70 pub const fn new(
71 thread_handle: JoinHandle<Result<(), GraphicsCaptureApiError<E>>>,
72 halt_handle: Arc<AtomicBool>,
73 callback: Arc<Mutex<T>>,
74 ) -> Self {
75 Self { thread_handle: Some(thread_handle), halt_handle, callback }
76 }
77
78 #[inline]
80 #[must_use]
81 pub fn is_finished(&self) -> bool {
82 self.thread_handle.as_ref().is_none_or(std::thread::JoinHandle::is_finished)
83 }
84
85 #[inline]
87 #[must_use]
88 pub fn into_thread_handle(self) -> JoinHandle<Result<(), GraphicsCaptureApiError<E>>> {
89 self.thread_handle.unwrap()
90 }
91
92 #[inline]
94 #[must_use]
95 pub fn halt_handle(&self) -> Arc<AtomicBool> {
96 self.halt_handle.clone()
97 }
98
99 #[inline]
101 #[must_use]
102 pub fn callback(&self) -> Arc<Mutex<T>> {
103 self.callback.clone()
104 }
105
106 #[inline]
114 pub fn wait(mut self) -> Result<(), CaptureControlError<E>> {
115 if let Some(thread_handle) = self.thread_handle.take() {
116 match thread_handle.join() {
117 Ok(result) => result?,
118 Err(_) => {
119 return Err(CaptureControlError::FailedToJoinThread);
120 }
121 }
122 } else {
123 return Err(CaptureControlError::ThreadHandleIsTaken);
124 }
125
126 Ok(())
127 }
128
129 #[inline]
141 pub fn stop(mut self) -> Result<(), CaptureControlError<E>> {
142 self.halt_handle.store(true, atomic::Ordering::Relaxed);
143
144 if let Some(thread_handle) = self.thread_handle.take() {
145 let handle = thread_handle.as_raw_handle();
146 let handle = HANDLE(handle);
147 let thread_id = unsafe { GetThreadId(handle) };
148
149 loop {
150 match unsafe { PostThreadMessageW(thread_id, WM_QUIT, WPARAM::default(), LPARAM::default()) } {
151 Ok(()) => break,
152 Err(e) => {
153 if thread_handle.is_finished() {
154 break;
155 }
156
157 if e.code().0 != -2_147_023_452 {
158 Err(e).map_err(|_| CaptureControlError::FailedToPostThreadMessage)?;
159 }
160 }
161 }
162 }
163
164 match thread_handle.join() {
165 Ok(result) => result?,
166 Err(_) => {
167 return Err(CaptureControlError::FailedToJoinThread);
168 }
169 }
170 } else {
171 return Err(CaptureControlError::ThreadHandleIsTaken);
172 }
173
174 Ok(())
175 }
176}
177
178#[derive(thiserror::Error, Eq, PartialEq, Clone, Debug)]
179pub enum GraphicsCaptureApiError<E> {
181 #[error("Failed to join thread")]
183 FailedToJoinThread,
184 #[error("Failed to initialize WinRT")]
188 FailedToInitWinRT,
189 #[error("Failed to create dispatcher queue controller")]
191 FailedToCreateDispatcherQueueController,
192 #[error("Failed to shut down dispatcher queue")]
194 FailedToShutdownDispatcherQueue,
195 #[error("Failed to set dispatcher queue completed handler")]
197 FailedToSetDispatcherQueueCompletedHandler,
198 #[error("Failed to convert item to `GraphicsCaptureItem`")]
204 ItemConvertFailed,
205 #[error("DirectX error: {0}")]
209 DirectXError(#[from] d3d11::Error),
210 #[error("Graphics capture error: {0}")]
214 GraphicsCaptureApiError(graphics_capture_api::Error),
215 #[error("New handler error: {0}")]
218 NewHandlerError(E),
219 #[error("Frame handler error: {0}")]
223 FrameHandlerError(E),
224}
225
226pub struct Context<Flags> {
228 pub flags: Flags,
230 pub device: ID3D11Device,
232 pub device_context: ID3D11DeviceContext,
234}
235
236pub trait GraphicsCaptureApiHandler: Sized {
238 type Flags;
240
241 type Error: Send + Sync;
244
245 #[inline]
247 fn start<T: TryInto<GraphicsCaptureItemType>>(
248 settings: Settings<Self::Flags, T>,
249 ) -> Result<(), GraphicsCaptureApiError<Self::Error>>
250 where
251 Self: Send + 'static,
252 <Self as GraphicsCaptureApiHandler>::Flags: Send,
253 {
254 let _winrt = WinRT::new();
256
257 let options = DispatcherQueueOptions {
259 dwSize: u32::try_from(mem::size_of::<DispatcherQueueOptions>()).unwrap(),
260 threadType: DQTYPE_THREAD_CURRENT,
261 apartmentType: DQTAT_COM_NONE,
262 };
263 let controller = unsafe {
264 CreateDispatcherQueueController(options)
265 .map_err(|_| GraphicsCaptureApiError::FailedToCreateDispatcherQueueController)?
266 };
267
268 let thread_id = unsafe { GetCurrentThreadId() };
270
271 let (d3d_device, d3d_device_context) = create_d3d_device()?;
273
274 let result = Arc::new(Mutex::new(None));
276
277 let ctx =
278 Context { flags: settings.flags, device: d3d_device.clone(), device_context: d3d_device_context.clone() };
279
280 let callback = Arc::new(Mutex::new(Self::new(ctx).map_err(GraphicsCaptureApiError::NewHandlerError)?));
281
282 let mut capture = GraphicsCaptureApi::new(
283 d3d_device,
284 d3d_device_context,
285 settings.item.try_into().map_err(|_| GraphicsCaptureApiError::ItemConvertFailed)?,
286 callback,
287 settings.cursor_capture_settings,
288 settings.draw_border_settings,
289 settings.secondary_window_settings,
290 settings.minimum_update_interval_settings,
291 settings.dirty_region_settings,
292 settings.color_format,
293 thread_id,
294 result.clone(),
295 )
296 .map_err(GraphicsCaptureApiError::GraphicsCaptureApiError)?;
297 capture.start_capture().map_err(GraphicsCaptureApiError::GraphicsCaptureApiError)?;
298
299 let mut message = MSG::default();
301 unsafe {
302 while GetMessageW(&mut message, None, 0, 0).as_bool() {
303 let _ = TranslateMessage(&message);
304 DispatchMessageW(&message);
305 }
306 }
307
308 let async_action =
310 controller.ShutdownQueueAsync().map_err(|_| GraphicsCaptureApiError::FailedToShutdownDispatcherQueue)?;
311
312 async_action
313 .SetCompleted(&AsyncActionCompletedHandler::new(move |_, _| -> WindowsResult<()> {
314 unsafe { PostQuitMessage(0) };
315 Ok(())
316 }))
317 .map_err(|_| GraphicsCaptureApiError::FailedToSetDispatcherQueueCompletedHandler)?;
318
319 let mut message = MSG::default();
321 unsafe {
322 while GetMessageW(&mut message, None, 0, 0).as_bool() {
323 let _ = TranslateMessage(&message);
324 DispatchMessageW(&message);
325 }
326 }
327
328 capture.stop_capture();
330
331 let result = result.lock().take();
333 if let Some(e) = result {
334 return Err(GraphicsCaptureApiError::FrameHandlerError(e));
335 }
336
337 Ok(())
338 }
339
340 #[inline]
342 fn start_free_threaded<T: TryInto<GraphicsCaptureItemType> + Send + 'static>(
343 settings: Settings<Self::Flags, T>,
344 ) -> Result<CaptureControl<Self, Self::Error>, GraphicsCaptureApiError<Self::Error>>
345 where
346 Self: Send + 'static,
347 <Self as GraphicsCaptureApiHandler>::Flags: Send,
348 {
349 let (halt_sender, halt_receiver) = mpsc::channel::<Arc<AtomicBool>>();
350 let (callback_sender, callback_receiver) = mpsc::channel::<Arc<Mutex<Self>>>();
351
352 let thread_handle = thread::spawn(move || -> Result<(), GraphicsCaptureApiError<Self::Error>> {
353 let _winrt = WinRT::new();
355
356 let options = DispatcherQueueOptions {
358 dwSize: u32::try_from(mem::size_of::<DispatcherQueueOptions>()).unwrap(),
359 threadType: DQTYPE_THREAD_CURRENT,
360 apartmentType: DQTAT_COM_NONE,
361 };
362 let controller = unsafe {
363 CreateDispatcherQueueController(options)
364 .map_err(|_| GraphicsCaptureApiError::FailedToCreateDispatcherQueueController)?
365 };
366
367 let thread_id = unsafe { GetCurrentThreadId() };
369
370 let (d3d_device, d3d_device_context) = create_d3d_device()?;
372
373 let result = Arc::new(Mutex::new(None));
375
376 let ctx = Context {
377 flags: settings.flags,
378 device: d3d_device.clone(),
379 device_context: d3d_device_context.clone(),
380 };
381
382 let callback = Arc::new(Mutex::new(Self::new(ctx).map_err(GraphicsCaptureApiError::NewHandlerError)?));
383
384 let mut capture = GraphicsCaptureApi::new(
385 d3d_device,
386 d3d_device_context,
387 settings.item.try_into().map_err(|_| GraphicsCaptureApiError::ItemConvertFailed)?,
388 callback.clone(),
389 settings.cursor_capture_settings,
390 settings.draw_border_settings,
391 settings.secondary_window_settings,
392 settings.minimum_update_interval_settings,
393 settings.dirty_region_settings,
394 settings.color_format,
395 thread_id,
396 result.clone(),
397 )
398 .map_err(GraphicsCaptureApiError::GraphicsCaptureApiError)?;
399
400 capture.start_capture().map_err(GraphicsCaptureApiError::GraphicsCaptureApiError)?;
401
402 let halt_handle = capture.halt_handle();
404 halt_sender.send(halt_handle).unwrap();
405
406 callback_sender.send(callback).unwrap();
408
409 let mut message = MSG::default();
411 unsafe {
412 while GetMessageW(&mut message, None, 0, 0).as_bool() {
413 let _ = TranslateMessage(&message);
414 DispatchMessageW(&message);
415 }
416 }
417
418 let async_action = controller
420 .ShutdownQueueAsync()
421 .map_err(|_| GraphicsCaptureApiError::FailedToShutdownDispatcherQueue)?;
422
423 async_action
424 .SetCompleted(&AsyncActionCompletedHandler::new(move |_, _| -> Result<(), windows::core::Error> {
425 unsafe { PostQuitMessage(0) };
426 Ok(())
427 }))
428 .map_err(|_| GraphicsCaptureApiError::FailedToSetDispatcherQueueCompletedHandler)?;
429
430 let mut message = MSG::default();
432 unsafe {
433 while GetMessageW(&mut message, None, 0, 0).as_bool() {
434 let _ = TranslateMessage(&message);
435 DispatchMessageW(&message);
436 }
437 }
438
439 capture.stop_capture();
441
442 let result = result.lock().take();
444 if let Some(e) = result {
445 return Err(GraphicsCaptureApiError::FrameHandlerError(e));
446 }
447
448 Ok(())
449 });
450
451 let Ok(halt_handle) = halt_receiver.recv() else {
452 match thread_handle.join() {
453 Ok(result) => return Err(result.err().unwrap()),
454 Err(_) => {
455 return Err(GraphicsCaptureApiError::FailedToJoinThread);
456 }
457 }
458 };
459
460 let Ok(callback) = callback_receiver.recv() else {
461 match thread_handle.join() {
462 Ok(result) => return Err(result.err().unwrap()),
463 Err(_) => {
464 return Err(GraphicsCaptureApiError::FailedToJoinThread);
465 }
466 }
467 };
468
469 Ok(CaptureControl::new(thread_handle, halt_handle, callback))
470 }
471
472 fn new(ctx: Context<Self::Flags>) -> Result<Self, Self::Error>;
475
476 fn on_frame_arrived(
478 &mut self,
479 frame: &mut Frame,
480 capture_control: InternalCaptureControl,
481 ) -> Result<(), Self::Error>;
482
483 #[inline]
485 fn on_closed(&mut self) -> Result<(), Self::Error> {
486 Ok(())
487 }
488}