-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
9 changed files
with
228 additions
and
37 deletions.
There are no files selected for viewing
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| use std::{ | ||
| cell::UnsafeCell, | ||
| ptr::NonNull, | ||
| sync::atomic::{AtomicUsize, Ordering}, | ||
| }; | ||
|
|
||
| pub static ALLOCATED: AtomicUsize = AtomicUsize::new(0); | ||
| const RESERVE: usize = 1000240; // 100 KiB | ||
| const MAX_RESERVE: usize = 2000480; // 200 KiB | ||
|
|
||
| struct ThreadLocalCounter { | ||
| reserved: usize, | ||
| } | ||
|
|
||
| impl ThreadLocalCounter { | ||
| fn add(&mut self, size: usize) { | ||
| if self.reserved >= size { | ||
| self.reserved -= size; | ||
| } else { | ||
| let offset = size - self.reserved + RESERVE; | ||
| self.reserved = RESERVE; | ||
| ALLOCATED.fetch_add(offset, Ordering::Relaxed); | ||
| } | ||
| } | ||
|
|
||
| fn remove(&mut self, size: usize) { | ||
| self.reserved += size; | ||
| if self.reserved > MAX_RESERVE { | ||
| let offset = self.reserved - RESERVE; | ||
| self.reserved = RESERVE; | ||
| ALLOCATED.fetch_sub(offset, Ordering::Relaxed); | ||
| } | ||
| } | ||
|
|
||
| fn unload(&mut self) { | ||
| if self.reserved > 0 { | ||
| ALLOCATED.fetch_sub(self.reserved, Ordering::Relaxed); | ||
| self.reserved = 0; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| thread_local! { | ||
| static LOCAL_COUNTER: UnsafeCell<ThreadLocalCounter> = UnsafeCell::new(ThreadLocalCounter { reserved: 0 }); | ||
| } | ||
|
|
||
| pub fn add(size: usize) { | ||
| LOCAL_COUNTER.with(|local| { | ||
| let ptr = local.get(); | ||
| // SAFETY: This is a thread local. | ||
| let mut local = unsafe { NonNull::new_unchecked(ptr) }; | ||
| unsafe { local.as_mut() }.add(size); | ||
| }) | ||
| } | ||
|
|
||
| pub fn remove(size: usize) { | ||
| LOCAL_COUNTER.with(|local| { | ||
| let ptr = local.get(); | ||
| // SAFETY: This is a thread local. | ||
| let mut local = unsafe { NonNull::new_unchecked(ptr) }; | ||
| unsafe { local.as_mut() }.remove(size); | ||
| }) | ||
| } | ||
|
|
||
| pub fn thread_stop() { | ||
| LOCAL_COUNTER.with(|local| { | ||
| let ptr = local.get(); | ||
| // SAFETY: This is a thread local. | ||
| let mut local = unsafe { NonNull::new_unchecked(ptr) }; | ||
| unsafe { local.as_mut() }.unload(); | ||
| }) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,43 +1,102 @@ | ||
| use std::alloc::{GlobalAlloc, Layout}; | ||
| mod counter; | ||
|
|
||
| use std::{ | ||
| alloc::{GlobalAlloc, Layout}, | ||
| sync::atomic::Ordering, | ||
| }; | ||
|
|
||
| use self::counter::{add, remove, thread_stop, ALLOCATED}; | ||
|
|
||
| /// Turbo's preferred global allocator. This is a new type instead of a type | ||
| /// alias because you can't use type aliases to instantiate unit types (E0423). | ||
| pub struct TurboMalloc; | ||
|
|
||
| #[cfg(not(all(target_os = "linux", target_arch = "aarch64")))] | ||
| impl TurboMalloc { | ||
| pub fn memory_usage() -> usize { | ||
| ALLOCATED.load(Ordering::Relaxed) | ||
| } | ||
|
|
||
| pub fn thread_stop() { | ||
| thread_stop(); | ||
| } | ||
| } | ||
|
|
||
| #[cfg(all( | ||
| feature = "custom_allocator", | ||
| not(all(target_os = "linux", target_arch = "aarch64")) | ||
| ))] | ||
| unsafe impl GlobalAlloc for TurboMalloc { | ||
| unsafe fn alloc(&self, layout: Layout) -> *mut u8 { | ||
| mimalloc::MiMalloc.alloc(layout) | ||
| let ret = mimalloc::MiMalloc.alloc(layout); | ||
| if !ret.is_null() { | ||
| add(layout.size()); | ||
| } | ||
| ret | ||
| } | ||
|
|
||
| unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { | ||
| mimalloc::MiMalloc.dealloc(ptr, layout) | ||
| mimalloc::MiMalloc.dealloc(ptr, layout); | ||
| remove(layout.size()); | ||
| } | ||
|
|
||
| unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { | ||
| mimalloc::MiMalloc.alloc_zeroed(layout) | ||
| let ret = mimalloc::MiMalloc.alloc_zeroed(layout); | ||
| if ret.is_null() { | ||
| add(layout.size()); | ||
| } | ||
| ret | ||
| } | ||
|
|
||
| unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { | ||
| mimalloc::MiMalloc.realloc(ptr, layout, new_size) | ||
| let ret = mimalloc::MiMalloc.realloc(ptr, layout, new_size); | ||
| if !ret.is_null() { | ||
| let old_size = layout.size(); | ||
| if old_size < new_size { | ||
| add(new_size - old_size); | ||
| } else { | ||
| remove(old_size - new_size); | ||
| } | ||
| } | ||
| ret | ||
| } | ||
| } | ||
|
|
||
| #[cfg(all(target_os = "linux", target_arch = "aarch64"))] | ||
| #[cfg(any( | ||
| not(feature = "custom_allocator"), | ||
| all(target_os = "linux", target_arch = "aarch64") | ||
| ))] | ||
| unsafe impl GlobalAlloc for TurboMalloc { | ||
| unsafe fn alloc(&self, layout: Layout) -> *mut u8 { | ||
| std::alloc::System.alloc(layout) | ||
| let ret = std::alloc::System.alloc(layout); | ||
| if !ret.is_null() { | ||
| add(layout.size()); | ||
| } | ||
| ret | ||
| } | ||
|
|
||
| unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { | ||
| std::alloc::System.dealloc(ptr, layout) | ||
| std::alloc::System.dealloc(ptr, layout); | ||
| remove(layout.size()); | ||
| } | ||
|
|
||
| unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { | ||
| std::alloc::System.alloc_zeroed(layout) | ||
| let ret = std::alloc::System.alloc_zeroed(layout); | ||
| if ret.is_null() { | ||
| add(layout.size()); | ||
| } | ||
| ret | ||
| } | ||
|
|
||
| unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { | ||
| std::alloc::System.realloc(ptr, layout, new_size) | ||
| let ret = std::alloc::System.realloc(ptr, layout, new_size); | ||
| if !ret.is_null() { | ||
| let old_size = layout.size(); | ||
| if old_size < new_size { | ||
| add(new_size - old_size); | ||
| } else { | ||
| remove(old_size - new_size); | ||
| } | ||
| } | ||
| ret | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters