mnemosyne-backend 0.6.0

//! Unix mmap/munmap memory backend.

use core::ffi::{c_int, c_void};
#[cfg(all(target_os = "linux", not(miri)))]
use mnemosyne_core::constants::SEGMENT_SIZE;

const PROT_NONE: c_int = 0;
const PROT_READ: c_int = 1;
const PROT_WRITE: c_int = 2;
const MAP_PRIVATE: c_int = 2;

#[cfg(any(
    target_os = "macos",
    target_os = "freebsd",
    target_os = "openbsd",
    target_os = "netbsd"
))]
const MAP_ANON: c_int = 0x1000;
#[cfg(not(any(
    target_os = "macos",
    target_os = "freebsd",
    target_os = "openbsd",
    target_os = "netbsd"
)))]
const MAP_ANON: c_int = 0x20;

const MAP_FAILED: *mut c_void = -1isize as *mut c_void;

/// Linux `MADV_DONTNEED` advice constant.
///
/// Instructs the kernel to drop the physical backing of the addressed
/// range while keeping the mapping itself valid. Subsequent reads return
/// zeroed pages produced by the standard demand-fault path. Defined as
/// `4` on Linux.
#[cfg(all(target_os = "linux", not(miri)))]
const MADV_DONTNEED: c_int = 4;

/// Linux `MADV_FREE` advice constant (kernel ≥ 4.5).
///
/// Marks the addressed range as reclaimable by the kernel, but the pages
/// stay resident until the kernel is under memory pressure.  Unlike
/// `MADV_DONTNEED` this does **not** trigger an IPI broadcast to flush the
/// TLB on all CPUs — critical for high-throughput scientific workloads that
/// free large working-set segments frequently.  Subsequent reads may return
/// the prior contents or zero.  Defined as `8` on Linux.
#[cfg(all(target_os = "linux", not(miri)))]
const MADV_FREE: c_int = 8;

/// Cache of whether the running kernel supports `MADV_FREE` (kernel ≥ 4.5).
///
/// Initialised lazily on the first `decommit` call. `2` = uninitialised,
/// `1` = supported, `0` = not supported (fall back to `MADV_DONTNEED`).
#[cfg(all(target_os = "linux", not(miri)))]
static MADV_FREE_SUPPORTED: core::sync::atomic::AtomicU8 = core::sync::atomic::AtomicU8::new(2);

/// macOS / BSD `MADV_FREE` advice constant.
///
/// Tells the kernel that the addressed range no longer needs to retain
/// its current contents; the kernel may reclaim the physical pages
/// lazily and a subsequent read may return either the prior contents or
/// zeros. Defined as `5` on the BSDs and macOS.
#[cfg(all(any(target_os = "macos", target_os = "freebsd"), not(miri)))]
const MADV_FREE: c_int = 5;

/// Linux `MADV_HUGEPAGE` advice constant.
///
/// Hints the kernel that the mapping is a good candidate for Transparent
/// Huge Pages (THP) promotion. On a 2 MiB-aligned, 2 MiB-multiple mapping
/// (matching `SEGMENT_SIZE` and `SEGMENT_ALIGN`) the kernel can typically
/// back the mapping with a single 2 MiB huge page, halving TLB pressure
/// for hot segment metadata access. Defined as `14` on Linux since 2.6.38.
#[cfg(all(target_os = "linux", not(miri)))]
const MADV_HUGEPAGE: c_int = 14;

unsafe extern "C" {
    fn mmap(
        addr: *mut c_void,
        length: usize,
        prot: c_int,
        flags: c_int,
        fd: c_int,
        offset: isize,
    ) -> *mut c_void;

    fn munmap(addr: *mut c_void, length: usize) -> c_int;

    #[cfg(all(
        any(target_os = "linux", target_os = "macos", target_os = "freebsd"),
        not(miri)
    ))]
    fn madvise(addr: *mut c_void, length: usize, advice: c_int) -> c_int;

    fn mprotect(addr: *mut c_void, length: usize, prot: c_int) -> c_int;
}

