low_latency_data_structures/
mem.rs

1//! Module for all things memory
2
3use std::mem::MaybeUninit;
4
5use crate::shim::alloc;
6use crate::shim::cell::UnsafeCell;
7
8/// Backing allocator for the queue slot buffer.
9///
10/// Implementors are called once at construction, before any producer or
11/// consumer touches the returned memory, and are expected to hand back a
12/// contiguous, aligned region big enough to hold `items` cells.
13pub trait Allocator {
14    /// Configuration passed to [`allocate`](Self::allocate). Each allocator
15    /// defines its own options type (e.g. mlock toggle, hugepage size).
16    type Options;
17
18    /// Allocates space for `items` cells of `T`.
19    ///
20    /// The returned [`Allocation`] owns the memory and releases it on drop.
21    fn allocate<T>(items: usize, options: Self::Options) -> impl Allocation<T>;
22}
23
24/// Handle for a slot buffer produced by an [`Allocator`].
25///
26/// The pointer returned by [`ptr`](Self::ptr) stays valid for the lifetime of
27/// the `Allocation`; dropping it releases the underlying memory.
28pub trait Allocation<T> {
29    /// Base pointer of the allocated region. Points to `items` contiguous,
30    /// aligned cells (see [`Allocator::allocate`]).
31    fn ptr(&self) -> *mut UnsafeCell<MaybeUninit<T>>;
32}
33
34#[cfg(test)]
35pub(crate) mod test_util {
36    use super::*;
37
38    /// A stub `Allocation` used only in `static_assertions` trait-bound
39    /// checks. Its `ptr()` must never be called.
40    pub(crate) struct NeverAlloc;
41
42    impl<T> Allocation<T> for NeverAlloc {
43        fn ptr(&self) -> *mut UnsafeCell<MaybeUninit<T>> {
44            unreachable!("NeverAlloc is a stub for trait checks only")
45        }
46    }
47}
48
49/// [`Allocator`] backed by Linux hugepages via `mmap(MAP_HUGETLB)`.
50pub mod hugepages {
51    use super::*;
52
53    /// Linux hugepage size, selected via `MAP_HUGE_*` mmap flags.
54    #[derive(Clone, Copy)]
55    pub enum HugepageSize {
56        /// Allocates using libc mmap with hugepage size of 64KB
57        H64KB,
58        /// Allocates using libc mmap with hugepage size of 512KB
59        H512KB,
60        /// Allocates using libc mmap with hugepage size of 1MB
61        H1MB,
62        /// Allocates using libc mmap with hugepage size of 2MB
63        H2MB,
64        /// Allocates using libc mmap with hugepage size of 8MB
65        H8MB,
66        /// Allocates using libc mmap with hugepage size of 16MB
67        H16MB,
68        /// Allocates using libc mmap with hugepage size of 32MB
69        H32MB,
70        /// Allocates using libc mmap with hugepage size of 256MB
71        H256MB,
72        /// Allocates using libc mmap with hugepage size of 512MB
73        H512MB,
74    }
75
76    impl HugepageSize {
77        fn as_usize_bytes(&self) -> usize {
78            match self {
79                HugepageSize::H64KB => 64 * 1024,
80                HugepageSize::H512KB => 512 * 1024,
81                HugepageSize::H1MB => 1024 * 1024,
82                HugepageSize::H2MB => 2 * 1024 * 1024,
83                HugepageSize::H8MB => 8 * 1024 * 1024,
84                HugepageSize::H16MB => 16 * 1024 * 1024,
85                HugepageSize::H32MB => 32 * 1024 * 1024,
86                HugepageSize::H256MB => 256 * 1024 * 1024,
87                HugepageSize::H512MB => 512 * 1024 * 1024,
88            }
89        }
90
91        fn as_mmap_size_flag(&self) -> libc::c_int {
92            match self {
93                HugepageSize::H64KB => libc::MAP_HUGE_64KB,
94                HugepageSize::H512KB => libc::MAP_HUGE_512KB,
95                HugepageSize::H1MB => libc::MAP_HUGE_1MB,
96                HugepageSize::H2MB => libc::MAP_HUGE_2MB,
97                HugepageSize::H8MB => libc::MAP_HUGE_8MB,
98                HugepageSize::H16MB => libc::MAP_HUGE_16MB,
99                HugepageSize::H32MB => libc::MAP_HUGE_32MB,
100                HugepageSize::H256MB => libc::MAP_HUGE_256MB,
101                HugepageSize::H512MB => libc::MAP_HUGE_512MB,
102            }
103        }
104    }
105
106    /// Allocates via `mmap(MAP_HUGETLB | MAP_HUGE_<size>)`. Requires the
107    /// kernel to have enough free hugepages of the requested size.
108    pub struct HugepageAllocator;
109
110    /// Options for [`HugepageAllocator`].
111    #[derive(bon::Builder)]
112    pub struct HugepageAllocatorOptions {
113        /// If true, `mlock` the allocated region so it stays resident.
114        mlock: bool,
115        /// Hugepage size passed to `mmap` via `MAP_HUGE_*` flags.
116        hugepage_size: HugepageSize,
117    }
118
119    impl Allocator for HugepageAllocator {
120        type Options = HugepageAllocatorOptions;
121
122        #[inline]
123        fn allocate<T>(items: usize, options: Self::Options) -> impl Allocation<T> {
124            let layout =
125                alloc::Layout::array::<UnsafeCell<MaybeUninit<T>>>(items).unwrap_or_else(|_| {
126                    panic!(
127                        "Layout::array overflow: {items} x {} bytes exceeds isize::MAX",
128                        std::mem::size_of::<UnsafeCell<MaybeUninit<T>>>()
129                    )
130                });
131            let page_size = options.hugepage_size.as_usize_bytes();
132            // Round up to page size
133            let alloc_size = (layout.size() + page_size - 1) & !(page_size - 1);
134            let ptr = unsafe {
135                libc::mmap(
136                    std::ptr::null_mut(),
137                    alloc_size,
138                    libc::PROT_READ | libc::PROT_WRITE,
139                    libc::MAP_POPULATE
140                        | libc::MAP_PRIVATE
141                        | libc::MAP_ANONYMOUS
142                        | libc::MAP_HUGETLB
143                        | options.hugepage_size.as_mmap_size_flag(),
144                    -1,
145                    0,
146                )
147            };
148            assert_ne!(
149                ptr,
150                libc::MAP_FAILED,
151                "hugepage mmap failed: {}",
152                std::io::Error::last_os_error()
153            );
154            if options.mlock {
155                let rc = unsafe { libc::mlock(ptr.cast(), alloc_size) };
156                assert_ne!(
157                    rc,
158                    -1,
159                    "mlock hugepage allocation failed: {}",
160                    std::io::Error::last_os_error()
161                );
162            }
163            HugePageAllocation {
164                ptr: ptr as *mut UnsafeCell<MaybeUninit<T>>,
165                length: alloc_size,
166            }
167        }
168    }
169
170    pub(crate) struct HugePageAllocation<T> {
171        ptr: *mut UnsafeCell<MaybeUninit<T>>,
172        length: usize,
173    }
174
175    unsafe impl<T: Send> Send for HugePageAllocation<T> {}
176
177    impl<T> Drop for HugePageAllocation<T> {
178        fn drop(&mut self) {
179            let rc = unsafe { libc::munmap(self.ptr.cast(), self.length) };
180            assert_ne!(
181                rc,
182                -1,
183                "failed to munmap hugepage allocation: {}",
184                std::io::Error::last_os_error()
185            );
186        }
187    }
188
189    impl<T> Allocation<T> for HugePageAllocation<T> {
190        #[inline]
191        fn ptr(&self) -> *mut UnsafeCell<MaybeUninit<T>> {
192            self.ptr
193        }
194    }
195}
196
197/// [`Allocator`] backed by the process's `#[global_allocator]`.
198pub mod global {
199    use super::*;
200
201    /// Allocates via the process's `#[global_allocator]`. This is the
202    /// default choice when hugepages aren't available or wanted.
203    pub struct GlobalAllocator;
204
205    /// Options for [`GlobalAllocator`].
206    #[derive(bon::Builder)]
207    pub struct GlobalAllocatorOptions {
208        /// If true, `mlock` the allocated region so it stays resident.
209        mlock: bool,
210    }
211
212    impl Allocator for GlobalAllocator {
213        type Options = GlobalAllocatorOptions;
214
215        #[inline]
216        fn allocate<T>(items: usize, options: Self::Options) -> impl Allocation<T> {
217            let layout =
218                alloc::Layout::array::<UnsafeCell<MaybeUninit<T>>>(items).unwrap_or_else(|_| {
219                    panic!(
220                        "Layout::array overflow: {items} x {} bytes exceeds isize::MAX",
221                        std::mem::size_of::<UnsafeCell<MaybeUninit<T>>>()
222                    )
223                });
224            let ptr = unsafe { alloc::alloc(layout) };
225            if ptr.is_null() {
226                alloc::handle_alloc_error(layout);
227            }
228            if options.mlock {
229                let rc = unsafe { libc::mlock(ptr.cast(), layout.size()) };
230                assert_ne!(
231                    rc,
232                    -1,
233                    "mlock global allocation failed: {}",
234                    std::io::Error::last_os_error()
235                );
236            }
237            GlobalAllocation {
238                ptr: ptr.cast(),
239                layout,
240            }
241        }
242    }
243
244    pub(crate) struct GlobalAllocation<T> {
245        ptr: *mut UnsafeCell<MaybeUninit<T>>,
246        layout: alloc::Layout,
247    }
248
249    unsafe impl<T: Send> Send for GlobalAllocation<T> {}
250
251    impl<T> Drop for GlobalAllocation<T> {
252        fn drop(&mut self) {
253            unsafe { alloc::dealloc(self.ptr.cast(), self.layout) };
254        }
255    }
256
257    impl<T> Allocation<T> for GlobalAllocation<T> {
258        #[inline]
259        fn ptr(&self) -> *mut UnsafeCell<MaybeUninit<T>> {
260            self.ptr
261        }
262    }
263}
264
265#[cfg(test)]
266#[cfg(feature = "tests_loom")]
267pub(crate) mod loom {
268    use super::*;
269
270    pub(crate) struct LoomVecAllocator;
271
272    pub(crate) struct LoomVecAllocatorOptions;
273
274    impl Allocator for LoomVecAllocator {
275        type Options = LoomVecAllocatorOptions;
276
277        #[inline]
278        fn allocate<T>(items: usize, _options: Self::Options) -> impl Allocation<T> {
279            type RealT<T> = UnsafeCell<MaybeUninit<T>>;
280            let mut buffer: Vec<RealT<T>> = Vec::with_capacity(items);
281            for _ in 0..items {
282                buffer.push(UnsafeCell::new(MaybeUninit::uninit()));
283            }
284            let ptr = buffer.as_mut_ptr();
285            std::mem::forget(buffer);
286            LoomVecAllocation {
287                ptr,
288                capacity: items,
289            }
290        }
291    }
292
293    pub(crate) struct LoomVecAllocation<T> {
294        ptr: *mut UnsafeCell<MaybeUninit<T>>,
295        capacity: usize,
296    }
297
298    unsafe impl<T: Send> Send for LoomVecAllocation<T> {}
299
300    impl<T> Drop for LoomVecAllocation<T> {
301        fn drop(&mut self) {
302            unsafe { drop(Vec::from_raw_parts(self.ptr, self.capacity, self.capacity)) };
303        }
304    }
305
306    impl<T> Allocation<T> for LoomVecAllocation<T> {
307        #[inline]
308        fn ptr(&self) -> *mut UnsafeCell<MaybeUninit<T>> {
309            self.ptr
310        }
311    }
312}
313
314/// Type of backing allocation
315pub enum BackingAllocation {
316    /// Allocates using whatever `#[global_allocator]` is set to.
317    Global,
318    /// Allocates via `mmap(MAP_HUGETLB)` at the given hugepage size.
319    Hugepages(hugepages::HugepageSize),
320}
low_latency_data_structures/mem.rs

low_latency_data_structures/
mem.rs
