1
0
Fork 0
forked from wry/wry

clientmem: move shm mapping into workspace crate

This commit is contained in:
kossLAN 2026-05-29 11:41:55 -04:00
parent bf3859a026
commit 663cfb3ca3
No known key found for this signature in database
10 changed files with 383 additions and 338 deletions

331
clientmem/src/lib.rs Normal file
View file

@ -0,0 +1,331 @@
use {
jay_cpu_worker::{AsyncCpuWork, CpuJob, CpuWork, CpuWorker},
jay_gfx_types::{ShmMemory, ShmMemoryBacking},
jay_utils::{
oserror::{OsError, OsErrorExt2},
page_size::page_size,
vec_ext::VecExt,
},
std::{
cell::Cell,
error::Error,
mem::{ManuallyDrop, MaybeUninit},
ops::Deref,
ptr,
rc::Rc,
sync::atomic::{Ordering, compiler_fence},
},
thiserror::Error,
uapi::{
OwnedFd, Pod,
c::{self, raise},
ftruncate,
},
};
#[derive(Copy, Clone, Debug)]
pub struct ClientMemClient<'a> {
pub comm: &'a str,
pub id: u64,
}
#[derive(Debug, Error)]
pub enum ClientMemError {
#[error("Could not install the sigbus handler")]
SigactionFailed(#[source] jay_utils::oserror::OsError),
#[error("A SIGBUS occurred while accessing mapped memory")]
Sigbus,
#[error("mmap failed")]
MmapFailed(#[source] jay_utils::oserror::OsError),
#[error("Length was not a multiple of the data element size")]
InvalidLength,
}
pub struct ClientMem {
fd: ManuallyDrop<Rc<OwnedFd>>,
failed: Cell<bool>,
sigbus_impossible: bool,
data: *const [Cell<u8>],
cpu: Option<Rc<CpuWorker>>,
}
#[derive(Clone)]
pub struct ClientMemOffset {
mem: Rc<ClientMem>,
offset: usize,
data: *const [Cell<u8>],
}
impl ClientMem {
pub fn new(
fd: &Rc<OwnedFd>,
len: usize,
read_only: bool,
client: Option<ClientMemClient<'_>>,
cpu: Option<&Rc<CpuWorker>>,
is_udmabuf: bool,
) -> Result<Self, ClientMemError> {
Self::new2(fd, len, read_only, client, cpu, c::MAP_SHARED, is_udmabuf)
}
pub fn new_private(
fd: &Rc<OwnedFd>,
len: usize,
read_only: bool,
client: Option<ClientMemClient<'_>>,
cpu: Option<&Rc<CpuWorker>>,
) -> Result<Self, ClientMemError> {
Self::new2(fd, len, read_only, client, cpu, c::MAP_PRIVATE, false)
}
fn new2(
fd: &Rc<OwnedFd>,
len: usize,
read_only: bool,
client: Option<ClientMemClient<'_>>,
cpu: Option<&Rc<CpuWorker>>,
flags: c::c_int,
is_udmabuf: bool,
) -> Result<Self, ClientMemError> {
let mut sigbus_impossible = is_udmabuf;
let mut real_size = None;
if !sigbus_impossible
&& let Ok(seals) = uapi::fcntl_get_seals(fd.raw())
&& seals & c::F_SEAL_SHRINK != 0
&& let Ok(stat) = uapi::fstat(fd.raw())
{
real_size = Some(stat.st_size as usize);
sigbus_impossible = stat.st_size as u64 >= len as u64;
}
if !sigbus_impossible && let Some(client) = client {
log::debug!(
"Client {} ({}) has created a shm buffer that might cause SIGBUS",
client.comm,
client.id,
);
}
let len = len.next_multiple_of(page_size());
if let Some(real_size) = real_size
&& real_size < len
{
let _ = ftruncate(fd.raw(), len as _);
}
let data = if len == 0 {
&mut [][..]
} else {
let prot = match read_only {
true => c::PROT_READ,
false => c::PROT_READ | c::PROT_WRITE,
};
unsafe {
let data = c::mmap64(ptr::null_mut(), len, prot, flags, fd.raw(), 0);
if data == c::MAP_FAILED {
return Err(ClientMemError::MmapFailed(OsError::default()));
}
std::slice::from_raw_parts_mut(data as *mut Cell<u8>, len)
}
};
Ok(Self {
fd: ManuallyDrop::new(fd.clone()),
failed: Cell::new(false),
sigbus_impossible,
data,
cpu: cpu.cloned(),
})
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn offset(self: &Rc<Self>, offset: usize, len: usize) -> ClientMemOffset {
let mem = unsafe { &*self.data };
ClientMemOffset {
mem: self.clone(),
offset,
data: &mem[offset..][..len],
}
}
pub fn fd(&self) -> &Rc<OwnedFd> {
&self.fd
}
pub fn is_sealed_memfd(&self) -> bool {
self.sigbus_impossible
}
}
impl ClientMemOffset {
pub fn pool(&self) -> &ClientMem {
&self.mem
}
pub fn offset(&self) -> usize {
self.offset
}
pub fn ptr(&self) -> *const [Cell<u8>] {
self.data
}
pub fn access<T, F: FnOnce(&[Cell<u8>]) -> T>(&self, f: F) -> Result<T, ClientMemError> {
unsafe {
if self.mem.sigbus_impossible {
return Ok(f(&*self.data));
}
let mref = MemRef {
mem: &*self.mem,
outer: MEM.get(),
};
MEM.set(&mref);
compiler_fence(Ordering::SeqCst);
let res = f(&*self.data);
MEM.set(mref.outer);
compiler_fence(Ordering::SeqCst);
match self.mem.failed.get() {
true => Err(ClientMemError::Sigbus),
_ => Ok(res),
}
}
}
pub fn read<T: Pod>(&self, dst: &mut Vec<T>) -> Result<(), ClientMemError> {
if self.data.len().checked_rem(std::mem::size_of::<T>()) != Some(0) {
return Err(ClientMemError::InvalidLength);
}
self.access(|v| {
let len_elements = v.len() / std::mem::size_of::<T>();
dst.reserve(len_elements);
let (_, unused) = dst.split_at_spare_mut_bytes_ext();
unused[..v.len()].copy_from_slice(uapi::as_maybe_uninit_bytes(v));
unsafe {
dst.set_len(dst.len() + len_elements);
}
})
}
}
impl Drop for ClientMem {
fn drop(&mut self) {
let fd = unsafe { ManuallyDrop::take(&mut self.fd) };
if let Some(cpu) = &self.cpu {
let pending = cpu.submit(Box::new(CloseMemWork {
fd: Rc::try_unwrap(fd).ok(),
data: self.data,
}));
pending.detach();
} else {
unsafe {
c::munmap(self.data as _, self.len());
}
}
}
}
struct MemRef {
mem: *const ClientMem,
outer: *const MemRef,
}
thread_local! {
static MEM: Cell<*const MemRef> = const { Cell::new(ptr::null()) };
}
unsafe fn kill() -> ! {
unsafe {
c::signal(c::SIGBUS, c::SIG_DFL);
raise(c::SIGBUS);
}
unreachable!();
}
unsafe extern "C" fn sigbus(sig: i32, info: &c::siginfo_t, _ucontext: *mut c::c_void) {
unsafe {
assert_eq!(sig, c::SIGBUS);
let mut memr_ptr = MEM.get();
while !memr_ptr.is_null() {
let memr = &*memr_ptr;
let mem = &*memr.mem;
let lo = mem.data as *mut u8 as usize;
let hi = lo + mem.len();
let fault_addr = info.si_addr() as usize;
if fault_addr < lo || fault_addr >= hi {
memr_ptr = memr.outer;
continue;
}
let res = c::mmap64(
lo as _,
hi - lo,
c::PROT_WRITE | c::PROT_READ,
c::MAP_ANONYMOUS | c::MAP_PRIVATE | c::MAP_FIXED,
-1,
0,
);
if res == c::MAP_FAILED {
kill();
}
mem.failed.set(true);
return;
}
kill();
}
}
pub fn init() -> Result<(), ClientMemError> {
unsafe {
let mut action: c::sigaction = MaybeUninit::zeroed().assume_init();
action.sa_sigaction =
sigbus as unsafe extern "C" fn(i32, &c::siginfo_t, *mut c::c_void) as _;
action.sa_flags = c::SA_NODEFER | c::SA_SIGINFO;
let res = c::sigaction(c::SIGBUS, &action, ptr::null_mut());
uapi::map_err!(res)
.map(drop)
.map_os_err(ClientMemError::SigactionFailed)
}
}
struct CloseMemWork {
fd: Option<OwnedFd>,
data: *const [Cell<u8>],
}
unsafe impl Send for CloseMemWork {}
impl CpuJob for CloseMemWork {
fn work(&mut self) -> &mut dyn CpuWork {
self
}
fn completed(self: Box<Self>) {
// nothing
}
}
impl CpuWork for CloseMemWork {
fn run(&mut self) -> Option<Box<dyn AsyncCpuWork>> {
jay_tracy::zone!("CloseMemWork");
self.fd.take();
unsafe {
c::munmap(self.data as _, self.data.len());
}
None
}
}
impl ShmMemory for ClientMemOffset {
fn len(&self) -> usize {
self.data.len()
}
fn safe_access(&self) -> ShmMemoryBacking {
match self.mem.is_sealed_memfd() {
true => ShmMemoryBacking::Ptr(self.data),
false => ShmMemoryBacking::Fd(self.mem.fd.deref().clone(), self.offset),
}
}
fn access(&self, f: &mut dyn FnMut(&[Cell<u8>])) -> Result<(), Box<dyn Error + Sync + Send>> {
self.access(f).map_err(|e| e.into())
}
}