/// Issues a Linux `MADV_HUGEPAGE` hint for a freshly mapped segment-sized
/// region. The advice is purely advisory: kernels without THP support or
/// userspace-disabled-THP simply ignore it, so a failure return is dropped
/// silently and does not affect mapping validity.
///
/// On non-Linux Unix targets the hint is a no-op because the same advice
/// constant does not exist or has different semantics.
///
/// # Safety
///
/// `ptr` must be the base of a mapping of at least `length` bytes, and
/// `length` must be the exact mapped length.
#[inline]
unsafe fn hint_hugepage(ptr: *mut u8, length: usize) {
    // Miri only emulates MADV_NORMAL/RANDOM/SEQUENTIAL/WILLNEED, not
    // MADV_HUGEPAGE ("unsupported operation"). This hint is best-effort and
    // already discards its result on real Linux, so skipping the actual
    // syscall under Miri changes no observable behavior — real runs are
    // unaffected since `cfg(miri)` never holds outside `cargo miri`.
    #[cfg(all(target_os = "linux", not(miri)))]
    {
        if length >= SEGMENT_SIZE
            && mnemosyne_core::options::ENABLE_HUGEPAGE_HINT
                .load(core::sync::atomic::Ordering::Relaxed)
        {
            // SAFETY: caller guarantees the mapping covers `length` bytes; madvise
            // is advisory and never invalidates the mapping on failure.
            let _ = unsafe { madvise(ptr as *mut c_void, length, MADV_HUGEPAGE) };
            // The advice discards its result, so this counter is the only
            // observable the hint has, and the only thing a test can assert
            // the decision by.
            crate::recorders::record_hugepage_hint();
        }
    }
    #[cfg(not(all(target_os = "linux", not(miri))))]
    {
        // Reference the arguments so the function signature stays stable
        // across Unix targets without a dead-argument warning.
        let _ = ptr;
        let _ = length;
    }
}

/// Unix virtual memory backend using `mmap`/`munmap`.
pub struct UnixBackend;

impl mnemosyne_core::MemoryBackend for UnixBackend {
    const SUPPORTS_PAGE_RESET: bool = cfg!(any(
        target_os = "linux",
        target_os = "macos",
        target_os = "freebsd"
    ));
    const SUPPORTS_MAKE_GUARD: bool = true;
    const SUPPORTS_DECOMMIT: bool = cfg!(any(
        target_os = "linux",
        target_os = "macos",
        target_os = "freebsd"
    ));

    /// Allocates virtual memory pages of the given size.
    ///
    /// # Safety
    ///
    /// The size must be a multiple of the system page size (usually 4KB).
    unsafe fn allocate(size: usize) -> *mut u8 {
        // SAFETY: Raw system call to mmap to establish a private anonymous page mapping.
        // Size must be page-aligned and non-zero.
        let ptr = unsafe {
            mmap(
                core::ptr::null_mut(),
                size,
                PROT_READ | PROT_WRITE,
                MAP_PRIVATE | MAP_ANON,
                -1,
                0,
            )
        };
        if ptr == MAP_FAILED {
            return core::ptr::null_mut();
        }
        let ptr = ptr as *mut u8;
        // SAFETY: ptr is a valid mapping of `size` bytes. The hint is advisory
        // and may be ignored by the kernel without affecting the mapping.
        unsafe { hint_hugepage(ptr, size) };
        ptr
    }

    /// Releases virtual memory pages previously allocated with `allocate`.
    ///
    /// # Safety
    ///
    /// The `ptr` must be the exact base address returned by `allocate` and
    /// cannot be used after release.
    unsafe fn deallocate(ptr: *mut u8, size: usize) -> bool {
        if ptr.is_null() {
            return false;
        }
        // SAFETY: Raw system call to munmap. The ptr must point to a valid mapped region
        // of the specified size.
        let res = unsafe { munmap(ptr as *mut c_void, size) };
        debug_assert_eq!(res, 0, "munmap failed");
        res == 0
    }

