rts_alloc/
header.rs

1use crate::cache_aligned::{CacheAligned, CacheAlignedU64};
2use crate::size_classes::NUM_SIZE_CLASSES;
3use crate::sync::{AtomicU32, AtomicU64, AtomicU8, AtomicUsize};
4
5/// Unique identifier for rts-alloc in shared memory.
6pub const MAGIC: u64 = u64::from_be_bytes(*b"\0rtsaloc");
7pub const VERSION_MAJOR: u16 = 4;
8pub const VERSION_PATCH: u16 = 0;
9/// Shared-memory format version, packed as `[u16 major | u16 patch]`.
10pub const VERSION: u32 = (VERSION_MAJOR as u32) << 16 | VERSION_PATCH as u32;
11
12#[repr(C)]
13pub struct WorkerLocalListPartialFullHeads {
14    pub partial: AtomicU32,
15    pub full: AtomicU32,
16}
17
18#[repr(C, align(64))]
19pub struct WorkerLocalListHeads {
20    pub claimed: AtomicU8,
21    pub outstanding_allocation_bytes: AtomicU64,
22    pub heads: [WorkerLocalListPartialFullHeads; NUM_SIZE_CLASSES],
23    pub remote_free_head: CacheAligned<AtomicUsize>,
24}
25
26#[repr(C)]
27pub struct Header {
28    pub magic: AtomicU64,
29    pub version: u32,
30    /// Maximum number of workers that can use this allocator.
31    pub num_workers: u32,
32    /// Number of slabs in the allocator.
33    pub num_slabs: u32,
34    /// The size in bytes of each slab.
35    pub slab_size: u32,
36
37    /// The offset in bytes to the free list elements.
38    pub free_list_elements_offset: u32,
39    /// The offset in bytes to the slab shared metadata.
40    pub slab_shared_meta_offset: u32,
41    /// The offset in bytes to the slab free stacks.
42    pub slab_free_stacks_offset: u32,
43    /// The offset in bytes to the slabs.
44    pub slabs_offset: u32,
45
46    /// The head of the global free list.
47    ///
48    /// Packed as `[u32 generation | u32 index]` to prevent ABA races.
49    pub global_free_list_head: CacheAlignedU64,
50    /// The heads of the per-worker local free lists.
51    pub worker_local_list_heads: [WorkerLocalListHeads; 0],
52}
53
54// Layout of the allocator.
55// Padding used to ensure proper alignment between components.
56//
57// [header]
58// [worker_local_list_heads; num_workers]
59// [free_list_elements; num_slabs]
60// [slab_shared_meta]
61// [slab_free_stacks]
62// [slabs]
63//
64// header:
65//     - contains metadata about the allocator as a whole.
66//
67// worker_local_list_heads:
68//     - contains heads of the worker local lists.
69//     - each worker has its own set of heads.
70//     - each worker has a partial and full head for each size class.
71//     - the heads store indexes into the free list elements.
72//     - each worker also has a cache-line-aligned remote-free head that tracks
73//       allocation offsets pending cleanup by that worker.
74//     - NULL_U32 is used for slab-list heads; NULL_USIZE is used for the
75//       remote-free offset head.
76//
77// free_list_elements:
78//     - list of free list elements, one per slab.
79//     - the free list elements, in conjunction with the `global_free_list_head`
80//       and `worker_local_list_heads`, form linked lists of slabs, in various states.
81//     - it is NOT valid for a slab to be in multiple lists at the same time.
82//     - NULL_U32 is used to indicate a null pointer in the linked list.
83//
84// slab_shared_meta:
85//     - shared metadata for each slab.
86//
87// slab_free_stacks:
88//     - each slab has its own free stack.
89//     - the free stack is used to track free indices within the slab.
90//
91// slabs:
92//     - the slabs themselves, containing chunks of `slab_size` bytes each.
93//     - guaranteed to be offset a multiple of `slab_size` bytes from the
94//       start of the file.
95pub mod layout {
96    use crate::{
97        align::round_to_next_alignment_of,
98        free_stack::FreeStack,
99        header::{Header, WorkerLocalListHeads},
100        linked_list_node::LinkedListNode,
101        size_classes::MIN_SIZE,
102        slab_meta::SlabMeta,
103    };
104    #[derive(Debug)]
105    pub struct AllocatorLayout {
106        /// The number of slabs in the allocator.
107        pub num_slabs: u32,
108        /// The offset in bytes to the free list elements.
109        pub free_list_elements_offset: u32,
110        /// The offset in bytes to the slab shared metadata.
111        pub slab_shared_meta_offset: u32,
112        /// The offset in bytes to the slab free stacks.
113        pub slab_free_stacks_offset: u32,
114        /// The offset in bytes to the slabs.
115        pub slabs_offset: u32,
116    }
117
118    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
119    pub struct WorkerLimits {
120        pub meta_slabs: u32,
121        pub usable_slabs: u32,
122        pub max_workers: u32,
123    }
124
125    pub fn max_workers(file_size: usize, slab_size: u32, min_workers: u32) -> Option<WorkerLimits> {
126        if slab_size == 0 {
127            return None;
128        }
129        let slab_size_usize = slab_size as usize;
130        let total_slabs = (file_size / slab_size_usize) as u32;
131        if total_slabs == 0 {
132            return None;
133        }
134
135        let mut meta_slabs: u32 = 1;
136        loop {
137            let usable_slabs = total_slabs.checked_sub(meta_slabs)?;
138            if usable_slabs == 0 {
139                return None;
140            }
141            let base_layout = layout_for(0, slab_size, usable_slabs);
142            let base_meta_bytes = base_layout.slabs_offset as usize;
143
144            let needed_meta_slabs = base_meta_bytes.div_ceil(slab_size_usize) as u32;
145            if needed_meta_slabs > meta_slabs {
146                meta_slabs = needed_meta_slabs;
147                continue;
148            }
149
150            let slack_bytes = meta_slabs as usize * slab_size_usize - base_meta_bytes;
151            let per_worker_bytes = core::mem::size_of::<WorkerLocalListHeads>();
152            let max_workers = (slack_bytes / per_worker_bytes) as u32;
153            if max_workers < min_workers {
154                meta_slabs = meta_slabs.saturating_add(1);
155                continue;
156            }
157
158            return Some(WorkerLimits {
159                meta_slabs,
160                usable_slabs,
161                max_workers,
162            });
163        }
164    }
165
166    fn layout_for(num_workers: u32, slab_size: u32, num_slabs: u32) -> AllocatorLayout {
167        let mut offset = header_size();
168        offset += worker_local_list_heads_size(num_workers);
169        offset = pad_for_free_list_elements(offset);
170        let free_list_elements_offset = offset as u32;
171        offset += free_list_elements_size(num_slabs);
172        offset = pad_for_slab_meta(offset);
173        let slab_shared_meta_offset = offset as u32;
174        offset += slab_meta_size(num_slabs);
175        offset = pad_for_slab_free_stacks(offset);
176        let slab_free_stacks_offset = offset as u32;
177        offset += free_stacks_size(num_slabs, slab_size);
178        let slabs_offset = pad_for_slabs(offset, slab_size) as u32;
179
180        AllocatorLayout {
181            num_slabs,
182            free_list_elements_offset,
183            slab_shared_meta_offset,
184            slab_free_stacks_offset,
185            slabs_offset,
186        }
187    }
188
189    pub fn layout_for_num_slabs(
190        num_workers: u32,
191        slab_size: u32,
192        num_slabs: u32,
193    ) -> AllocatorLayout {
194        layout_for(num_workers, slab_size, num_slabs)
195    }
196
197    /// The size of the header in bytes.
198    pub const fn header_size() -> usize {
199        core::mem::size_of::<Header>()
200    }
201
202    /// The size of the worker local list heads in bytes.
203    pub const fn worker_local_list_heads_size(num_workers: u32) -> usize {
204        core::mem::size_of::<WorkerLocalListHeads>() * num_workers as usize
205    }
206
207    /// Update offset to padd for free list elements.
208    pub const fn pad_for_free_list_elements(offset: usize) -> usize {
209        const FREE_LIST_ELEMENT_ALIGNMENT: usize = core::mem::align_of::<LinkedListNode>();
210        round_to_next_alignment_of::<FREE_LIST_ELEMENT_ALIGNMENT>(offset)
211    }
212
213    /// The size of the free list elements in bytes.
214    pub const fn free_list_elements_size(num_slabs: u32) -> usize {
215        core::mem::size_of::<LinkedListNode>() * num_slabs as usize
216    }
217
218    /// Update offset to pad for slab shared metadata.
219    pub const fn pad_for_slab_meta(offset: usize) -> usize {
220        const SLAB_META_ALIGNMENT: usize = core::mem::align_of::<SlabMeta>();
221        round_to_next_alignment_of::<SLAB_META_ALIGNMENT>(offset)
222    }
223
224    /// The size of the slab meta in bytes with trailing padding.
225    pub const fn slab_meta_size(num_slabs: u32) -> usize {
226        core::mem::size_of::<SlabMeta>() * num_slabs as usize
227    }
228
229    /// Update offset to pad for slab free stacks.
230    pub const fn pad_for_slab_free_stacks(offset: usize) -> usize {
231        const FREE_STACK_ALIGNMENT: usize = core::mem::align_of::<FreeStack>();
232        round_to_next_alignment_of::<FREE_STACK_ALIGNMENT>(offset)
233    }
234
235    /// The size of an individual free stack in bytes.
236    pub const fn single_free_stack_size(slab_size: u32) -> usize {
237        let max_capacity = slab_size / MIN_SIZE;
238        FreeStack::byte_size(max_capacity as u16)
239    }
240
241    /// The size of the free stacks in bytes WITHOUT trailing padding.
242    pub const fn free_stacks_size(num_slabs: u32, slab_size: u32) -> usize {
243        single_free_stack_size(slab_size) * num_slabs as usize
244    }
245
246    /// Update offset to the next multiple of `slab_size`.
247    pub const fn pad_for_slabs(offset: usize, slab_size: u32) -> usize {
248        debug_assert!(slab_size.is_power_of_two());
249        let slab_size = slab_size as usize;
250        (offset + slab_size - 1) & !(slab_size - 1)
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257
258    #[test]
259    fn test_worker_local_list_heads_layout() {
260        assert_eq!(core::mem::align_of::<WorkerLocalListHeads>(), 64);
261        assert_eq!(core::mem::size_of::<WorkerLocalListHeads>(), 128);
262        assert_eq!(core::mem::align_of::<CacheAligned<AtomicUsize>>(), 64);
263        assert_eq!(
264            core::mem::offset_of!(WorkerLocalListHeads, remote_free_head),
265            64
266        );
267    }
268
269    #[test]
270    fn test_layout() {
271        let num_workers = 4;
272        let num_slabs = 8;
273        let slab_size = 4096;
274
275        let mut offset = layout::header_size();
276        assert_eq!(offset, core::mem::size_of::<Header>());
277        assert_eq!(
278            offset,
279            core::mem::offset_of!(Header, worker_local_list_heads)
280        );
281        assert_eq!(offset, 128);
282
283        offset += layout::worker_local_list_heads_size(num_workers);
284        assert_eq!(offset, 640);
285
286        offset = layout::pad_for_free_list_elements(offset);
287        assert_eq!(offset, 640);
288
289        offset += layout::free_list_elements_size(num_slabs);
290        assert_eq!(offset, 736);
291
292        offset = layout::pad_for_slab_meta(offset);
293        assert_eq!(offset, 736);
294
295        offset += layout::slab_meta_size(num_slabs);
296        assert_eq!(offset, 864);
297
298        offset = layout::pad_for_slab_free_stacks(offset);
299        assert_eq!(offset, 864);
300
301        offset += layout::free_stacks_size(num_slabs, slab_size);
302        assert_eq!(offset, 1152);
303
304        offset = layout::pad_for_slabs(offset, slab_size);
305        assert_eq!(offset, 4096);
306    }
307}
rts_alloc/header.rs

rts_alloc/
header.rs
