alloc_madvise/
memory.rs

1//! This module provides functionality for aligned memory allocation with support for huge pages and sequential access patterns.
2//!
3//! The `Memory` struct represents an allocated memory block with various allocation flags and methods for allocation and deallocation.
4//!
5//! # Constants
6//! - `ALLOC_FLAGS_NONE`: No special instructions.
7//! - `ALLOC_FLAGS_HUGE_PAGES`: Indicates that huge pages should be used.
8//! - `ALLOC_FLAGS_SEQUENTIAL`: Indicates that memory access is mainly sequential rather than random-access.
9//!
10//! # Structs
11//! - `Memory`: Represents an allocated memory block with methods for allocation, deallocation, and accessing the memory as slices.
12//!
13//! # Methods
14//! - `Memory::allocate`: Allocates memory of the specified number of bytes with optional sequential access pattern and zeroing out.
15//! - `Memory::free`: Frees the allocated memory.
16//! - `Memory::len`: Returns the number of bytes allocated.
17//! - `Memory::is_empty`: Returns whether this instance has zero bytes allocated.
18//! - `Memory::as_ptr`: Returns a pointer to the data buffer.
19//! - `Memory::as_ptr_mut`: Returns a mutable pointer to the data buffer.
20//!
21//! # Macros
22//! - `impl_asref_slice`: Implements `AsRef` and `AsMut` traits for slices of various types.
23//!
24//! # Examples
25//! ```
26//! # use alloc_madvise::Memory;
27//! const FOUR_MEGABYTES: usize = 4 * 1024 * 1024;
28//!
29//! // Allocate 2 MiB of aligned, zeroed-out, sequential read memory.
30//! // The memory will be automatically freed when it leaves scope.
31//! let mut memory = Memory::allocate(FOUR_MEGABYTES, true, true).unwrap();
32//!
33//! // Get a reference to a mutable slice.
34//! let data: &mut [f32] = memory.as_mut();
35//! data[0] = 1.234;
36//! data[1] = 5.678;
37//!
38//! // Get a reference to an immutable slice.
39//! let reference: &[f32] = memory.as_ref();
40//! assert_eq!(reference[0], 1.234);
41//! assert_eq!(reference[1], 5.678);
42//! assert_eq!(reference[2], 0.0);
43//! assert_eq!(reference.len(), memory.len() / std::mem::size_of::<f32>());
44//! ```
45//!
46//! # Safety
47//! - The `madvise` function is used to give advice about the use of memory. The safety of this function relies on the correctness of the pointer and size provided.
48//! - The `free` method ensures that the memory is properly deallocated and the fields are zeroed out to prevent use-after-free errors.
49
50use crate::alignment::AlignmentHint;
51use crate::alloc_free::{alloc_aligned, free_aligned};
52use crate::alloc_result::{AllocResult, AllocationError};
53use libc::madvise;
54use std::ffi::c_void;
55use std::ptr::{null_mut, NonNull};
56
57/// No special instructions.
58const ALLOC_FLAGS_NONE: u32 = 0;
59
60/// Indicates that huge pages should be used.
61const ALLOC_FLAGS_HUGE_PAGES: u32 = 1 << 0;
62
63/// Indicates that memory access is mainly sequential rather than random-access.
64const ALLOC_FLAGS_SEQUENTIAL: u32 = 1 << 1;
65
66/// Allocated memory.
67///
68/// ## Example
69/// ```
70/// # use alloc_madvise::Memory;
71/// const FOUR_MEGABYTES: usize = 4 * 1024 * 1024;
72///
73/// // Allocate 2 MiB of aligned, zeroed-out, sequential read memory.
74/// // The memory will be automatically freed when it leaves scope.
75/// let mut memory = Memory::allocate(FOUR_MEGABYTES, true, true).unwrap();
76///
77/// // Get a reference to a mutable slice.
78/// let data: &mut [f32] = memory.as_mut();
79/// data[0] = 1.234;
80/// data[1] = 5.678;
81///
82/// // Get a reference to an immutable slice.
83/// let reference: &[f32] = memory.as_ref();
84/// assert_eq!(reference[0], 1.234);
85/// assert_eq!(reference[1], 5.678);
86/// assert_eq!(reference[2], 0.0);
87/// assert_eq!(reference.len(), memory.len() / std::mem::size_of::<f32>());
88/// ```
89#[derive(Debug)]
90pub struct Memory {
91    pub(crate) flags: u32,
92    pub(crate) num_bytes: usize,
93    pub(crate) address: *mut c_void,
94}
95
96impl Memory {
97    /// Allocates memory of the specified number of bytes.
98    ///
99    /// The optimal alignment will be determined by the number of bytes provided.
100    /// If the amount of bytes is a multiple of 2MB, Huge/Large Page support is enabled.
101    ///
102    /// ## Arguments
103    /// * `num_bytes` - The number of bytes to allocate.
104    /// * `sequential` - Whether or not the memory access pattern is sequential mostly.
105    /// * `clear` - Whether or not to zero out the allocated memory.
106    pub fn allocate(
107        num_bytes: usize,
108        sequential: bool,
109        clear: bool,
110    ) -> Result<Self, AllocationError> {
111        if num_bytes == 0 {
112            return Err(AllocationError::EmptyAllocation);
113        }
114
115        let alignment = AlignmentHint::new(num_bytes);
116        let ptr = alloc_aligned(num_bytes, alignment.alignment, clear)?;
117
118        let ptr: *mut c_void = ptr.as_ptr().cast::<c_void>();
119
120        let mut advice = if sequential {
121            libc::MADV_SEQUENTIAL
122        } else {
123            libc::MADV_NORMAL
124        };
125
126        let mut flags = if sequential {
127            ALLOC_FLAGS_SEQUENTIAL
128        } else {
129            ALLOC_FLAGS_NONE
130        };
131
132        if alignment.use_huge_pages {
133            advice |= libc::MADV_HUGEPAGE;
134            flags |= ALLOC_FLAGS_HUGE_PAGES;
135        };
136
137        if advice != 0 {
138            // See https://www.man7.org/linux/man-pages/man2/madvise.2.html
139            // SAFETY: `ptr` came from alloc_aligned(num_bytes, alignment)
140            unsafe {
141                madvise(ptr, num_bytes, advice);
142            }
143        }
144
145        Ok(Self::new(AllocResult::Ok, flags, num_bytes, ptr))
146    }
147
148    /// Frees memory of the specified number of bytes.
149    ///
150    /// The memory instance is required to be created by `allocate`.
151    pub fn free(&mut self) {
152        if self.address.is_null() {
153            return;
154        }
155
156        let alignment = AlignmentHint::new(self.num_bytes);
157
158        debug_assert_ne!(self.address, null_mut());
159        let ptr = core::ptr::NonNull::new(self.address);
160
161        if (self.flags & ALLOC_FLAGS_HUGE_PAGES) == ALLOC_FLAGS_HUGE_PAGES {
162            debug_assert!(alignment.use_huge_pages);
163
164            // See https://www.man7.org/linux/man-pages/man2/madvise.2.html
165            // SAFETY: `ptr` came from alloc_aligned(num_bytes, alignment)
166            unsafe {
167                madvise(self.address, self.num_bytes, libc::MADV_FREE);
168            }
169        }
170
171        // SAFETY:
172        // - `ptr` is checked for null before
173        // - `num_bytes` and `alignment` are required to be correct by the caller
174        unsafe {
175            free_aligned(ptr, self.num_bytes, alignment.alignment);
176        }
177
178        // Zero out the fields.
179        self.address = null_mut();
180        self.num_bytes = 0;
181    }
182
183    pub(crate) fn new(
184        status: AllocResult,
185        flags: u32,
186        num_bytes: usize,
187        address: *mut c_void,
188    ) -> Self {
189        debug_assert!(
190            status == AllocResult::Ok && !address.is_null() || address.is_null(),
191            "Found null pointer when allocation status was okay"
192        );
193        Memory {
194            flags,
195            num_bytes,
196            address,
197        }
198    }
199
200    pub(crate) fn from_error(status: AllocResult) -> Self {
201        assert_ne!(status, AllocResult::Ok);
202        Memory {
203            flags: 0,
204            num_bytes: 0,
205            address: null_mut(),
206        }
207    }
208
209    /// Returns the number of bytes allocated.
210    #[inline(always)]
211    pub fn len(&self) -> usize {
212        self.num_bytes
213    }
214
215    /// Returns whether this instance has zero bytes allocated.
216    pub fn is_empty(&self) -> bool {
217        debug_assert!(self.num_bytes > 0 || self.address.is_null());
218        self.num_bytes == 0
219    }
220
221    /// See [`Memory::to_ptr_const`] or [`Memory::to_ptr`].
222    #[inline(always)]
223    #[deprecated(note = "Use to_const_ptr or to_ptr instead", since = "0.5.0")]
224    pub fn as_ptr(&self) -> *const c_void {
225        self.to_ptr_const()
226    }
227
228    /// Returns a pointer to the constant data buffer.
229    ///
230    /// ## Returns
231    /// A valid pointer.
232    ///
233    /// ## Safety
234    /// If the memory is freed while the pointer is in use, access to the address pointed
235    /// at is undefined behavior.
236    #[inline(always)]
237    pub fn to_ptr_const(&self) -> *const c_void {
238        self.address.cast_const()
239    }
240
241    /// See [`Memory::to_ptr_mut`] or [`Memory::to_ptr`].
242    #[inline(always)]
243    #[deprecated(note = "Use to_ptr_mut or to_ptr instead", since = "0.5.0")]
244    pub fn as_ptr_mut(&mut self) -> *mut c_void {
245        self.to_ptr_mut()
246    }
247
248    /// Returns a mutable pointer to the mutable data buffer.
249    ///
250    /// ## Returns
251    /// A valid pointer.
252    ///
253    /// ## Safety
254    /// If the memory is freed while the pointer is in use, access to the address pointed
255    /// at is undefined behavior.
256    #[inline(always)]
257    pub fn to_ptr_mut(&mut self) -> *mut c_void {
258        self.address
259    }
260
261    /// Returns a non-null pointer to the data buffer.
262    ///
263    /// ## Returns
264    /// A pointer that is guaranteed to be non-null if the [`Memory`] was properly
265    /// initialized (i.e., is non-default) and wasn't freed.
266    ///
267    /// ## Safety
268    /// If the memory is freed while the pointer is in use, access to the address pointed
269    /// at is undefined behavior.
270    ///
271    /// # Example
272    ///
273    /// ```
274    /// use alloc_madvise::{Memory, AllocationError};
275    ///
276    /// fn main() -> Result<(), AllocationError> {
277    ///     // Allocate 1024 bytes aligned to 64 bytes
278    ///     const SIZE: usize = 1024;
279    ///     const SEQUENTIAL: bool = true;
280    ///     const CLEAR: bool = true;
281    ///     let memory = Memory::allocate(SIZE, SEQUENTIAL, CLEAR)?;
282    ///     let ptr = memory.to_ptr().expect("pointer was allocated");
283    ///     
284    ///     // Use the allocated memory...
285    ///     assert_ne!(ptr.as_ptr(), std::ptr::null_mut());
286    ///     
287    ///     // Memory is automatically freed when dropped
288    ///     Ok(())
289    /// }
290    /// ```
291    #[inline(always)]
292    pub fn to_ptr(&self) -> Option<NonNull<c_void>> {
293        NonNull::new(self.address)
294    }
295}
296
297impl Default for Memory {
298    fn default() -> Self {
299        Memory::from_error(AllocResult::Empty)
300    }
301}
302
303impl Drop for Memory {
304    #[inline(always)]
305    fn drop(&mut self) {
306        self.free()
307    }
308}
309
310/// Implements AsRef and AsMut
311macro_rules! impl_asref_slice {
312    ($type:ty) => {
313        impl AsRef<[$type]> for Memory {
314            #[inline(always)]
315            fn as_ref(&self) -> &[$type] {
316                let ptr: *const $type = self.address.cast();
317                let len = self.num_bytes / std::mem::size_of::<$type>();
318                unsafe { &*std::ptr::slice_from_raw_parts(ptr, len) }
319            }
320        }
321
322        impl AsMut<[$type]> for Memory {
323            #[inline(always)]
324            fn as_mut(&mut self) -> &mut [$type] {
325                let ptr: *mut $type = self.address.cast();
326                let len = self.num_bytes / std::mem::size_of::<$type>();
327                unsafe { &mut *std::ptr::slice_from_raw_parts_mut(ptr, len) }
328            }
329        }
330    };
331    ($first:ty, $($rest:ty),+) => {
332        impl_asref_slice!($first);
333        impl_asref_slice!($($rest),+);
334    };
335}
336
337impl_asref_slice!(c_void);
338impl_asref_slice!(i8, u8, i16, u16, i32, u32, i64, u64);
339impl_asref_slice!(isize, usize);
340impl_asref_slice!(f32, f64);
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    const TWO_MEGABYTES: usize = 2 * 1024 * 1024;
347    const SIXTY_FOUR_BYTES: usize = 64;
348
349    #[test]
350    fn alloc_4mb_is_2mb_aligned_hugepage() {
351        const SIZE: usize = TWO_MEGABYTES * 2;
352        let memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
353
354        assert_ne!(memory.address, null_mut());
355        assert_eq!((memory.address as usize) % TWO_MEGABYTES, 0);
356        assert_eq!(memory.len(), SIZE);
357        assert!(!memory.is_empty());
358        assert_eq!(
359            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
360            ALLOC_FLAGS_HUGE_PAGES
361        );
362        assert_eq!(
363            memory.flags & ALLOC_FLAGS_SEQUENTIAL,
364            ALLOC_FLAGS_SEQUENTIAL
365        );
366    }
367
368    #[test]
369    fn alloc_4mb_nonsequential_is_2mb_aligned_hugepage() {
370        const SIZE: usize = TWO_MEGABYTES * 2;
371        let memory = Memory::allocate(SIZE, false, false).expect("allocation failed");
372
373        assert_ne!(memory.address, null_mut());
374        assert_eq!((memory.address as usize) % TWO_MEGABYTES, 0);
375        assert_eq!(memory.len(), SIZE);
376        assert!(!memory.is_empty());
377        assert_eq!(
378            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
379            ALLOC_FLAGS_HUGE_PAGES
380        );
381        assert_ne!(
382            memory.flags & ALLOC_FLAGS_SEQUENTIAL,
383            ALLOC_FLAGS_SEQUENTIAL
384        );
385    }
386
387    #[test]
388    fn alloc_2mb_is_2mb_aligned_hugepage() {
389        const SIZE: usize = TWO_MEGABYTES;
390        let memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
391
392        assert_ne!(memory.address, null_mut());
393        assert_eq!((memory.address as usize) % TWO_MEGABYTES, 0);
394        assert_eq!(memory.len(), SIZE);
395        assert!(!memory.is_empty());
396        assert_eq!(
397            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
398            ALLOC_FLAGS_HUGE_PAGES
399        );
400    }
401
402    #[test]
403    fn alloc_1mb_is_64b_aligned() {
404        const SIZE: usize = TWO_MEGABYTES / 2;
405        let memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
406
407        assert_ne!(memory.address, null_mut());
408        assert_eq!((memory.address as usize) % SIXTY_FOUR_BYTES, 0);
409        assert_eq!(memory.len(), SIZE);
410        assert!(!memory.is_empty());
411        assert_ne!(
412            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
413            ALLOC_FLAGS_HUGE_PAGES
414        );
415    }
416
417    #[test]
418    fn alloc_63kb_is_64b_aligned() {
419        const SIZE: usize = 63 * 1024;
420        let memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
421
422        assert_ne!(memory.address, null_mut());
423        assert_eq!((memory.address as usize) % SIXTY_FOUR_BYTES, 0);
424        assert_eq!(memory.len(), SIZE);
425        assert!(!memory.is_empty());
426        assert_ne!(
427            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
428            ALLOC_FLAGS_HUGE_PAGES
429        );
430    }
431
432    #[test]
433    fn alloc_64kb_is_64b_aligned() {
434        const SIZE: usize = 64 * 1024;
435        let memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
436
437        assert_ne!(memory.address, null_mut());
438        assert_eq!((memory.address as usize) % SIXTY_FOUR_BYTES, 0);
439        assert_eq!(memory.len(), SIZE);
440        assert!(!memory.is_empty());
441        assert_ne!(
442            memory.flags & ALLOC_FLAGS_HUGE_PAGES,
443            ALLOC_FLAGS_HUGE_PAGES
444        );
445    }
446
447    #[test]
448    fn alloc_0b_is_not_allocated() {
449        const SIZE: usize = 0;
450        let err = Memory::allocate(SIZE, true, true).expect_err("the allocation was empty");
451
452        assert_eq!(err, AllocationError::EmptyAllocation);
453    }
454
455    #[test]
456    fn deref_works() {
457        const SIZE: usize = TWO_MEGABYTES * 2;
458        let mut memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
459
460        let addr: *mut u8 = memory.to_ptr_mut() as *mut u8;
461        unsafe {
462            *addr = 0x42;
463        }
464
465        let reference: &[u8] = memory.as_ref();
466        assert_eq!(reference[0], 0x42);
467        assert_eq!(reference[1], 0x00);
468        assert_eq!(reference.len(), memory.len());
469    }
470
471    #[test]
472    fn deref_mut_works() {
473        const SIZE: usize = TWO_MEGABYTES * 2;
474        let mut memory = Memory::allocate(SIZE, true, true).expect("allocation failed");
475
476        let addr: &mut [f32] = memory.as_mut();
477        addr[0] = 1.234;
478        addr[1] = 5.678;
479
480        let reference: &[f32] = memory.as_ref();
481        assert_eq!(reference[0], 1.234);
482        assert_eq!(reference[1], 5.678);
483        assert_eq!(reference[2], 0.0);
484        assert_eq!(reference.len(), memory.len() / std::mem::size_of::<f32>());
485    }
486}
alloc_madvise/memory.rs

alloc_madvise/
memory.rs
