use libc::{c_void, io_uring_register, IORING_REGISTER_BUFFERS};
use std::os::unix::io::RawFd;
use std::alloc::{alloc, dealloc, Layout};
#[repr(C)] struct Iovec {
iov_base: *mut c_void,
iov_len: usize,
}
pub struct FixedMmuBufferPool {
ring_fd: RawFd,
fixed_buffers: Vec<Iovec>,
}
#[derive(Debug)]
pub enum IoUringError {
RegistrationFailed { os_error: i32, context: &'static str },
AllocationFailed { count: usize, alignment: usize },
}
impl std::fmt::Display for IoUringError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::RegistrationFailed { os_error, context } => write!(f, "io_uring_register failed: {} (OS error: {})", context, os_error),
Self::AllocationFailed { count, alignment } => write!(f, "failed to allocate io_uring pinned buffer pool: {} elements aligned to {}", count, alignment),
}
}
}
impl std::error::Error for IoUringError {}
impl FixedMmuBufferPool {
pub fn register_physical_buffers(ring_fd: RawFd, capacity: usize, count: usize) -> Result<Self, IoUringError> {
let mut buffers = Vec::with_capacity(count);
let align = 4096; let layout = Layout::from_size_align(capacity, align).map_err(|_| {
IoUringError::AllocationFailed { count: 1, alignment: align }
})?;
for _ in 0..count {
let ptr = unsafe { alloc(layout) };
if ptr.is_null() {
for b in &buffers {
unsafe { dealloc(b.iov_base as *mut u8, layout) };
}
return Err(IoUringError::AllocationFailed { count, alignment: align });
}
buffers.push(Iovec {
iov_base: ptr as *mut c_void,
iov_len: capacity,
});
}
let res = unsafe {
io_uring_register(
ring_fd,
IORING_REGISTER_BUFFERS,
buffers.as_ptr() as *const c_void,
count as u32
)
};
if res < 0 {
let err = std::io::Error::last_os_error().raw_os_error().unwrap_or(0);
for b in &buffers {
unsafe { dealloc(b.iov_base as *mut u8, layout) };
}
return Err(IoUringError::RegistrationFailed { os_error: err, context: "kernel rejected pinned buffer hardware restrictions natively" });
}
tracing::info!("Cudagrep flawlessly registered {} native Hardware Buffer mappings inside the Linux Kernel.", count);
Ok(Self { ring_fd, fixed_buffers: buffers })
}
}
impl Drop for FixedMmuBufferPool {
fn drop(&mut self) {
unsafe {
let res = libc::io_uring_register(self.ring_fd, 1, std::ptr::null(), 0);
if res < 0 {
tracing::error!("failed to unregister fixed IORING mappings cleanly natively");
}
let align = 4096;
for b in &self.fixed_buffers {
let layout = Layout::from_size_align_unchecked(b.iov_len, align);
dealloc(b.iov_base as *mut u8, layout);
}
}
}
}