    /// Drops the physical backing of the addressed range while keeping the
    /// virtual mapping valid. Uses `MADV_DONTNEED` on Linux (subsequent
    /// reads return zero) and `MADV_FREE` on macOS/FreeBSD.
    ///
    /// # THP-aware behaviour
    ///
    /// On Linux, calling `MADV_DONTNEED` on a sub-2MB range inside a 2MB-aligned
    /// region forces the kernel to split a Transparent Huge Page (THP), incurring
    /// significant TLB overhead on multi-core systems.  This path applies the
    /// following guard:
    ///
    /// - If `size >= SEGMENT_SIZE` AND `ptr` is 2MB-aligned: the range covers a
    ///   complete THP unit; `MADV_DONTNEED` is safe and zeroes pages.
    /// - Otherwise (sub-segment range, e.g. pool recycling skipping page-0):
    ///   prefer `MADV_FREE` when available to avoid THP splitting.  The pool
    ///   recycle path (`purge_segment_pool`) only needs the physical pages
    ///   returned, not guaranteed zeroing; `MADV_FREE` satisfies this without
    ///   splitting the THP.
    unsafe fn page_reset(ptr: *mut u8, size: usize) -> bool {
        if ptr.is_null() || size == 0 {
            return false;
        }
        // Miri only emulates MADV_NORMAL/RANDOM/SEQUENTIAL/WILLNEED, not
        // MADV_DONTNEED/MADV_FREE ("unsupported operation"). Route Miri
        // through the same `false` fallback already used for genuinely
        // unsupported Unix targets below — callers already treat `false` as
        // "the OS-level optimization is unavailable here", which is exactly
        // Miri's situation for this syscall.
        #[cfg(all(target_os = "linux", not(miri)))]
        {
            // THP guard: prefer MADV_FREE for sub-THP-unit ranges so the kernel
            // does not need to split a 2MB huge page.  Full-segment resets
            // (size >= SEGMENT_SIZE, SEGMENT_ALIGN-aligned base) are always
            // THP-safe and use MADV_DONTNEED for guaranteed zeroing.
            let is_full_segment =
                size >= SEGMENT_SIZE && (ptr as usize).is_multiple_of(SEGMENT_SIZE);
            let advice = if is_full_segment {
                MADV_DONTNEED
            } else {
                // Sub-THP range: use MADV_FREE if available to avoid THP
                // splitting.  Fall back to MADV_DONTNEED if the kernel is
                // too old (MADV_FREE_SUPPORTED == 0 after a failed probe).
                use core::sync::atomic::Ordering;
                let supported = MADV_FREE_SUPPORTED.load(Ordering::Relaxed);
                if supported != 0 {
                    MADV_FREE
                } else {
                    MADV_DONTNEED
                }
            };
            // SAFETY: caller guarantees `ptr` is page-aligned inside an
            // active mapping and `size` is a non-zero multiple of the
            // system page size; madvise never invalidates the mapping.
            let res = unsafe { madvise(ptr as *mut c_void, size, advice) };
            if res == 0 && advice == MADV_FREE {
                // Lazy purge: not a strict reset (may retain contents).
                crate::recorders::record_purge_only(size);
            }
            if res != 0 && advice == MADV_FREE {
                // MADV_FREE failed (unsupported kernel or transient error);
                // mark as unavailable and fall back to MADV_DONTNEED.
                use core::sync::atomic::Ordering;
                MADV_FREE_SUPPORTED.store(0, Ordering::Relaxed);
                let res2 = unsafe { madvise(ptr as *mut c_void, size, MADV_DONTNEED) };
                return res2 == 0;
            }
            res == 0
        }
        #[cfg(all(any(target_os = "macos", target_os = "freebsd"), not(miri)))]
        {
            // SAFETY: same contract as the Linux branch; macOS/FreeBSD
            // MADV_FREE has identical "do not invalidate the mapping"
            // semantics.
            let res = unsafe { madvise(ptr as *mut c_void, size, MADV_FREE) };
            res == 0
        }
        #[cfg(any(
            miri,
            not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
        ))]
        {
            let _ = ptr;
            let _ = size;
            false
        }
    }

    /// Installs a `PROT_NONE` guard region via `mprotect`. Every Unix
    /// target implements `mprotect`, so the impl applies uniformly.
    /// Returns `true` when the kernel confirmed the protection change.
    unsafe fn make_guard(ptr: *mut u8, size: usize) -> bool {
        if ptr.is_null() || size == 0 {
            return false;
        }
        // SAFETY: caller guarantees `ptr` is page-aligned inside an active
        // mapping and `size` is a non-zero multiple of the system page
        // size. `mprotect` does not invalidate the mapping; it only
        // changes access permissions.
        let res = unsafe { mprotect(ptr as *mut c_void, size, PROT_NONE) };
        res == 0
    }

    /// Releases the resident physical pages of the addressed range while
    /// keeping the mapping, so the base `munmap` on release still covers it.
    ///
    /// On Unix there is no separate commit charge (anonymous mappings are
    /// lazily backed under overcommit), so decommitting untouched alignment
    /// slack is largely a no-op; this still drops any resident pages via
    /// `MADV_DONTNEED` (Linux) / `MADV_FREE` (macOS/FreeBSD), matching the
    /// `page_reset` mechanism. Other Unix targets fall back to `false`.
    ///
    /// # Safety
    ///
    /// Same contract as `page_reset`: `ptr` page-aligned inside an active
    /// mapping, `size` a non-zero multiple of the page size, range holding no
    /// live data.
    unsafe fn decommit(ptr: *mut u8, size: usize) -> bool {
        if ptr.is_null() || size == 0 {
            return false;
        }
        // See `page_reset`: Miri doesn't emulate MADV_DONTNEED/MADV_FREE,
        // so it takes the same "unsupported target" `false` fallback.
        #[cfg(all(target_os = "linux", not(miri)))]
        {
            // Prefer MADV_FREE (kernel >= 4.5): pages are reclaimed lazily
            // under pressure without an IPI broadcast, which is critical for
            // scientific workloads that release large segments frequently.
            // MADV_DONTNEED triggers an IPI to all CPUs to flush TLBs; MADV_FREE
            // defers that cost to the kernel's reclaim path or avoids it entirely
            // when the process re-touches the range before reclaim.
            //
            // Probe the kernel once; subsequent calls skip the EINVAL check.
            // SAFETY: see `page_reset`; madvise never invalidates the mapping.
            use core::sync::atomic::Ordering;
            let cached = MADV_FREE_SUPPORTED.load(Ordering::Relaxed);
            let use_free = if cached == 2 {
                // First call: probe whether MADV_FREE is available.
                let r = unsafe { madvise(ptr as *mut c_void, size, MADV_FREE) };
                if r == 0 {
                    MADV_FREE_SUPPORTED.store(1, Ordering::Relaxed);
                    return true;
                }
                // EINVAL means MADV_FREE is not supported on this kernel.
                MADV_FREE_SUPPORTED.store(0, Ordering::Relaxed);
                false
            } else {
                cached == 1
            };
            let advice = if use_free { MADV_FREE } else { MADV_DONTNEED };
            // SAFETY: as for the probe above -- `ptr`/`size` satisfy this
            // function's documented contract, and `advice` is whichever of
            // MADV_FREE / MADV_DONTNEED the probe found this kernel accepts.
            // madvise discards page contents; it never invalidates the
            // mapping, which the caller still owns until `munmap`.
            let res = unsafe { madvise(ptr as *mut c_void, size, advice) };
            if res == 0 && use_free {
                // Track the lazy-purge subset: decommit_calls/bytes are
                // recorded centrally by reset::do_decommit after this
                // returns true; only the MADV_FREE portion needs a separate
                // counter here.
                crate::recorders::record_purge_only(size);
            }
            res == 0
        }
        #[cfg(all(any(target_os = "macos", target_os = "freebsd"), not(miri)))]
        {
            // SAFETY: see `page_reset`.
            let res = unsafe { madvise(ptr as *mut c_void, size, MADV_FREE) };
            res == 0
        }
        #[cfg(any(
            miri,
            not(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))
        ))]
        {
            let _ = ptr;
            let _ = size;
            false
        }
    }
}

#[cfg(all(test, target_os = "linux"))]
#[path = "unix_tests.rs"]
mod tests;