nvfbc/
cuda.rs

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	/// Capturing does not wait for a new frame nor a mouse move.
31	///
32	/// It is therefore possible to capture the same frame multiple times.
33	/// When this occurs, the current_frame parameter of the
34	/// CudaFrameInfo struct is not incremented.
35	NoWait = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT as isize,
36
37	/// Similar to NoWait, except that the capture will not wait if there
38	/// is already a frame available that the client has never seen yet.
39	NoWaitIfNewFrame = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOWAIT_IF_NEW_FRAME_READY as isize,
40
41	/// Capturing waits for a new frame or mouse move.
42	Blocking = NVFBC_TOCUDA_FLAGS_NVFBC_TOCUDA_GRAB_FLAGS_NOFLAGS as isize,
43}
44
45/// Contains information about a frame captured in a CUDA device.
46#[derive(Copy, Clone)]
47pub struct CudaFrameInfo {
48	/// Address of the CUDA buffer where the frame is grabbed.
49	///
50	/// Note that this an address in CUDA memory, not in system memory.
51	pub device_buffer: usize,
52	/// Size of the frame in bytes.
53	pub device_buffer_len: u32,
54	/// Width of the captured frame.
55	pub width: u32,
56	/// Height of the captured frame.
57	pub height: u32,
58	/// Incremental ID of the current frame.
59	///
60	/// This can be used to identify a frame.
61	pub current_frame: u32,
62	/// Whether this frame is a new frame.
63	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
78/// Uses NVFBC to capture frames in the form of a CUDA device pointer.
79pub struct CudaCapturer {
80	/// A handle to the internal NVFBC instance used for FFI interaction.
81	handle: Handle,
82}
83
84impl CudaCapturer {
85	/// Create a new CUDA capture object.
86	///
87	/// CUDA must be initialized before creating this object.
88	pub fn new() -> Result<Self, Error> {
89		Ok(Self { handle: create_handle()? })
90	}
91
92	/// Retrieve the status of NVFBC.
93	pub fn status(&self) -> Result<Status, Error> {
94		status(self.handle)
95	}
96
97	/// Start a capture session with the desired buffer format.
98	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	/// Stop a capture session.
112	pub fn stop(&self) -> Result<(), Error> {
113		destroy_capture_session(self.handle)
114	}
115
116	/// Retrieve the next frame from the GPU.
117	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	/// Releases the FBC context from the calling thread.
141	///
142	/// If the FBC context is bound to a different thread, nvfbc_sys::_NVFBCSTATUS_NVFBC_ERR_CONTEXT is
143	/// returned.
144	///
145	/// If the FBC context is already released, this function has no effect.
146	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	/// Binds the FBC context to the calling thread.
156	///
157	/// The NvFBC library internally relies on objects that must be bound to a
158	/// thread. Such objects are OpenGL contexts and CUDA contexts.
159	///
160	/// This function binds these objects to the calling thread.
161	///
162	/// The FBC context must be bound to the calling thread for most NvFBC entry
163	/// points, otherwise nvfbc_sys::_NVFBCSTATUS_NVFBC_ERR_CONTEXT is returned.
164	///
165	/// If the FBC context is already bound to a different thread,
166	/// nvfbc_sys::_NVFBCSTATUS_NVFBC_ERR_CONTEXT is returned.
167	/// The other thread must release the context first by calling the release_context().
168	///
169	/// If the FBC context is already bound to the current thread, this function has
170	/// no effects.
171	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		// TODO: Figure out why this crashes (nvfbc examples also fail here..)
184		destroy_handle(self.handle).ok();
185	}
186}