ultraslayer 0.2.4

//! src/shm.rs
//!
//! Minimal shared‑memory wrapper that can be used by multiple processes
//! to access the same UltraSlayer slab.  The implementation is deliberately
//! simple – it creates (or opens) a file in `/dev/shm`, mmaps it with
//! `MAP_SHARED | MAP_HUGETLB` when huge pages are available, and then hands
//! the raw pointer to `UltraSlayer` for its internal operations.
//!
//! The public API mirrors the tiny subset that the README refers to:
//!   * `ShmSlab::<T>::create(name, channels, size_bytes)` – allocate a new slab.
//!   * `ShmSlab::<T>::open(name)` – open an existing slab created elsewhere.
//!   * `read(idx)`, `write(idx, val)` – volatile operations.
//!   * `into_ultraslayer(self) -> UltraSlayer<T>` – consume the wrapper and get the
//!     fully‑featured slab object.
//!
//! The code works on Linux only (uses `memfd_create`/`shm_open` and `mmap`).

#[cfg(unix)]
use std::ffi::CString;
#[cfg(unix)]
use std::fs::OpenOptions;
#[cfg(unix)]
use std::io::{Error, ErrorKind, Result};
#[cfg(unix)]
use std::os::unix::io::AsRawFd;
#[cfg(unix)]
use std::path::Path;
#[cfg(unix)]
use std::ptr::NonNull;
#[cfg(unix)]
use std::sync::Arc;

// Import from reader, as UltraSlayer and SpinPolicy live there now
#[cfg(unix)]
use crate::reader::{UltraSlayer, SpinPolicy};

/// Helper to create a *named* shared‑memory object.  On most Linux systems
/// `/dev/shm/<name>` is a `tmpfs` mount that is perfect for this purpose.
#[cfg(unix)]
fn shm_file_path(name: &str) -> Result<std::fs::File> {
    let path = Path::new("/dev/shm").join(name);
    let file = OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(&path)?;
    Ok(file)
}

/// The public wrapper.  It holds the raw mmap pointer and the length, and
/// forwards all operations to an inner `UltraSlayer<T>`.
#[cfg(unix)]
pub struct ShmSlab<T> {
    ptr: NonNull<T>,
    len: usize,
    _file: std::fs::File,
    slayer: UltraSlayer<T>,
}

#[cfg(unix)]
impl<T> ShmSlab<T>
where
    T: Copy + Send + Sync + 'static,
{
    /// Create a brand‑new shared slab.
    pub fn create(name: &str, channels: usize, size_bytes: usize) -> Result<Self> {
        let file = shm_file_path(name)?;
        file.set_len(size_bytes as u64)?;

        let prot = libc::PROT_READ | libc::PROT_WRITE;
        let flags = libc::MAP_SHARED
            | libc::MAP_LOCKED
            | libc::MAP_POPULATE
            | libc::MAP_HUGETLB;

        let raw_ptr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                size_bytes,
                prot,
                flags,
                file.as_raw_fd(),
                0,
            )
        };

        if raw_ptr == libc::MAP_FAILED {
            return Err(Error::last_os_error());
        }

        let ptr = NonNull::new(raw_ptr as *mut T).ok_or_else(|| {
            Error::new(ErrorKind::Other, "mmap returned a null pointer")
        })?;

        // Note: UltraSlayer::new expects (num_replicas, capacity_in_elements)
        let slayer = UltraSlayer::<T>::new(channels, size_bytes / std::mem::size_of::<T>());

        Ok(Self {
            ptr,
            len: size_bytes / std::mem::size_of::<T>(),
            _file: file,
            slayer,
        })
    }

    /// Open an already‑existing shared slab.
    pub fn open(name: &str, channels: usize, size_bytes: usize) -> Result<Self> {
        let file = shm_file_path(name)?;
        let raw_ptr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                size_bytes,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_SHARED,
                file.as_raw_fd(),
                0,
            )
        };
        if raw_ptr == libc::MAP_FAILED {
            return Err(Error::last_os_error());
        }
        let ptr = NonNull::new(raw_ptr as *mut T).ok_or_else(|| {
            Error::new(ErrorKind::Other, "mmap returned a null pointer")
        })?;
        let slayer = UltraSlayer::<T>::new(channels, size_bytes / std::mem::size_of::<T>());
        Ok(Self {
            ptr,
            len: size_bytes / std::mem::size_of::<T>(),
            _file: file,
            slayer,
        })
    }

    pub fn into_ultraslayer(self) -> UltraSlayer<T> {
        self.slayer
    }

    #[inline(always)]
    pub fn read(&self, idx: usize) -> T {
        self.slayer.read(idx)
    }

    #[inline(always)]
    pub fn write(&self, idx: usize, val: T) {
        self.slayer.insert(idx, val);
    }

    #[inline(always)]
    pub fn set_spin_policy(&self, policy: SpinPolicy) {
        self.slayer.set_spin_policy(policy);
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.len
    }
}

#[cfg(unix)]
impl<T> Drop for ShmSlab<T> {
    fn drop(&mut self) {
        let size_bytes = self.len * std::mem::size_of::<T>();
        unsafe {
            libc::munmap(self.ptr.as_ptr() as *mut _, size_bytes);
        }
    }
}