turbo_tasks_malloc/
lib.rs

1mod counter;
2
3use std::{
4    alloc::{GlobalAlloc, Layout},
5    marker::PhantomData,
6    ops::{Add, AddAssign},
7};
8
9use self::counter::{add, flush, get, remove, update};
10
11#[derive(Default, Clone, Debug)]
12pub struct AllocationInfo {
13    pub allocations: usize,
14    pub deallocations: usize,
15    pub allocation_count: usize,
16    pub deallocation_count: usize,
17}
18
19impl AllocationInfo {
20    pub const ZERO: Self = Self {
21        allocations: 0,
22        deallocations: 0,
23        allocation_count: 0,
24        deallocation_count: 0,
25    };
26
27    pub fn is_empty(&self) -> bool {
28        self.allocations == 0
29            && self.deallocations == 0
30            && self.allocation_count == 0
31            && self.deallocation_count == 0
32    }
33
34    pub fn memory_usage(&self) -> usize {
35        self.allocations.saturating_sub(self.deallocations)
36    }
37}
38
39impl Add<Self> for AllocationInfo {
40    type Output = Self;
41
42    fn add(self, other: Self) -> Self {
43        Self {
44            allocations: self.allocations + other.allocations,
45            deallocations: self.deallocations + other.deallocations,
46            allocation_count: self.allocation_count + other.allocation_count,
47            deallocation_count: self.deallocation_count + other.deallocation_count,
48        }
49    }
50}
51
52impl AddAssign<Self> for AllocationInfo {
53    fn add_assign(&mut self, other: Self) {
54        self.allocations += other.allocations;
55        self.deallocations += other.deallocations;
56        self.allocation_count += other.allocation_count;
57        self.deallocation_count += other.deallocation_count;
58    }
59}
60
61#[derive(Default, Clone, Debug)]
62pub struct AllocationCounters {
63    pub allocations: usize,
64    pub deallocations: usize,
65    pub allocation_count: usize,
66    pub deallocation_count: usize,
67    _not_send: PhantomData<*mut ()>,
68}
69
70impl AllocationCounters {
71    const fn new() -> Self {
72        Self {
73            allocation_count: 0,
74            deallocation_count: 0,
75            allocations: 0,
76            deallocations: 0,
77            _not_send: PhantomData {},
78        }
79    }
80}
81
82/// Turbo's preferred global allocator. This is a new type instead of a type
83/// alias because you can't use type aliases to instantiate unit types (E0423).
84pub struct TurboMalloc;
85
86impl TurboMalloc {
87    // Returns the current amount of memory
88    pub fn memory_usage() -> usize {
89        get()
90    }
91
92    pub fn thread_stop() {
93        flush();
94    }
95
96    pub fn allocation_counters() -> AllocationCounters {
97        self::counter::allocation_counters()
98    }
99
100    pub fn reset_allocation_counters(start: AllocationCounters) {
101        self::counter::reset_allocation_counters(start);
102    }
103}
104
105/// Get the allocator for this platform that we should wrap with TurboMalloc.
106#[inline]
107fn base_alloc() -> &'static impl GlobalAlloc {
108    #[cfg(all(feature = "custom_allocator", not(target_family = "wasm")))]
109    return &mimalloc::MiMalloc;
110    #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]
111    return &std::alloc::System;
112}
113
114#[allow(unused_variables)]
115unsafe fn base_alloc_size(ptr: *const u8, layout: Layout) -> usize {
116    #[cfg(all(feature = "custom_allocator", not(target_family = "wasm")))]
117    return unsafe { mimalloc::MiMalloc.usable_size(ptr) };
118    #[cfg(not(all(feature = "custom_allocator", not(target_family = "wasm"))))]
119    return layout.size();
120}
121
122unsafe impl GlobalAlloc for TurboMalloc {
123    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
124        let ret = unsafe { base_alloc().alloc(layout) };
125        if !ret.is_null() {
126            let size = unsafe { base_alloc_size(ret, layout) };
127            add(size);
128        }
129        ret
130    }
131
132    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
133        let size = unsafe { base_alloc_size(ptr, layout) };
134        unsafe { base_alloc().dealloc(ptr, layout) };
135        remove(size);
136    }
137
138    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
139        let ret = unsafe { base_alloc().alloc_zeroed(layout) };
140        if !ret.is_null() {
141            let size = unsafe { base_alloc_size(ret, layout) };
142            add(size);
143        }
144        ret
145    }
146
147    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
148        let old_size = unsafe { base_alloc_size(ptr, layout) };
149        let ret = unsafe { base_alloc().realloc(ptr, layout, new_size) };
150        if !ret.is_null() {
151            // SAFETY: the caller must ensure that the `new_size` does not overflow.
152            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid.
153            let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
154            let new_size = unsafe { base_alloc_size(ret, new_layout) };
155            update(old_size, new_size);
156        }
157        ret
158    }
159}
turbo_tasks_malloc/lib.rs

turbo_tasks_malloc/
lib.rs
