mnemosyne_arena/segment/alloc/

allocate.rs

//! Handing a segment out: from a pool if one is retained, from the OS
//! otherwise, then initialized and guarded.
//!
//! Returning one is the other direction and lives in [`super::release`],
//! which reaches `try_return_to_pool` here through the parent.

use super::super::alignment::checked_align_up;
use super::super::pool::HasSegmentPool;
#[cfg(feature = "segment-header-guards")]
use super::SEGMENT_HEADER_GUARD_SIZE;
use super::{SEGMENT_MAPPING_SIZE, SEGMENT_TAIL_GUARD_SIZE};
use crate::numa::current_numa_node;
#[cfg(feature = "segment-header-guards")]
use mnemosyne_core::constants::PAGE_SIZE;
use mnemosyne_core::constants::{SEGMENT_ALIGN, SEGMENT_SIZE};
use mnemosyne_core::types::Segment;

/// Pops a retained free segment from the global segment pool and
/// re-initializes it, monomorphized per backend `B` (`#[inline(never)]` keeps
/// this cold pool path out of the hot caller).
///
/// Only the free pool is consulted: a retained segment holds no live
/// allocation, so re-initializing it erases nothing a caller still reads.
///
/// # Safety
///
/// The caller must ensure that the global segment pool contains valid,
/// initialized `Segment` structures. The returned segment (if any) is owned by
/// the caller.
#[inline(never)]
unsafe fn pop_free_segment<B: HasSegmentPool>() -> Option<*mut Segment> {
    let segment = B::global_segment_pool().pop()?;
    // SAFETY: segment points to a valid allocated Segment. We re-initialize
    // the segment to erase stale epoch metadata and reset it for new allocations.
    unsafe {
        let raw_ptr = (*segment).raw_alloc_ptr;
        let node = (*segment).numa_node;
        Segment::initialize(segment, raw_ptr, node);
    }
    Some(segment)
}

/// Helper to return a segment to the global segment pool, monomorphized per
/// backend `B`.
///
/// # Safety
///
/// The `segment` pointer must point to a valid, initialized `Segment` exclusively owned
/// by the caller.
#[inline(always)]
pub(super) unsafe fn try_return_to_pool<B: HasSegmentPool>(segment: *mut Segment) -> bool {
    debug_assert!(
        !segment.is_null(),
        "try_return_to_pool received null segment"
    );
    // SAFETY: by this function's contract `segment` is a valid, initialized,
    // exclusively-owned `Segment`, satisfying `try_push_retained`'s own contract.
    unsafe { B::global_segment_pool().try_push_retained(segment) }
}

/// Non-generic helper to initialize an allocated segment header and establish alignment bounds.
///
/// # Safety
///
/// The caller must guarantee:
/// - `raw_ptr` must point to a valid, exclusive, and page-aligned allocation of size `SEGMENT_MAPPING_SIZE`.
/// - The memory range must be writable to initialize the `Segment` structure.
#[inline(never)]
unsafe fn initialize_allocated_segment(
    raw_ptr: *mut u8,
    numa_node: u32,
) -> Option<(*mut Segment, usize, usize, usize)> {
    let aligned_addr = checked_align_up(raw_ptr as usize, SEGMENT_ALIGN)?;
    let aligned_ptr = raw_ptr.map_addr(|_| aligned_addr).cast::<Segment>();

    // SAFETY: aligned_ptr is within the allocated region.
    unsafe {
        Segment::initialize(aligned_ptr, raw_ptr, numa_node);
    }

    let tail_slack_start = if cfg!(feature = "segment-tail-guards") {
        aligned_addr + SEGMENT_SIZE + SEGMENT_TAIL_GUARD_SIZE
    } else {
        aligned_addr + SEGMENT_SIZE
    };
    let mapping_end = raw_ptr as usize + SEGMENT_MAPPING_SIZE;

    Some((aligned_ptr, aligned_addr, tail_slack_start, mapping_end))
}

