rulloc/allocator.rs
1use std::{
2 alloc::{AllocError, Allocator, GlobalAlloc, Layout},
3 ptr::{self, NonNull},
4 sync::Mutex,
5};
6
7use crate::{bucket::Bucket, realloc::Realloc, AllocResult};
8
9/// This is the main allocator, it contains multiple allocation buckets for
10/// different sizes. Once you've read [`crate::header`], [`crate::block`],
11/// [`crate::region`], [`crate::freelist`] and [`crate::bucket`], this is where
12/// the circle gets completed:
13///
14/// ```text
15/// Next Free Block Next Free Block
16/// +------------------------------------+ +-----------------------+
17/// | | | |
18/// +--------+-------|----------------+ +--------+---|---|-----------------------|-----+
19/// | | +-----|-+ +-------+ | | | +-|---|-+ +-------+ +---|---+ |
20/// buckets[0] -> | Region | | Free | -> | Block | | ---> | Region | | Free | -> | Block | -> | Free | |
21/// | | +-------+ +-------+ | | | +-------+ +-------+ +-------+ |
22/// +--------+------------------------+ +--------+-------------------------------------+
23///
24/// Next Free Block
25/// +----------------------------------------+
26/// | |
27/// +--------+------------------|-----+ +--------+------------------|------------------+
28/// | | +-------+ +---|---+ | | | +-------+ +---|---+ +-------+ |
29/// buckets[1] -> | Region | | Block | -> | Free | | ---> | Region | | Block | -> | Free | -> | Block | |
30/// | | +-------+ +-------+ | | | +-------+ +-------+ +-------+ |
31/// +--------+------------------------+ +--------+-------------------------------------+
32///
33/// .......................................................................................................
34///
35/// Next Free Block
36/// +---------------------------+
37/// | |
38/// +--------+------------------|-----+ +--------+-----|-------------------------------+
39/// | | +-------+ +---|---+ | | | +---|---+ +-------+ +-------+ |
40/// buckets[N-1] -> | Region | | Block | -> | Free | | ---> | Region | | Free | -> | Block | -> | Block | |
41/// | | +-------+ +-------+ | | | +-------+ +-------+ +-------+ |
42/// +--------+------------------------+ +--------+-------------------------------------+
43///
44/// Next Free Block
45/// +-----------------------------------------------------+
46/// | |
47/// +--------+-----|------------------+ +--------+------------------|------------------+
48/// | | +---|---+ +-------+ | | | +-------+ +---|---+ +-------+ |
49/// dyn_bucket -> | Region | | Free | -> | Block | | ---> | Region | | Block | -> | Free | -> | Block | |
50/// | | +-------+ +-------+ | | | +-------+ +-------+ +-------+ |
51/// +--------+------------------------+ +--------+-------------------------------------+
52/// ```
53///
54/// Number of buckets and size of each bucket can be configured at compile
55/// time. This struct is not thread safe and it also needs mutable borrows to
56/// operate, so it has to be wrapped in some container like [`Mutex`] to satisfy
57/// [`std::alloc::Allocator`] trait. See [`MmapAllocator`] for the public API.
58///
59/// # Drop
60///
61/// This struct doesn't implement [`Drop`] because region unmapping is
62/// implemented by [`Bucket`]. If we don't implement [`Drop`], the compiler will
63/// just call [`Drop::drop`] on all the struct members one by one, so all the
64/// buckets will be dropped automatically.
65struct InternalAllocator<const N: usize> {
66 /// Size of each bucket, in bytes.
67 sizes: [usize; N],
68 /// Fixed size buckets.
69 buckets: [Bucket; N],
70 /// Any allocation request of `size > sizes[N - 1]` will use this bucket.
71 dyn_bucket: Bucket,
72}
73
74impl<const N: usize> InternalAllocator<N> {
75 /// Builds a new allocator configured with the given bucket sizes.
76 pub const fn with_bucket_sizes(sizes: [usize; N]) -> Self {
77 const BUCKET: Bucket = Bucket::new();
78 InternalAllocator::<N> {
79 sizes,
80 buckets: [BUCKET; N],
81 dyn_bucket: Bucket::new(),
82 }
83 }
84
85 /// Returns the index of the [`Bucket`] where `layout` should be allocated.
86 fn bucket_index_of(&self, layout: Layout) -> usize {
87 for (i, size) in self.sizes.iter().enumerate() {
88 if layout.size() <= *size {
89 return i;
90 }
91 }
92
93 self.buckets.len()
94 }
95
96 /// Returns a mutable reference to the [`Bucket`] at `index`.
97 fn bucket_mut(&mut self, index: usize) -> &mut Bucket {
98 if index == self.buckets.len() {
99 &mut self.dyn_bucket
100 } else {
101 &mut self.buckets[index]
102 }
103 }
104
105 /// Returns a mutable reference to the [`Bucket`] where `layout` should be
106 /// allocated.
107 #[inline]
108 fn dispatch(&mut self, layout: Layout) -> &mut Bucket {
109 self.bucket_mut(self.bucket_index_of(layout))
110 }
111
112 /// Returns an address where `layout.size()` bytes can be safely written or
113 /// [`AllocError`] if it fails to allocate.
114 #[inline]
115 pub unsafe fn allocate(&mut self, layout: Layout) -> AllocResult {
116 self.dispatch(layout).allocate(layout)
117 }
118
119 /// Deallocates the memory block at `address`.
120 #[inline]
121 pub unsafe fn deallocate(&mut self, address: NonNull<u8>, layout: Layout) {
122 // We can find the bucket that has allocated the pointer because we also
123 // know the layout. If the allocator trait changes and the layout is
124 // no longer required, we can still obtain the block header given any
125 // valid address and check the size to find the bucket. Let's hope it
126 // doesn't change though, layouts are useful information for allocators!
127 self.dispatch(layout).deallocate(address, layout)
128 }
129
130 /// Reallocation algorithm. Whether shrinking or growing, we'll try to
131 /// preserve the maximum allocation size of each bucket as it was defined
132 /// when creating the struct. So if `new_layout` should be allocated in a
133 /// different bucket, we'll move the user contents there. Otherwise just
134 /// delegate the call to the current bucket and handle reallocation
135 /// internally.
136 pub unsafe fn reallocate(&mut self, realloc: &Realloc) -> AllocResult {
137 let current_bucket = self.bucket_index_of(realloc.old_layout);
138 let ideal_bucket = self.bucket_index_of(realloc.new_layout);
139
140 if current_bucket == ideal_bucket {
141 return self.bucket_mut(current_bucket).reallocate(realloc);
142 }
143
144 let new_address = self.bucket_mut(ideal_bucket).allocate(realloc.new_layout)?;
145 ptr::copy_nonoverlapping(
146 realloc.address.as_ptr(),
147 new_address.as_mut_ptr(),
148 realloc.count(),
149 );
150 self.bucket_mut(current_bucket)
151 .deallocate(realloc.address, realloc.old_layout);
152
153 Ok(new_address)
154 }
155}
156
157/// This struct exposes the public interface by implementing
158/// [`std::alloc::Allocator`].
159///
160/// # Examples
161///
162/// ## Standalone allocator
163///
164/// ```rust
165/// #![feature(allocator_api)]
166/// #![feature(slice_ptr_get)]
167///
168/// use std::alloc::{Allocator, Layout};
169///
170/// use rulloc::Rulloc;
171///
172/// let rulloc = Rulloc::default();
173/// let (size, align) = (128, 8);
174/// let layout = Layout::from_size_align(size, align).unwrap();
175///
176/// unsafe {
177/// let address = rulloc.allocate(layout).unwrap();
178/// // The allocator can return more space than requested.
179/// assert!(address.len() >= size);
180/// // Alignment is guaranteed for any power of two.
181/// assert_eq!(address.as_mut_ptr() as usize % align, 0);
182/// // Deallocate the pointer.
183/// rulloc.deallocate(address.cast(), layout);
184/// }
185/// ```
186///
187/// ## Collections and [`Box`]
188///
189/// ```no_run
190/// #![feature(allocator_api)]
191///
192/// use std::alloc::Allocator;
193///
194/// use rulloc::Rulloc;
195///
196/// let rulloc = Rulloc::default();
197///
198/// // Any struct that supports the allocator API works with Rulloc.
199/// let mut num = Box::new_in(12, &rulloc);
200/// assert_eq!(*num, 12);
201///
202/// let mut vec = Vec::new_in(&rulloc);
203/// vec.push(5);
204/// assert_eq!(vec[0], 5);
205/// ```
206///
207/// ## Global allocator
208///
209/// ```no_run
210/// #![feature(allocator_api)]
211///
212/// use rulloc::Rulloc;
213///
214/// #[global_allocator]
215/// static ALLOCATOR: Rulloc = Rulloc::with_default_config();
216///
217/// fn main() {
218/// let num = Box::new(5);
219/// assert_eq!(*num, 5);
220/// }
221/// ```
222pub struct Rulloc<const N: usize = 3> {
223 /// Currently we use a global [`Mutex`] to access the allocator, but here
224 /// are some ideas to further optimize multithreaded allocations:
225 ///
226 /// 1. Use one [`Mutex`] per [`Bucket`]. That way different size allocations
227 /// don't have to wait on each other. Note that reallocations might try to
228 /// "move" a pointer from one [`Bucket`] to another if the requested new
229 /// size changes drastically. If each [`Bucket`] has its own lock, we have
230 /// to handle deadlocks properly with [`Mutex::try_lock`].
231 ///
232 /// 2. Use a fixed number of allocators and distribute requests from
233 /// different threads between them (round-robin, for example). Each
234 /// allocator could have a global [`Mutex`] or one [`Mutex`] per [`Bucket`]
235 /// like mentioned above.
236 ///
237 /// 3. Don't use any [`Mutex`] at all, have one entire allocator per thread.
238 /// Conceptually, we would need a mapping of [`std::thread::ThreadId`] to
239 /// [`InternalAllocator`]. Instead of using general data structures that
240 /// need to allocate memory, such as hash maps, we could use a fixed size
241 /// array and store a tuple of `(ThreadId, Bucket)`. Each allocation will
242 /// perform a linear scan to find the [`Bucket`] where we should allocate.
243 /// This is technically O(n) but as long as we don't have thousands of
244 /// threads it won't be an issue. If we end up needing to allocate memory
245 /// for ourselves, we can just use [`crate::platform::request_memory`]. The
246 /// issue with this approach is that we have to deal with threads that
247 /// deallocate memory which was not allocated by themselves, so we need more
248 /// than a simple mapping.
249 allocator: Mutex<InternalAllocator<N>>,
250}
251
252unsafe impl<const N: usize> Sync for Rulloc<N> {}
253
254impl Rulloc {
255 /// Default configuration includes 3 buckets of sizes 128, 1024 and 8192.
256 /// See [`Rulloc::<N>::with_bucket_sizes`] for details.
257 pub const fn with_default_config() -> Self {
258 Self {
259 allocator: Mutex::new(InternalAllocator::with_bucket_sizes([128, 1024, 8192])),
260 }
261 }
262}
263
264impl<const N: usize> Rulloc<N> {
265 /// Builds a new allocator configured with the given bucket sizes.
266 ///
267 /// # Examples
268 ///
269 /// ```rust
270 /// #![feature(allocator_api)]
271 ///
272 /// use std::alloc::{Allocator, Layout};
273 ///
274 /// use rulloc::Rulloc;
275 ///
276 /// // 3 fixed size buckets. First one will contain allocations less than
277 /// // or equal to 64 bytes in size, second one will contain allocations
278 /// // less than or equal to 128 bytes, and so forth. Allocations larger
279 /// // than the last bucket size will be allocated separately.
280 /// let rulloc = Rulloc::<3>::with_bucket_sizes([64, 128, 256]);
281 ///
282 /// // Allocated in the first bucket.
283 /// let p1 = rulloc.allocate(Layout::from_size_align(64, 8).unwrap()).unwrap();
284 /// // Allocated in the second bucket.
285 /// let p2 = rulloc.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap();
286 /// // Allocated in the third bucket.
287 /// let p3 = rulloc.allocate(Layout::from_size_align(210, 8).unwrap()).unwrap();
288 /// // Allocated in a dynamic bucket that can allocate any size.
289 /// let p4 = rulloc.allocate(Layout::from_size_align(512, 8).unwrap()).unwrap();
290 ///
291 /// assert!(p1.len() >= 64);
292 /// assert!(p2.len() >= 100);
293 /// assert!(p3.len() >= 210);
294 /// assert!(p4.len() >= 512);
295 /// ```
296 pub fn with_bucket_sizes(sizes: [usize; N]) -> Self {
297 Self {
298 allocator: Mutex::new(InternalAllocator::with_bucket_sizes(sizes)),
299 }
300 }
301}
302
303impl Default for Rulloc {
304 fn default() -> Self {
305 Rulloc::with_default_config()
306 }
307}
308
309unsafe impl<const N: usize> Allocator for Rulloc<N> {
310 fn allocate(&self, layout: Layout) -> AllocResult {
311 unsafe {
312 match self.allocator.lock() {
313 Ok(mut allocator) => allocator.allocate(layout),
314 Err(_) => Err(AllocError),
315 }
316 }
317 }
318
319 unsafe fn deallocate(&self, address: NonNull<u8>, layout: Layout) {
320 if let Ok(mut allocator) = self.allocator.lock() {
321 allocator.deallocate(address, layout)
322 }
323 }
324
325 unsafe fn shrink(
326 &self,
327 address: NonNull<u8>,
328 old_layout: Layout,
329 new_layout: Layout,
330 ) -> AllocResult {
331 match self.allocator.lock() {
332 Ok(mut allocator) => {
333 allocator.reallocate(&Realloc::shrink(address, old_layout, new_layout))
334 }
335 Err(_) => Err(AllocError),
336 }
337 }
338
339 unsafe fn grow(
340 &self,
341 address: NonNull<u8>,
342 old_layout: Layout,
343 new_layout: Layout,
344 ) -> AllocResult {
345 match self.allocator.lock() {
346 Ok(mut allocator) => {
347 allocator.reallocate(&Realloc::grow(address, old_layout, new_layout))
348 }
349 Err(_) => Err(AllocError),
350 }
351 }
352
353 unsafe fn grow_zeroed(
354 &self,
355 address: NonNull<u8>,
356 old_layout: Layout,
357 new_layout: Layout,
358 ) -> AllocResult {
359 let new_address = self.grow(address, old_layout, new_layout)?;
360 let zero_from = new_address
361 .as_mut_ptr()
362 .map_addr(|addr| addr + old_layout.size());
363 zero_from.write_bytes(0, new_layout.size() - old_layout.size());
364
365 Ok(new_address)
366 }
367}
368
369unsafe impl GlobalAlloc for Rulloc {
370 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
371 match self.allocate(layout) {
372 Ok(address) => address.cast().as_ptr(),
373 Err(_) => ptr::null_mut(),
374 }
375 }
376
377 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
378 self.deallocate(NonNull::new_unchecked(ptr), layout)
379 }
380
381 unsafe fn realloc(&self, address: *mut u8, old_layout: Layout, new_size: usize) -> *mut u8 {
382 let new_layout = Layout::from_size_align(new_size, old_layout.align()).unwrap();
383 let address = NonNull::new_unchecked(address);
384
385 let result = if old_layout.size() <= new_size {
386 self.shrink(address, old_layout, new_layout)
387 } else {
388 self.grow(address, old_layout, new_layout)
389 };
390
391 match result {
392 Ok(new_address) => new_address.as_mut_ptr(),
393 Err(_) => ptr::null_mut(),
394 }
395 }
396}
397
398#[cfg(test)]
399mod tests {
400 use std::{
401 sync,
402 thread::{self, ThreadId},
403 };
404
405 use super::*;
406 use crate::platform::PAGE_SIZE;
407
408 #[test]
409 fn internal_allocator_wrapper() {
410 let allocator = Rulloc::with_default_config();
411 unsafe {
412 let layout1 = Layout::array::<u8>(8).unwrap();
413 let mut addr1 = allocator.allocate(layout1).unwrap();
414
415 addr1.as_mut().fill(69);
416
417 let layout2 = Layout::array::<u8>(PAGE_SIZE * 2).unwrap();
418 let mut addr2 = allocator.allocate(layout2).unwrap();
419
420 addr2.as_mut().fill(42);
421
422 for value in addr1.as_ref() {
423 assert_eq!(value, &69);
424 }
425
426 allocator.deallocate(addr1.cast(), layout1);
427
428 for value in addr2.as_ref() {
429 assert_eq!(value, &42);
430 }
431
432 allocator.deallocate(addr2.cast(), layout2);
433 }
434 }
435
436 #[test]
437 fn buckets() {
438 unsafe {
439 let sizes = [8, 16, 24];
440 let mut allocator = InternalAllocator::<3>::with_bucket_sizes(sizes);
441
442 macro_rules! verify_number_of_regions_per_bucket {
443 ($expected:expr) => {
444 for i in 0..sizes.len() {
445 assert_eq!(allocator.buckets[i].regions().len(), $expected[i]);
446 }
447 };
448 }
449
450 let layout1 = Layout::array::<u8>(sizes[0]).unwrap();
451 let addr1 = allocator.allocate(layout1).unwrap().cast();
452 verify_number_of_regions_per_bucket!([1, 0, 0]);
453
454 let layout2 = Layout::array::<u8>(sizes[1]).unwrap();
455 let addr2 = allocator.allocate(layout2).unwrap().cast();
456 verify_number_of_regions_per_bucket!([1, 1, 0]);
457
458 let layout3 = Layout::array::<u8>(sizes[2]).unwrap();
459 let addr3 = allocator.allocate(layout3).unwrap().cast();
460 verify_number_of_regions_per_bucket!([1, 1, 1]);
461
462 allocator.deallocate(addr1, layout1);
463 verify_number_of_regions_per_bucket!([0, 1, 1]);
464
465 allocator.deallocate(addr2, layout2);
466 verify_number_of_regions_per_bucket!([0, 0, 1]);
467
468 allocator.deallocate(addr3, layout3);
469 verify_number_of_regions_per_bucket!([0, 0, 0]);
470
471 let layout4 = Layout::array::<u8>(sizes[2] + 128).unwrap();
472 let addr4 = allocator.allocate(layout4).unwrap().cast();
473 verify_number_of_regions_per_bucket!([0, 0, 0]);
474 assert_eq!(allocator.dyn_bucket.regions().len(), 1);
475
476 allocator.deallocate(addr4, layout4);
477 assert_eq!(allocator.dyn_bucket.regions().len(), 0);
478
479 // Now let's try some reallocs
480 let mut realloc_addr = allocator.allocate(layout1).unwrap();
481 let corruption_check = 213;
482 realloc_addr.as_mut().fill(corruption_check);
483
484 realloc_addr = allocator
485 .reallocate(&Realloc::grow(realloc_addr.cast(), layout1, layout2))
486 .unwrap();
487 verify_number_of_regions_per_bucket!([0, 1, 0]);
488
489 realloc_addr = allocator
490 .reallocate(&Realloc::grow(realloc_addr.cast(), layout2, layout3))
491 .unwrap();
492 verify_number_of_regions_per_bucket!([0, 0, 1]);
493
494 for value in &realloc_addr.as_ref()[..layout1.size()] {
495 assert_eq!(*value, corruption_check);
496 }
497 }
498 }
499
500 fn verify_buckets_are_empty(allocator: Rulloc) {
501 let internal = allocator.allocator.lock().unwrap();
502 for bucket in &internal.buckets {
503 assert_eq!(bucket.regions().len(), 0);
504 }
505 assert_eq!(internal.dyn_bucket.regions().len(), 0);
506 }
507
508 /// We'll make all the threads do only allocs at the same time, then wait
509 /// and do only deallocs at the same time.
510 #[test]
511 fn multiple_threads_synchronized_allocs_and_deallocs() {
512 let allocator = Rulloc::with_default_config();
513
514 let num_threads = 8;
515
516 let barrier = sync::Barrier::new(num_threads);
517
518 thread::scope(|scope| {
519 for _ in 0..num_threads {
520 scope.spawn(|| unsafe {
521 let num_elements = 1024;
522 let layout = Layout::array::<ThreadId>(num_elements).unwrap();
523 let addr = allocator.allocate(layout).unwrap().cast::<ThreadId>();
524 let id = thread::current().id();
525
526 for i in 0..num_elements {
527 *addr.as_ptr().add(i) = id;
528 }
529
530 barrier.wait();
531
532 // Check memory corruption.
533 for i in 0..num_elements {
534 assert_eq!(*addr.as_ptr().add(i), id);
535 }
536
537 allocator.deallocate(addr.cast(), layout);
538 });
539 }
540 });
541
542 verify_buckets_are_empty(allocator);
543 }
544
545 /// In this case we'll make the threads do allocs and deallocs
546 /// interchangeably.
547 #[test]
548 fn multiple_threads_unsynchronized_allocs_and_deallocs() {
549 let allocator = Rulloc::with_default_config();
550
551 let num_threads = 8;
552
553 let barrier = sync::Barrier::new(num_threads);
554
555 thread::scope(|scope| {
556 for _ in 0..num_threads {
557 scope.spawn(|| unsafe {
558 // We'll use different sizes to make sure that contention
559 // over a single region or multiple regions doesn't cause
560 // issues.
561 let layouts = [16, 256, 1024, 2048, 4096, 8192]
562 .map(|size| Layout::array::<u8>(size).unwrap());
563
564 // Miri is really slow, but we don't need as many operations
565 // to find bugs with it.
566 let num_allocs = if cfg!(miri) { 20 } else { 1000 };
567
568 for layout in layouts {
569 barrier.wait();
570 for _ in 0..num_allocs {
571 let addr = allocator.allocate(layout).unwrap().cast::<u8>();
572 if cfg!(miri) {
573 // Since Miri is slow we won't write all the
574 // bytes, just a few to check data races. If
575 // somehow two threads receive the same address,
576 // Miri will catch that.
577 let offsets = [0, layout.size() / 2, layout.size() - 1];
578 let values = [1, 5, 10];
579 for (offset, value) in offsets.iter().zip(values) {
580 *addr.as_ptr().add(*offset) = value;
581 }
582 for (offset, value) in offsets.iter().zip(values) {
583 assert_eq!(*addr.as_ptr().add(*offset), value);
584 }
585 } else {
586 // If we're not using Miri then write all the
587 // bytes and check them again later.
588 for i in 0..layout.size() {
589 *addr.as_ptr().add(i) = (i % 256) as u8;
590 }
591 for i in 0..layout.size() {
592 assert_eq!(*addr.as_ptr().add(i), (i % 256) as u8);
593 }
594 }
595 allocator.deallocate(addr, layout);
596 }
597 }
598 });
599 }
600 });
601
602 verify_buckets_are_empty(allocator);
603 }
604}