mnemosyne/
lib.rs

1//! The Mnemosyne high-performance memory allocator global interface.
2
3#![no_std]
4
5use core::alloc::{GlobalAlloc, Layout};
6use mnemosyne_core::NUM_SIZE_CLASSES;
7use mnemosyne_local::{thread_alloc_layout, thread_free_layout, thread_realloc};
8
9pub use mnemosyne_backend::{
10    CudaDeviceBackend, CudaGddrBackend, CudaHbmBackend, CudaHostPinnedBackend, CudaUnifiedBackend,
11    MemoryBackendWrapper, is_cuda_available,
12};
13pub use mnemosyne_core::{AllocPolicy, StandardPolicy, options::MnemosyneOptions};
14pub use mnemosyne_hardened::{HardenedPolicy, SecurePolicy};
15#[cfg(feature = "branded")]
16pub use mnemosyne_heap::{
17    BrandedBlock, BrandedBox, BrandedCell, BrandedVec, Heap, InvariantLifetime, ReallocError,
18    ReallocFailure, ThreadLocalToken, scope as branded_scope,
19};
20pub use mnemosyne_local::{LocalAllocatorSelector, SizeClassOccupancy, usable_size};
21pub use mnemosyne_prof::{
22    disable_leak_detector, disable_profiling, dump_leaks, dump_profile, enable_leak_detector,
23    enable_profiling, is_leak_detector_enabled, is_profiling_enabled, register_alloc_hook,
24    register_free_hook,
25};
26
27/// Returns the current allocator configuration options snapshot.
28#[inline]
29pub fn get_options() -> MnemosyneOptions {
30    mnemosyne_core::options::get_options()
31}
32/// Configures the allocator runtime settings programmatically.
33///
34/// Modifies the global settings. Can be called at runtime; changes apply
35/// to subsequent allocator operations. If the purge cadence is changed
36/// to a non-zero value and background purger was inactive, starts the
37/// background decay engine thread.
38#[inline]
39pub fn configure(options: MnemosyneOptions) {
40    let old_cadence =
41        mnemosyne_core::options::PURGE_CADENCE_MS.load(core::sync::atomic::Ordering::Acquire);
42    mnemosyne_core::options::set_options(options);
43    mnemosyne_local::mark_options_initialized();
44
45    if options.purge_cadence_ms > 0 && old_cadence == 0 {
46        mnemosyne_decay::init_decay_engine();
47    }
48}
49
50/// Snapshot of Mnemosyne memory mapping and segment cache state.
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub struct MemoryStats {
53    pub current_mapped_bytes: usize,
54    pub peak_mapped_bytes: usize,
55    pub map_calls: usize,
56    pub unmap_calls: usize,
57    /// Number of confirmed backend `page_reset` calls (Linux `MADV_DONTNEED`,
58    /// macOS/FreeBSD `MADV_FREE`, Windows `VirtualAlloc(MEM_RESET)`).
59    pub page_reset_calls: usize,
60    /// Cumulative byte count passed to confirmed `page_reset` calls.
61    pub page_reset_bytes: usize,
62    /// Number of confirmed backend `make_guard` calls (Unix `mprotect(PROT_NONE)`,
63    /// Windows `VirtualProtect(PAGE_NOACCESS)`).
64    pub guard_install_calls: usize,
65    /// Cumulative byte count passed to confirmed `make_guard` calls.
66    pub guard_install_bytes: usize,
67    pub retained_free_segments: usize,
68    pub max_retained_free_segments: usize,
69    pub retained_free_bytes: usize,
70    pub purged_segments: usize,
71    pub purge_calls: usize,
72    pub purged_bytes: usize,
73    /// Number of segments whose physical backing was released by a
74    /// confirmed `page_reset` while the segment itself remained cached
75    /// in the retained pool.
76    pub reset_segments: usize,
77    /// Number of `reset_segment_pool` invocations.
78    pub reset_calls: usize,
79    /// Number of huge blocks currently retained in the huge-allocation cache
80    /// across all NUMA nodes.
81    pub retained_huge_blocks: usize,
82    /// Total bytes of huge blocks currently retained in the huge-allocation
83    /// cache across all NUMA nodes.
84    pub retained_huge_bytes: usize,
85    pub current_thread_live_allocations: usize,
86    pub current_thread_owned_segments: usize,
87    pub cross_thread_reclaimed_blocks: usize,
88    pub page_refills: usize,
89    pub recycled_pages: usize,
90    pub fresh_pages: usize,
91    pub fresh_segments: usize,
92    pub orphan_segments_adopted: usize,
93    pub recycle_sweeps: usize,
94    pub size_class_occupancy: [SizeClassOccupancy; NUM_SIZE_CLASSES],
95}
96
97impl Default for MemoryStats {
98    fn default() -> Self {
99        Self {
100            current_mapped_bytes: 0,
101            peak_mapped_bytes: 0,
102            map_calls: 0,
103            unmap_calls: 0,
104            page_reset_calls: 0,
105            page_reset_bytes: 0,
106            guard_install_calls: 0,
107            guard_install_bytes: 0,
108            retained_free_segments: 0,
109            max_retained_free_segments: 0,
110            retained_free_bytes: 0,
111            purged_segments: 0,
112            purge_calls: 0,
113            purged_bytes: 0,
114            reset_segments: 0,
115            reset_calls: 0,
116            retained_huge_blocks: 0,
117            retained_huge_bytes: 0,
118            current_thread_live_allocations: 0,
119            current_thread_owned_segments: 0,
120            cross_thread_reclaimed_blocks: 0,
121            page_refills: 0,
122            recycled_pages: 0,
123            fresh_pages: 0,
124            fresh_segments: 0,
125            orphan_segments_adopted: 0,
126            recycle_sweeps: 0,
127            size_class_occupancy: [SizeClassOccupancy::default(); NUM_SIZE_CLASSES],
128        }
129    }
130}
131
132/// Returns current Mnemosyne allocator memory counters for a specific backend.
133pub fn memory_stats_generic<B: mnemosyne_arena::HasSegmentPool + LocalAllocatorSelector<B>>()
134-> MemoryStats {
135    let backend = mnemosyne_backend::backend_memory_stats();
136    let arena = mnemosyne_arena::arena_memory_stats::<B>();
137    let local = mnemosyne_local::thread_allocator_stats::<B>();
138    MemoryStats {
139        current_mapped_bytes: backend.current_mapped_bytes,
140        peak_mapped_bytes: backend.peak_mapped_bytes,
141        map_calls: backend.map_calls,
142        unmap_calls: backend.unmap_calls,
143        page_reset_calls: backend.page_reset_calls,
144        page_reset_bytes: backend.page_reset_bytes,
145        guard_install_calls: backend.guard_install_calls,
146        guard_install_bytes: backend.guard_install_bytes,
147        retained_free_segments: arena.retained_free_segments,
148        max_retained_free_segments: arena.max_retained_free_segments,
149        retained_free_bytes: arena.retained_free_bytes,
150        purged_segments: arena.purged_segments,
151        purge_calls: arena.purge_calls,
152        purged_bytes: arena.purged_bytes,
153        reset_segments: arena.reset_segments,
154        reset_calls: arena.reset_calls,
155        retained_huge_blocks: arena.retained_huge_blocks,
156        retained_huge_bytes: arena.retained_huge_bytes,
157        current_thread_live_allocations: local.current_thread_live_allocations,
158        current_thread_owned_segments: local.current_thread_owned_segments,
159        cross_thread_reclaimed_blocks: local.cross_thread_reclaimed_blocks,
160        page_refills: local.page_refills,
161        recycled_pages: local.recycled_pages,
162        fresh_pages: local.fresh_pages,
163        fresh_segments: local.fresh_segments,
164        orphan_segments_adopted: local.orphan_segments_adopted,
165        recycle_sweeps: local.recycle_sweeps,
166        size_class_occupancy: local.size_class_occupancy,
167    }
168}
169
170/// Returns current Mnemosyne allocator memory counters.
171pub fn memory_stats() -> MemoryStats {
172    memory_stats_generic::<mnemosyne_backend::MemoryBackendWrapper>()
173}
174
175/// Purges the global segment pool for a specific backend, releasing all retained/cached segments back to the OS.
176pub fn purge_generic<B: mnemosyne_arena::HasSegmentPool>() {
177    // Safety: Purging the segment pool releases only free segments that are
178    // no longer actively referenced by any thread allocator cache.
179    unsafe {
180        mnemosyne_arena::purge_segment_pool::<B>();
181    }
182}
183
184/// Purges the global segment pool, releasing all retained/cached segments back to the OS.
185pub fn purge() {
186    purge_generic::<mnemosyne_backend::MemoryBackendWrapper>();
187}
188
189/// Asks the OS to drop the physical backing of every retained free
190/// segment for a specific backend without removing them from the cache.
191///
192/// Use this as a lighter-weight RSS-reduction knob than `purge`: the
193/// segment cache stays warm so subsequent allocations skip the OS
194/// mapping syscall, while the resident memory footprint of idle
195/// segments drops to the kernel's demand-fault baseline.
196pub fn reset_generic<B: mnemosyne_arena::HasSegmentPool>() {
197    // Safety: reset_segment_pool drains the retained pool, issues
198    // page_reset on each segment's mapping, and pushes them back into
199    // the cache; no segment is released or accessed by another path.
200    unsafe {
201        mnemosyne_arena::reset_segment_pool::<B>();
202    }
203}
204
205/// Asks the OS to drop the physical backing of every retained free
206/// segment without removing them from the cache.
207pub fn reset() {
208    reset_generic::<mnemosyne_backend::MemoryBackendWrapper>();
209}
210
211/// Triggers a manual background decay and defragmentation cycle across all active memory backends.
212pub fn decay() {
213    mnemosyne_decay::decay_step();
214}
215
216/// The Mnemosyne global allocator structure.
217///
218/// Implements `core::alloc::GlobalAlloc` and routes allocations to the
219/// thread-local cache or global arena.
220pub struct Mnemosyne;
221
222unsafe impl GlobalAlloc for Mnemosyne {
223    // Safety: thread_alloc handles alignment constraints, size validation, and
224    // OS mapping, returning null on failure or a valid memory block pointer on success.
225    #[inline(always)]
226    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
227        // `thread_alloc_layout` rejects `size == 0` through
228        // `is_valid_layout_alloc_request`, so an explicit zero guard here
229        // would be a redundant branch on the hottest path. The
230        // single-source validation returns null for size 0, which is a
231        // valid `GlobalAlloc::alloc` result.
232        // Safety: size and alignment are derived from a valid Layout, and
233        // the returned pointer is verified or null.
234        unsafe {
235            thread_alloc_layout::<StandardPolicy, mnemosyne_backend::MemoryBackendWrapper>(
236                layout.size(),
237                layout.align(),
238            )
239        }
240    }
241
242    // Safety: The ptr must be valid and previously returned by alloc.
243    // thread_free determines the owner segment/page and returns blocks safely.
244    #[inline(always)]
245    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
246        // Safety: thread_free is safe because ptr is guaranteed by the GlobalAlloc
247        // contract to be a valid pointer allocated by this allocator.
248        unsafe {
249            thread_free_layout::<StandardPolicy, mnemosyne_backend::MemoryBackendWrapper>(
250                ptr,
251                layout.size(),
252                layout.align(),
253            )
254        }
255    }
256
257    /// In-place `realloc` shortcut for within-class size changes.
258    ///
259    /// When the new size fits inside the size-class block already
260    /// reserved for `ptr`, return `ptr` unchanged — the allocation
261    /// already covers the request. This eliminates the alloc/copy/free
262    /// round trip that the default `GlobalAlloc::realloc` performs and
263    /// is the common case for `Vec<T>::push` capacity-rounding because
264    /// Mnemosyne rounds small requests up to the next size class.
265    ///
266    /// Falls through to the default `alloc + copy + dealloc` path when:
267    ///   - `ptr` is null (treated as a fresh allocation),
268    ///   - `new_size` is 0 (treated as a deallocation),
269    ///   - `new_size` exceeds the current usable size and a new size
270    ///     class is required,
271    ///   - `new_size` is less than 50% of the current size (capacity-shrink
272    ///     heuristic), forcing a real shrink to release memory.
273    ///
274    /// # Safety
275    ///
276    /// `ptr` must be a previously-returned Mnemosyne allocation with
277    /// the given `layout`; `new_size` must be a valid `Layout` size
278    /// when paired with `layout.align()`. Same contract as the default
279    /// `GlobalAlloc::realloc`.
280    #[inline(always)]
281    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
282        unsafe {
283            thread_realloc::<StandardPolicy, mnemosyne_backend::MemoryBackendWrapper>(
284                ptr, layout, new_size,
285            )
286        }
287    }
288}
289
290/// Generic global allocator that is parameterized by an allocation policy `P` and a memory backend `B`.
291///
292/// This permits zero-cost compile-time configuration of allocator behaviors
293/// (e.g. `SecurePolicy` for memory zeroing and poisoning) and backends (e.g. `CudaUnifiedBackend`).
294pub struct MnemosyneAllocator<
295    P: AllocPolicy,
296    B: mnemosyne_arena::HasSegmentPool + LocalAllocatorSelector<B> = mnemosyne_backend::MemoryBackendWrapper,
297>(core::marker::PhantomData<(P, B)>);
298
299impl<P: AllocPolicy, B: mnemosyne_arena::HasSegmentPool + LocalAllocatorSelector<B>>
300    MnemosyneAllocator<P, B>
301{
302    /// Creates a new `MnemosyneAllocator` with the specified policy and backend.
303    pub const fn new() -> Self {
304        Self(core::marker::PhantomData)
305    }
306}
307
308impl<P: AllocPolicy, B: mnemosyne_arena::HasSegmentPool + LocalAllocatorSelector<B>> Default
309    for MnemosyneAllocator<P, B>
310{
311    fn default() -> Self {
312        Self::new()
313    }
314}
315
316unsafe impl<P: AllocPolicy, B: mnemosyne_arena::HasSegmentPool + LocalAllocatorSelector<B>>
317    GlobalAlloc for MnemosyneAllocator<P, B>
318{
319    // Safety: thread_alloc handles alignment constraints, size validation, and
320    // OS mapping, returning null on failure or a valid memory block pointer on success.
321    #[inline(always)]
322    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
323        // `thread_alloc_layout` rejects `size == 0` via
324        // `is_valid_layout_alloc_request`; the explicit zero guard would be
325        // a redundant hot-path branch (see `Mnemosyne::alloc`).
326        // Safety: size and alignment are derived from a valid Layout, and
327        // the returned pointer is verified or null.
328        unsafe { thread_alloc_layout::<P, B>(layout.size(), layout.align()) }
329    }
330
331    // Safety: The ptr must be valid and previously returned by alloc.
332    // thread_free determines the owner segment/page and returns blocks safely.
333    #[inline(always)]
334    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
335        // Safety: thread_free is safe because ptr is guaranteed by the GlobalAlloc
336        // contract to be a valid pointer allocated by this allocator.
337        unsafe { thread_free_layout::<P, B>(ptr, layout.size(), layout.align()) }
338    }
339
340    /// In-place `realloc` shortcut. See `Mnemosyne::realloc` for the
341    /// full rationale (including capacity-shrink heuristic details); the
342    /// generic variant uses the policy-aware `thread_alloc_layout` and
343    /// `thread_free` paths so a `SecurePolicy` realloc still zeroes/poisons
344    /// the slow-path replacement.
345    ///
346    /// # Safety
347    ///
348    /// Same contract as `Mnemosyne::realloc`.
349    #[inline(always)]
350    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
351        unsafe { thread_realloc::<P, B>(ptr, layout, new_size) }
352    }
353}
354
355/// Aligned scratch pool for temporary buffers.
356pub mod scratch {
357    pub use mnemosyne_arena::scratch::*;
358}
mnemosyne/lib.rs

mnemosyne/
lib.rs