/// Returns the head (`[raw_ptr, aligned_addr)`) and tail
/// (`[tail_slack_start, mapping_end)`) alignment-slack subranges of a segment
/// mapping to the OS via `B::decommit`.
///
/// Both slack regions exist only to satisfy `SEGMENT_ALIGN` rounding and never
/// hold allocator or user data. On Windows `VirtualAlloc` eagerly commits the
/// whole mapping, so decommitting drops the slack's commit charge (up to
/// ~`SEGMENT_ALIGN` ≈ 2 MiB of head slack per segment) for the mapping's
/// lifetime; on Unix the slack is lazily backed, so this is typically a no-op.
/// Best-effort: a backend without decommit support (`SUPPORTS_DECOMMIT ==
/// false`) skips entirely, and both subranges stay inside the reservation,
/// which the base `B::deallocate(raw_ptr, ..)` releases in full.
///
/// # Safety
///
/// `raw_ptr` must name the base of the live backend mapping that contains both
/// subranges, with `raw_ptr as usize <= aligned_addr` and
/// `tail_slack_start <= mapping_end`. Neither subrange may hold allocator or
/// user data, and the head bounds are page-aligned (`raw_ptr` comes from the
/// backend allocator and `aligned_addr` is a `SEGMENT_ALIGN` multiple), as
/// `decommit` requires.
#[inline]
pub(crate) unsafe fn decommit_mapping_slack<B: mnemosyne_core::MemoryBackend>(
    raw_ptr: *mut u8,
    aligned_addr: usize,
    tail_slack_start: usize,
    mapping_end: usize,
) {
    if !B::SUPPORTS_DECOMMIT {
        return;
    }
    let head_slack = aligned_addr - raw_ptr as usize;
    if head_slack > 0 {
        // SAFETY: per this function's contract, `[raw_ptr, aligned_addr)` is a
        // page-aligned, data-free subrange of the live mapping.
        let _ = unsafe { B::decommit(raw_ptr, head_slack) };
    }
    if tail_slack_start < mapping_end {
        // SAFETY: per this function's contract, `[tail_slack_start,
        // mapping_end)` is a data-free subrange of the live mapping. Deriving
        // the subrange pointer from `raw_ptr` retains the mapping provenance.
        let tail_slack = raw_ptr.map_addr(|_| tail_slack_start);
        let _ = unsafe { B::decommit(tail_slack, mapping_end - tail_slack_start) };
    }
}

/// A segment handed out by [`acquire_segment`], tagged with whether it may
/// still hold live allocations.
///
/// The distinction decides what the new owner may do with the segment: a
/// [`Free`](Self::Free) segment holds nothing, so it may be carved up or handed
/// back through [`deallocate_segment`](super::deallocate_segment); an
/// [`Orphan`](Self::Orphan) holds blocks a dead thread cache allocated and
/// other threads may still read, write, or free, so it may be adopted but
/// never returned to the free pool or the OS until every one of its pages is
/// empty.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AcquiredSegment {
    /// A freshly mapped or pool-reinitialized segment with no live allocation.
    Free(*mut Segment),
    /// A segment from the orphan pool, handed out as is with its live
    /// allocations, per-page free lists, and pending cross-thread frees.
    Orphan(*mut Segment),
}

/// Allocates an empty, aligned segment, either from the free pool or from the
/// OS.
///
/// Never returns an orphan: an orphaned segment still holds live allocations,
/// and a caller of this function may hand the segment straight back through
/// [`deallocate_segment`](super::deallocate_segment), which would put those
/// allocations in the free pool for reuse or for a purge to unmap. Thread
/// caches that can adopt an orphan use [`acquire_segment`].
///
/// # Monomorphization and ZST Static Routing
///
/// The backend parameter `B` acts as a Zero-Sized Type (ZST) policy marker. Calls
/// to this function are fully monomorphized by the compiler into direct machine-code
/// calls for the target backend, preserving the zero-cost abstraction invariant.
///
/// # Safety
///
/// This function is unsafe because it allocates virtual memory from the OS/backend,
/// aligns and initializes a raw `Segment` pointer. The caller must guarantee:
/// - The backend `B` must be a valid implementor of `HasSegmentPool`.
/// - The returned pointer must eventually be returned to the pool via
///   `deallocate_segment` or released to the OS via `release_segment_mapping`.
#[inline]
pub unsafe fn allocate_segment<B: HasSegmentPool>() -> Option<*mut Segment> {
    // SAFETY: forwarded to the caller — the free pool holds valid segments.
    if let Some(segment) = unsafe { pop_free_segment::<B>() } {
        return Some(segment);
    }
    // SAFETY: forwarded to the caller.
    unsafe { map_fresh_segment::<B>() }
}

