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}