1use std::ffi::c_void;
2use std::mem::MaybeUninit;
3use std::ptr::null_mut;
4use std::time::Duration;
5
6use nvfbc_sys::{
7 NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT,
8 NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOFLAGS,
9 NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT_IF_NEW_FRAME_READY
10};
11
12use crate::{
13 BufferFormat,
14 CaptureType,
15 Error,
16 Status,
17};
18
19use crate::common::{
20 Handle,
21 check_ret,
22 create_capture_session,
23 create_handle,
24 destroy_capture_session,
25 destroy_handle,
26 status,
27};
28
29pub enum CaptureMethod {
30 NoWait = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT as isize,
36
37 NoWaitIfNewFrame = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT_IF_NEW_FRAME_READY as isize,
40
41 Blocking = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOFLAGS as isize,
43}
44
45#[derive(Copy, Clone)]
47pub struct CudaFrameInfo {
48 pub device_buffer: usize,
52 pub device_buffer_len: u32,
54 pub width: u32,
56 pub height: u32,
58 pub current_frame: u32,
62 pub is_new_frame: bool,
64}
65
66impl std::fmt::Debug for CudaFrameInfo {
67 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
68 f.debug_struct("CudaFrameInfo")
69 .field("device_buffer", &(&self.device_buffer as *const usize))
70 .field("device_buffer_len", &self.device_buffer_len)
71 .field("width", &self.width)
72 .field("height", &self.height)
73 .field("current_frame", &self.current_frame)
74 .finish()
75 }
76}
77
78pub struct CudaCapturer {
80 handle: Handle,
82}
83
84impl CudaCapturer {
85 pub fn new() -> Result<Self, Error> {
89 Ok(Self { handle: create_handle()? })
90 }
91
92 pub fn status(&self) -> Result<Status, Error> {
94 status(self.handle)
95 }
96
97 pub fn start(&self, buffer_format: BufferFormat, fps: u32) -> Result<(), Error> {
99 create_capture_session(
100 self.handle,
101 CaptureType::SharedCuda,
102 std::time::Duration::from_millis(1000 / fps as u64),
103 )?;
104
105 let mut params: nvfbc_sys::NVFBC_TOCUDA_SETUP_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
106 params.dwVersion = nvfbc_sys::NVFBC_TOCUDA_SETUP_PARAMS_VER;
107 params.eBufferFormat = buffer_format as u32;
108 check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToCudaSetUp(self.handle, &mut params) })
109 }
110
111 pub fn stop(&self) -> Result<(), Error> {
113 destroy_capture_session(self.handle)
114 }
115
116 pub fn next_frame(&mut self, capture_method: CaptureMethod, timeout: Option<Duration>) -> Result<CudaFrameInfo, Error> {
118 let mut device_buffer: *mut c_void = null_mut();
119 let mut frame_info: nvfbc_sys::NVFBC_FRAME_GRAB_INFO = unsafe { MaybeUninit::zeroed().assume_init() };
120 let mut params: nvfbc_sys::NVFBC_TOCUDA_GRAB_FRAME_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
121 params.dwVersion = nvfbc_sys::NVFBC_TOCUDA_GRAB_FRAME_PARAMS_VER;
122 params.dwFlags = capture_method as u32;
123 params.pFrameGrabInfo = &mut frame_info;
124 params.pCUDADeviceBuffer = &mut device_buffer as *mut _ as *mut c_void;
125 if let Some(timeout) = timeout {
126 params.dwTimeoutMs = timeout.as_millis() as u32;
127 }
128 check_ret(self.handle, unsafe { nvfbc_sys::NvFBCToCudaGrabFrame(self.handle, &mut params) })?;
129
130 Ok(CudaFrameInfo {
131 device_buffer: device_buffer as usize,
132 device_buffer_len: frame_info.dwByteSize,
133 width: frame_info.dwWidth,
134 height: frame_info.dwHeight,
135 current_frame: frame_info.dwCurrentFrame,
136 is_new_frame: frame_info.bIsNewFrame != 0,
137 })
138 }
139
140 pub fn release_context(&self) -> Result<(), Error> {
147 let mut params: nvfbc_sys::NVFBC_RELEASE_CONTEXT_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
148 params.dwVersion = nvfbc_sys::NVFBC_RELEASE_CONTEXT_PARAMS_VER;
149 check_ret(
150 self.handle,
151 unsafe { nvfbc_sys::NvFBCReleaseContext(self.handle, &mut params) }
152 )
153 }
154
155 pub fn bind_context(&self) -> Result<(), Error> {
172 let mut params: nvfbc_sys::NVFBC_BIND_CONTEXT_PARAMS = unsafe { MaybeUninit::zeroed().assume_init() };
173 params.dwVersion = nvfbc_sys::NVFBC_BIND_CONTEXT_PARAMS_VER;
174 check_ret(
175 self.handle,
176 unsafe { nvfbc_sys::NvFBCBindContext(self.handle, &mut params) }
177 )
178 }
179}
180
181impl Drop for CudaCapturer {
182 fn drop(&mut self) {
183 destroy_handle(self.handle).ok();
185 }
186}