blink_alloc/
sync.rs

1//! This module provides single-threaded blink allocator.
2
3use core::{
4    alloc::Layout,
5    mem::ManuallyDrop,
6    ptr::NonNull,
7    sync::atomic::{AtomicUsize, Ordering},
8};
9
10use allocator_api2::alloc::{AllocError, Allocator};
11
12#[cfg(feature = "alloc")]
13use allocator_api2::alloc::Global;
14
15use crate::{
16    api::BlinkAllocator,
17    arena::{ArenaLocal, ArenaSync},
18};
19
20switch_alloc_default! {
21    /// Multi-threaded blink allocator.
22    ///
23    /// Blink-allocator is arena-based allocator that
24    /// allocates memory in growing chunks and serve allocations from them.
25    /// When chunk is exhausted a new larger chunk is allocated.
26    ///
27    /// Deallocation is no-op. [`BlinkAllocator`] can be reset
28    /// to free all chunks except the last one, that will be reused.
29    ///
30    /// Blink allocator aims to allocate a chunk large enough to
31    /// serve all allocations between resets.
32    ///
33    /// A shared and mutable reference to the [`SyncBlinkAlloc`] implement
34    /// [`Allocator`] trait.
35    /// When "nightly" feature is enabled, [`Allocator`] trait is
36    /// [`core::alloc::Allocator`]. Otherwise it is duplicated trait defined
37    /// in [`allocator-api2`](allocator_api2).
38    ///
39    /// Resetting blink allocator requires mutable borrow, so it is not possible
40    /// to do while shared borrow is alive. That matches requirement of
41    /// [`Allocator`] trait - while [`Allocator`] instance
42    /// (a shared reference to `BlinkAlloc`) or any of its clones are alive,
43    /// allocated memory must be valid.
44    ///
45    /// This version of blink-allocator is multi-threaded.
46    /// It can be used from multiple threads concurrently to allocate memory.
47    /// As mutable borrow is required to reset the allocator,
48    /// it is not possible to do when shared.
49    /// Internally it uses [`RwLock`] and [`AtomicUsize`] for synchronized
50    /// interior mutability. [`RwLock`] is only write-locked when new chunk
51    /// must be allocated. The arena allocation is performed using lock-free
52    /// algorithm.
53    ///
54    /// Still it is slower than single-threaded version [`BlinkAlloc`].
55    ///
56    /// For best of both worlds [`LocalBlinkAlloc`] can be created from
57    /// this allocator. [`LocalBlinkAlloc`] will allocate chunks from this
58    /// allocator, but is single-threaded by itself.
59    ///
60    /// [`RwLock`]: parking_lot::RwLock
61    /// [`AtomicUsize`]: core::sync::atomic::AtomicUsize
62    /// [`BlinkAlloc`]: crate::local::BlinkAlloc
63    /// [`LocalBlinkAlloc`]: crate::sync::LocalBlinkAlloc
64    ///
65    /// # Example
66    ///
67    /// ```
68    /// # #![cfg_attr(feature = "nightly", feature(allocator_api))]
69    /// # use blink_alloc::SyncBlinkAlloc;
70    /// # use std::ptr::NonNull;
71    ///
72    /// let mut blink = SyncBlinkAlloc::new();
73    /// let layout = std::alloc::Layout::new::<[u32; 8]>();
74    /// let ptr = blink.allocate(layout).unwrap();
75    /// let ptr = NonNull::new(ptr.as_ptr() as *mut u8).unwrap(); // Method for this is unstable.
76    ///
77    /// unsafe {
78    ///     std::ptr::write(ptr.as_ptr().cast(), [1, 2, 3, 4, 5, 6, 7, 8]);
79    /// }
80    ///
81    /// blink.reset();
82    /// ```
83    ///
84    /// # Example that uses nightly's `allocator_api`
85    ///
86    /// ```
87    /// # #![cfg_attr(feature = "nightly", feature(allocator_api))]
88    /// # use blink_alloc::SyncBlinkAlloc;
89    /// # use allocator_api2::vec::Vec;
90    /// # fn main() {
91    /// let mut blink = SyncBlinkAlloc::new();
92    /// let mut vec = Vec::new_in(&blink);
93    /// vec.push(1);
94    /// vec.extend(1..3);
95    /// vec.extend(3..10);
96    /// drop(vec);
97    /// blink.reset();
98    /// # }
99    /// ```
100    pub struct SyncBlinkAlloc<A: Allocator = +Global> {
101        arena: ArenaSync,
102        allocator: A,
103        max_local_alloc: AtomicUsize,
104    }
105}
106
107impl<A: Allocator> Drop for SyncBlinkAlloc<A> {
108    fn drop(&mut self) {
109        unsafe {
110            self.arena.reset(false, &self.allocator);
111        }
112    }
113}
114
115#[test]
116fn check_sync() {
117    fn for_sync_alloc<A: Allocator + Sync>() {
118        fn is_sink<T: Sync>() {}
119        is_sink::<SyncBlinkAlloc<A>>();
120    }
121    for_sync_alloc::<Global>();
122}
123
124impl<A> Default for SyncBlinkAlloc<A>
125where
126    A: Allocator + Default,
127{
128    #[inline(always)]
129    fn default() -> Self {
130        Self::new_in(Default::default())
131    }
132}
133
134#[cfg(feature = "alloc")]
135impl SyncBlinkAlloc<Global> {
136    /// Creates new blink allocator that uses global allocator
137    /// to allocate memory chunks.
138    ///
139    /// See [`SyncBlinkAlloc::new_in`] for using custom allocator.
140    #[inline(always)]
141    pub const fn new() -> Self {
142        SyncBlinkAlloc::new_in(Global)
143    }
144}
145
146impl<A> SyncBlinkAlloc<A>
147where
148    A: Allocator,
149{
150    /// Creates new blink allocator that uses provided allocator
151    /// to allocate memory chunks.
152    ///
153    /// See [`SyncBlinkAlloc::new`] for using global allocator.
154    #[inline(always)]
155    pub const fn new_in(allocator: A) -> Self {
156        SyncBlinkAlloc {
157            arena: ArenaSync::new(),
158            allocator,
159            max_local_alloc: AtomicUsize::new(0),
160        }
161    }
162
163    /// Returns reference to the underlying allocator used by this blink allocator.
164    #[inline(always)]
165    pub const fn inner(&self) -> &A {
166        &self.allocator
167    }
168
169    /// Creates new blink allocator that uses global allocator
170    /// to allocate memory chunks.
171    /// With this method you can specify initial chunk size.
172    ///
173    /// See [`SyncBlinkAlloc::new_in`] for using custom allocator.
174    #[inline(always)]
175    pub const fn with_chunk_size_in(chunk_size: usize, allocator: A) -> Self {
176        SyncBlinkAlloc {
177            arena: ArenaSync::with_chunk_size(chunk_size),
178            allocator,
179            max_local_alloc: AtomicUsize::new(0),
180        }
181    }
182
183    /// Creates a new thread-local blink allocator proxy
184    /// that borrows from this multi-threaded allocator.
185    ///
186    /// The local proxy allocator works faster and
187    /// allows more consistent memory reuse.
188    /// It can be recreated without resetting the multi-threaded allocator,
189    /// allowing [`SyncBlinkAlloc`] to be warm-up and serve all allocations
190    /// from a single chunk without ever blocking.
191    ///
192    /// Best works for fork-join style of parallelism.
193    /// Create a local allocator for each thread/task.
194    /// Reset after all threads/tasks are finished.
195    ///
196    /// # Examples
197    ///
198    /// ```
199    /// # #![cfg_attr(feature = "nightly", feature(allocator_api))]
200    /// # use blink_alloc::SyncBlinkAlloc;
201    /// # use allocator_api2::vec::Vec;
202    /// # #[cfg(feature = "alloc")] fn main() {
203    /// let mut blink = SyncBlinkAlloc::new();
204    /// for _ in 0..3 {
205    ///     for i in 0..16 {
206    ///         std::thread::scope(|_| {
207    ///             let blink = blink.local();
208    ///             let mut vec = Vec::new_in(&blink);
209    ///             vec.push(i);
210    ///             for j in i*2..i*30 {
211    ///                 vec.push(j); // Proxy will allocate enough memory to grow vec without reallocating on 2nd iteration and later.
212    ///             }
213    ///         });
214    ///     }
215    ///     blink.reset();
216    /// }
217    /// # }
218    /// # #[cfg(not(feature = "alloc"))] fn main() {}
219    /// ```
220    #[inline(always)]
221    pub fn local(&self) -> LocalBlinkAlloc<A> {
222        LocalBlinkAlloc {
223            arena: ArenaLocal::with_chunk_size(self.max_local_alloc.load(Ordering::Relaxed)),
224            shared: self,
225        }
226    }
227
228    /// Allocates memory with specified layout from this allocator.
229    /// If needed it will allocate new chunk using underlying allocator.
230    /// If chunk allocation fails, it will return `Err`.
231    #[inline(always)]
232    pub fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
233        // Safety:
234        // Same instance is used for all allocations and resets.
235        if let Some(ptr) = unsafe { self.arena.alloc_fast(layout) } {
236            return Ok(ptr);
237        }
238        unsafe { self.arena.alloc_slow(layout, &self.allocator) }
239    }
240
241    /// Resizes memory allocation.
242    /// Potentially happens in-place.
243    ///
244    /// # Safety
245    ///
246    /// `ptr` must be a pointer previously returned by [`allocate`](SyncBlinkAlloc::allocate).
247    /// `old_size` must be in range `layout.size()..=slice.len()`
248    /// where `layout` is the layout used in the call to [`allocate`](SyncBlinkAlloc::allocate).
249    /// and `slice` is the slice pointer returned by [`allocate`](SyncBlinkAlloc::allocate).
250    ///
251    /// On success, the old pointer is invalidated and the new pointer is returned.
252    /// On error old allocation is still valid.
253    #[inline(always)]
254    pub unsafe fn resize(
255        &self,
256        ptr: NonNull<u8>,
257        old_layout: Layout,
258        new_layout: Layout,
259    ) -> Result<NonNull<[u8]>, AllocError> {
260        if let Some(ptr) = unsafe { self.arena.resize_fast(ptr, old_layout, new_layout) } {
261            return Ok(ptr);
262        }
263
264        // Safety:
265        // Same instance is used for all allocations and resets.
266        // `ptr` was allocated by this allocator.
267        unsafe {
268            self.arena
269                .resize_slow(ptr, old_layout, new_layout, &self.allocator)
270        }
271    }
272
273    /// Deallocates memory previously allocated from this allocator.
274    ///
275    /// This call may not actually free memory.
276    /// All memory is guaranteed to be freed on [`reset`](SyncBlinkAlloc::reset) call.
277    ///
278    /// # Safety
279    ///
280    /// `ptr` must be a pointer previously returned by [`allocate`](SyncBlinkAlloc::allocate).
281    /// `size` must be in range `layout.size()..=slice.len()`
282    /// where `layout` is the layout used in the call to [`allocate`](SyncBlinkAlloc::allocate).
283    /// and `slice` is the slice pointer returned by [`allocate`](SyncBlinkAlloc::allocate).
284    #[inline(always)]
285    pub unsafe fn deallocate(&self, ptr: NonNull<u8>, size: usize) {
286        // Safety:
287        // `ptr` was allocated by this allocator.
288        unsafe {
289            self.arena.dealloc(ptr, size);
290        }
291    }
292
293    /// Resets this allocator, deallocating all chunks except the last one.
294    /// Last chunk will be reused.
295    /// With steady memory usage after few iterations
296    /// one chunk should be sufficient for all allocations between resets.
297    #[inline(always)]
298    pub fn reset(&mut self) {
299        // Safety:
300        // Same instance is used for all allocations and resets.
301        unsafe {
302            self.arena.reset(true, &self.allocator);
303        }
304    }
305
306    /// Resets this allocator, deallocating all chunks.
307    #[inline(always)]
308    pub fn reset_final(&mut self) {
309        // Safety:
310        // Same instance is used for all allocations and resets.
311        unsafe {
312            self.arena.reset(false, &self.allocator);
313        }
314    }
315
316    /// Resets this allocator, deallocating all chunks except the last one.
317    /// Last chunk will be reused.
318    /// With steady memory usage after few iterations
319    /// one chunk should be sufficient for all allocations between resets.
320    ///
321    /// # Safety
322    ///
323    /// Blink-allocators guarantee that memory can be used while shared
324    /// borrow to the allocator is held, preventing safe `fn reset` call.
325    ///
326    /// With this method it becomes caller responsibility to ensure
327    /// that allocated memory won't be used after reset.
328    #[inline(always)]
329    pub unsafe fn reset_unchecked(&self) {
330        // Safety:
331        // Same instance is used for all allocations and resets.
332        unsafe {
333            self.arena.reset_unchecked(true, &self.allocator);
334        }
335    }
336
337    /// Unwrap this allocator, returning the underlying allocator.
338    /// Leaks allocated chunks.
339    ///
340    /// To deallocate all chunks call [`reset_final`](BlinkAlloc::reset_final) first.
341    ///
342    /// The second returned value will use global allocator, so
343    /// use with care if this method is used inside global allocator.
344    pub fn into_inner(self) -> A {
345        let me = ManuallyDrop::new(self);
346        unsafe { core::ptr::read(&me.allocator) }
347    }
348
349    /// Update maximum local allocation size.
350    /// Can be used by thread-local blink-allocators that use
351    /// this shared blink-allocator.
352    #[inline(always)]
353    pub fn update_max_local_alloc(&self, max_local_alloc: usize) {
354        self.max_local_alloc
355            .fetch_max(max_local_alloc, Ordering::Relaxed);
356    }
357}
358
359unsafe impl<A> Allocator for SyncBlinkAlloc<A>
360where
361    A: Allocator,
362{
363    #[inline(always)]
364    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
365        SyncBlinkAlloc::allocate(self, layout)
366    }
367
368    #[inline(always)]
369    unsafe fn shrink(
370        &self,
371        ptr: NonNull<u8>,
372        old_layout: Layout,
373        new_layout: Layout,
374    ) -> Result<NonNull<[u8]>, AllocError> {
375        SyncBlinkAlloc::resize(self, ptr, old_layout, new_layout)
376    }
377
378    #[inline(always)]
379    unsafe fn grow(
380        &self,
381        ptr: NonNull<u8>,
382        old_layout: Layout,
383        new_layout: Layout,
384    ) -> Result<NonNull<[u8]>, AllocError> {
385        SyncBlinkAlloc::resize(self, ptr, old_layout, new_layout)
386    }
387
388    #[inline(always)]
389    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
390        SyncBlinkAlloc::deallocate(self, ptr, layout.size());
391    }
392}
393
394unsafe impl<A> Allocator for &mut SyncBlinkAlloc<A>
395where
396    A: Allocator,
397{
398    #[inline(always)]
399    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
400        SyncBlinkAlloc::allocate(self, layout)
401    }
402
403    #[inline(always)]
404    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
405        SyncBlinkAlloc::allocate_zeroed(self, layout)
406    }
407
408    #[inline(always)]
409    unsafe fn shrink(
410        &self,
411        ptr: NonNull<u8>,
412        old_layout: Layout,
413        new_layout: Layout,
414    ) -> Result<NonNull<[u8]>, AllocError> {
415        SyncBlinkAlloc::resize(self, ptr, old_layout, new_layout)
416    }
417
418    #[inline(always)]
419    unsafe fn grow(
420        &self,
421        ptr: NonNull<u8>,
422        old_layout: Layout,
423        new_layout: Layout,
424    ) -> Result<NonNull<[u8]>, AllocError> {
425        SyncBlinkAlloc::resize(self, ptr, old_layout, new_layout)
426    }
427
428    #[inline(always)]
429    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
430        SyncBlinkAlloc::deallocate(self, ptr, layout.size());
431    }
432}
433
434unsafe impl<A> BlinkAllocator for SyncBlinkAlloc<A>
435where
436    A: Allocator,
437{
438    #[inline(always)]
439    fn reset(&mut self) {
440        SyncBlinkAlloc::reset(self)
441    }
442}
443
444switch_alloc_default! {
445    /// Thread-local proxy for [`SyncBlinkAlloc`].
446    ///
447    /// Using proxy can yield better performance when
448    /// it is possible to create proxy once to use for many allocations.
449    ///
450    /// See [`SyncBlinkAlloc::local`] for more details.
451    pub struct LocalBlinkAlloc<'a, A: Allocator = +Global> {
452        arena: ArenaLocal,
453        shared: &'a SyncBlinkAlloc<A>,
454    }
455}
456
457impl<A> Drop for LocalBlinkAlloc<'_, A>
458where
459    A: Allocator,
460{
461    fn drop(&mut self) {
462        self.shared
463            .update_max_local_alloc(self.arena.last_chunk_size());
464        self.arena.reset_leak(false);
465    }
466}
467
468impl<A> LocalBlinkAlloc<'_, A>
469where
470    A: Allocator,
471{
472    /// Allocates memory with specified layout from this allocator.
473    /// If needed it will allocate new chunk using underlying allocator.
474    /// If chunk allocation fails, it will return `Err`.
475    #[inline(always)]
476    pub fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
477        // Safety:
478        // Same instance is used for all allocations and resets.
479        if let Some(ptr) = unsafe { self.arena.alloc_fast(layout) } {
480            return Ok(ptr);
481        }
482        unsafe { self.arena.alloc_slow(layout, &self.shared) }
483    }
484
485    /// Resizes memory allocation.
486    /// Potentially happens in-place.
487    ///
488    /// # Safety
489    ///
490    /// `ptr` must be a pointer previously returned by [`allocate`](LocalBlinkAlloc::allocate).
491    /// `old_size` must be in range `layout.size()..=slice.len()`
492    /// where `layout` is the layout used in the call to [`allocate`](LocalBlinkAlloc::allocate).
493    /// and `slice` is the slice pointer returned by [`allocate`](LocalBlinkAlloc::allocate).
494    ///
495    /// On success, the old pointer is invalidated and the new pointer is returned.
496    /// On error old allocation is still valid.
497    #[inline(always)]
498    pub unsafe fn resize(
499        &self,
500        ptr: NonNull<u8>,
501        old_layout: Layout,
502        new_layout: Layout,
503    ) -> Result<NonNull<[u8]>, AllocError> {
504        if let Some(ptr) = unsafe { self.arena.resize_fast(ptr, old_layout, new_layout) } {
505            return Ok(ptr);
506        }
507
508        // Safety:
509        // Same instance is used for all allocations and resets.
510        // `ptr` was allocated by this allocator.
511        unsafe {
512            self.arena
513                .resize_slow(ptr, old_layout, new_layout, &self.shared)
514        }
515    }
516
517    /// Deallocates memory previously allocated from this allocator.
518    ///
519    /// This call may not actually free memory.
520    /// All memory is guaranteed to be freed on [`reset`](LocalBlinkAlloc::reset) call.
521    ///
522    /// # Safety
523    ///
524    /// `ptr` must be a pointer previously returned by [`allocate`](LocalBlinkAlloc::allocate).
525    /// `size` must be in range `layout.size()..=slice.len()`
526    /// where `layout` is the layout used in the call to [`allocate`](LocalBlinkAlloc::allocate).
527    /// and `slice` is the slice pointer returned by [`allocate`](LocalBlinkAlloc::allocate).
528    #[inline(always)]
529    pub unsafe fn deallocate(&self, ptr: NonNull<u8>, size: usize) {
530        // Safety:
531        // `ptr` was allocated by this allocator.
532        unsafe {
533            self.arena.dealloc(ptr, size);
534        }
535    }
536
537    /// Resets this allocator, deallocating all chunks except the last one.
538    /// Last chunk will be reused.
539    /// With steady memory usage after few iterations
540    /// one chunk should be sufficient for all allocations between resets.
541    #[inline(always)]
542    pub fn reset(&mut self) {
543        self.shared
544            .update_max_local_alloc(self.arena.last_chunk_size());
545        self.arena.reset_leak(true);
546    }
547
548    /// Resets this allocator, deallocating all chunks except the last one.
549    /// Last chunk will be reused.
550    /// With steady memory usage after few iterations
551    /// one chunk should be sufficient for all allocations between resets.
552    ///
553    /// # Safety
554    ///
555    /// Blink-allocators guarantee that memory can be used while shared
556    /// borrow to the allocator is held, preventing safe `fn reset` call.
557    ///
558    /// With this method it becomes caller responsibility to ensure
559    /// that allocated memory won't be used after reset.
560    #[inline(always)]
561    pub unsafe fn reset_unchecked(&self) {
562        // Safety:
563        // Same instance is used for all allocations and resets.
564        unsafe {
565            self.arena.reset_unchecked(true, self.shared);
566        }
567    }
568}
569
570unsafe impl<A> Allocator for LocalBlinkAlloc<'_, A>
571where
572    A: Allocator,
573{
574    #[inline(always)]
575    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
576        LocalBlinkAlloc::allocate(self, layout)
577    }
578
579    #[inline(always)]
580    unsafe fn shrink(
581        &self,
582        ptr: NonNull<u8>,
583        old_layout: Layout,
584        new_layout: Layout,
585    ) -> Result<NonNull<[u8]>, AllocError> {
586        LocalBlinkAlloc::resize(self, ptr, old_layout, new_layout)
587    }
588
589    #[inline(always)]
590    unsafe fn grow(
591        &self,
592        ptr: NonNull<u8>,
593        old_layout: Layout,
594        new_layout: Layout,
595    ) -> Result<NonNull<[u8]>, AllocError> {
596        LocalBlinkAlloc::resize(self, ptr, old_layout, new_layout)
597    }
598
599    #[inline(always)]
600    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
601        LocalBlinkAlloc::deallocate(self, ptr, layout.size())
602    }
603}
604
605unsafe impl<A> Allocator for &mut LocalBlinkAlloc<'_, A>
606where
607    A: Allocator,
608{
609    #[inline(always)]
610    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
611        LocalBlinkAlloc::allocate(self, layout)
612    }
613
614    #[inline(always)]
615    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
616        LocalBlinkAlloc::allocate_zeroed(self, layout)
617    }
618
619    #[inline(always)]
620    unsafe fn shrink(
621        &self,
622        ptr: NonNull<u8>,
623        old_layout: Layout,
624        new_layout: Layout,
625    ) -> Result<NonNull<[u8]>, AllocError> {
626        LocalBlinkAlloc::resize(self, ptr, old_layout, new_layout)
627    }
628
629    #[inline(always)]
630    unsafe fn grow(
631        &self,
632        ptr: NonNull<u8>,
633        old_layout: Layout,
634        new_layout: Layout,
635    ) -> Result<NonNull<[u8]>, AllocError> {
636        LocalBlinkAlloc::resize(self, ptr, old_layout, new_layout)
637    }
638
639    #[inline(always)]
640    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
641        LocalBlinkAlloc::deallocate(self, ptr, layout.size())
642    }
643}
644
645unsafe impl<A> BlinkAllocator for LocalBlinkAlloc<'_, A>
646where
647    A: Allocator,
648{
649    #[inline(always)]
650    fn reset(&mut self) {
651        LocalBlinkAlloc::reset(self)
652    }
653}