/// Acquires a segment for a thread cache: a retained free segment first, then
/// an orphan to adopt, then a fresh OS mapping.
///
/// # Safety
///
/// As [`allocate_segment`]. Additionally, an [`AcquiredSegment::Orphan`] must
/// be adopted — its live pages kept intact — or pushed back to the orphan
/// pool; it must not reach [`deallocate_segment`](super::deallocate_segment)
/// or [`release_segment_mapping`](super::release_segment_mapping) while any of
/// its pages holds a live allocation.
#[inline]
pub unsafe fn acquire_segment<B: HasSegmentPool>() -> Option<AcquiredSegment> {
    // SAFETY: forwarded to the caller — the free pool holds valid segments.
    if let Some(segment) = unsafe { pop_free_segment::<B>() } {
        return Some(AcquiredSegment::Free(segment));
    }
    if let Some(segment) = B::global_orphan_pool().pop() {
        return Some(AcquiredSegment::Orphan(segment));
    }
    // SAFETY: forwarded to the caller.
    unsafe { map_fresh_segment::<B>() }.map(AcquiredSegment::Free)
}

/// Maps, aligns, and initializes a fresh segment from the OS, purging the free
/// pool and retrying once when the first mapping request fails.
///
/// # Safety
///
/// As [`allocate_segment`].
#[inline(never)]
unsafe fn map_fresh_segment<B: HasSegmentPool>() -> Option<*mut Segment> {
    // We allocate twice the segment size to ensure we can find an aligned boundary.
    // SAFETY: SEGMENT_MAPPING_SIZE is non-zero and aligned. We call B::allocate.
    let mut raw_ptr = unsafe { B::allocate(SEGMENT_MAPPING_SIZE) };
    if raw_ptr.is_null() {
        // First OS allocation failed: release every retained free segment
        // back to the OS to reclaim the address space / commit charge a
        // transient working-set spike may be holding, then retry exactly
        // once. Bounded to one purge and one retry: a still-exhausted OS
        // fails fast on the second call rather than looping.
        //
        // SAFETY: segments retained in the pool at this point are
        // exclusively pool-owned — a segment a thread-local allocator still
        // references is "active", not "retained" — so releasing them
        // invalidates no live pointer, and the retried `B::allocate` call
        // follows the same contract as the first.
        unsafe { super::release::purge_segment_pool::<B>() };
        raw_ptr = unsafe { B::allocate(SEGMENT_MAPPING_SIZE) };
        // Surfaces the recovery attempt and its outcome as a stats counter
        // (`ArenaMemoryStats::oom_retries` / `oom_retry_successes`) — the
        // crate has no tracing dependency, so telemetry follows the
        // existing `purge_calls`/`reset_calls` counter convention.
        B::global_segment_pool().record_oom_retry(!raw_ptr.is_null());
        if raw_ptr.is_null() {
            return None;
        }
    }

    let numa_node = current_numa_node();
    // SAFETY: `raw_ptr` is the non-null `SEGMENT_MAPPING_SIZE` mapping just
    // returned by `B::allocate`, which is exclusively owned and writable —
    // exactly `initialize_allocated_segment`'s contract.
    let (aligned_ptr, aligned_addr, tail_slack_start, mapping_end) =
        match unsafe { initialize_allocated_segment(raw_ptr, numa_node) } {
            Some(val) => val,
            None => {
                // SAFETY: Releasing raw memory back to the backend because alignment check overflowed.
                let _released = unsafe { B::deallocate(raw_ptr, SEGMENT_MAPPING_SIZE) };
                return None;
            }
        };

    // Bind the aligned segment mapping to the allocating thread's NUMA node.
    // `mbind(MPOL_BIND)` enforces first-touch locality so pages are allocated
    // from the correct socket rather than relying on the OS default policy.
    // This call goes to the kernel only once per fresh segment (cold path);
    // segments recycled from the pool retain their prior NUMA binding and never
    // reach this function (they return early via `pop_free_segment`).
    //
    // SAFETY: `aligned_ptr` points to the segment-aligned base of the
    // `SEGMENT_MAPPING_SIZE` mapping that is exclusively owned at this point
    // (not yet published to any pool or thread), `SEGMENT_SIZE` is the
    // usable extent of the aligned portion, and `numa_node` was read from
    // Themis's thread-local cache immediately before this call.
    unsafe {
        crate::numa::bind_segment_to_numa_node(
            raw_ptr.map_addr(|_| aligned_ptr as usize),
            SEGMENT_SIZE,
            numa_node,
        )
    };

    #[cfg(feature = "segment-header-guards")]
    {
        if B::SUPPORTS_MAKE_GUARD {
            // Install a header guard at the end of Page 0.
            // Underflows (backward OOB writes) from Page 1 land in this guard region
            // instead of overwriting the segment metadata at the start of Page 0.
            //
            // SAFETY: aligned_addr + PAGE_SIZE - SEGMENT_HEADER_GUARD_SIZE is inside the mapping
            // and Page 0 is reserved strictly for Segment metadata (ending far before the guard).
            let header_guard_addr = aligned_addr + PAGE_SIZE - SEGMENT_HEADER_GUARD_SIZE;
            let header_guard = raw_ptr.map_addr(|_| header_guard_addr);
            let _guarded = unsafe { B::make_guard(header_guard, SEGMENT_HEADER_GUARD_SIZE) };
        }
    }

    #[cfg(feature = "segment-tail-guards")]
    {
        if B::SUPPORTS_MAKE_GUARD {
            // Install a tail guard immediately after the segment's user-page
            // region. Forward OOB writes that walk past Page 31 land in this
            // guard region instead of an unrelated mapping. The address lives
            // inside the `SEGMENT_MAPPING_SIZE - SEGMENT_SIZE` slack the arena
            // reserves to satisfy `SEGMENT_ALIGN` rounding, so it is always
            // part of the same backend allocation and is released together
            // with the segment by `B::deallocate(raw_ptr, SEGMENT_MAPPING_SIZE)`.
            // The install is best-effort: a backend without a `make_guard`
            // implementation (default `false`) or a kernel that declines the
            // request (e.g. macOS-arm64 where the OS page size exceeds 4 KiB)
            // silently skips, leaving the slack accessible. Backend telemetry
            // (`guard_install_calls`) surfaces the actual install count.
            //
            // SAFETY: aligned_addr + SEGMENT_SIZE is inside the raw mapping
            // because slack-after >= OS_PAGE_SIZE >= SEGMENT_TAIL_GUARD_SIZE on
            // supported targets. `make_guard` never invalidates the mapping.
            let tail_guard_addr = aligned_addr + SEGMENT_SIZE;
            let tail_guard = raw_ptr.map_addr(|_| tail_guard_addr);
            let _guarded = unsafe { B::make_guard(tail_guard, SEGMENT_TAIL_GUARD_SIZE) };
        }
    }

    // The mapping over-reserves `SEGMENT_MAPPING_SIZE = 2 * SEGMENT_SIZE` so a
    // `SEGMENT_ALIGN`-aligned base can always be found; return the resulting
    // head and tail slack to the OS (guard pages installed above lie outside
    // both slack subranges, so ordering relative to the guards is immaterial).
    //
    // SAFETY: `[raw_ptr, aligned_addr)` precedes the header and
    // `[tail_slack_start, mapping_end)` succeeds the segment pages (and tail
    // guard, when enabled), so neither holds allocator data; both stay inside
    // the live reservation covered by the base release, and the head bounds
    // are page-aligned (`raw_ptr` from `allocate`, `aligned_addr` a
    // `SEGMENT_ALIGN` multiple).
    unsafe { decommit_mapping_slack::<B>(raw_ptr, aligned_addr, tail_slack_start, mapping_end) };

    Some(aligned_ptr)
}