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    pub fn until_now(&self) -> AllocationInfo {
81        let new = TurboMalloc::allocation_counters();
82        AllocationInfo {
83            allocations: new.allocations - self.allocations,
84            deallocations: new.deallocations - self.deallocations,
85            allocation_count: new.allocation_count - self.allocation_count,
86            deallocation_count: new.deallocation_count - self.deallocation_count,
87        }
88    }
89}
90
91/// Turbo's preferred global allocator. This is a new type instead of a type
92/// alias because you can't use type aliases to instantiate unit types (E0423).
93pub struct TurboMalloc;
94
95impl TurboMalloc {
96    // Returns the current amount of memory
97    pub fn memory_usage() -> usize {
98        get()
99    }
100
101    pub fn thread_stop() {
102        flush();
103    }
104
105    pub fn allocation_counters() -> AllocationCounters {
106        self::counter::allocation_counters()
107    }
108
109    pub fn reset_allocation_counters(start: AllocationCounters) {
110        self::counter::reset_allocation_counters(start);
111    }
112}
113
114/// Get the allocator for this platform that we should wrap with TurboMalloc.
115#[inline]
116fn base_alloc() -> &'static impl GlobalAlloc {
117    #[cfg(all(
118        feature = "custom_allocator",
119        not(any(target_family = "wasm", target_env = "musl"))
120    ))]
121    return &mimalloc_rspack::MiMalloc;
122    #[cfg(any(
123        not(feature = "custom_allocator"),
124        any(target_family = "wasm", target_env = "musl")
125    ))]
126    return &std::alloc::System;
127}
128
129#[allow(unused_variables)]
130unsafe fn base_alloc_size(ptr: *const u8, layout: Layout) -> usize {
131    #[cfg(all(
132        feature = "custom_allocator",
133        not(any(target_family = "wasm", target_env = "musl"))
134    ))]
135    return unsafe { mimalloc_rspack::MiMalloc.usable_size(ptr) };
136    #[cfg(any(
137        not(feature = "custom_allocator"),
138        any(target_family = "wasm", target_env = "musl")
139    ))]
140    return layout.size();
141}
142
143unsafe impl GlobalAlloc for TurboMalloc {
144    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
145        let ret = unsafe { base_alloc().alloc(layout) };
146        if !ret.is_null() {
147            let size = unsafe { base_alloc_size(ret, layout) };
148            add(size);
149        }
150        ret
151    }
152
153    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
154        let size = unsafe { base_alloc_size(ptr, layout) };
155        unsafe { base_alloc().dealloc(ptr, layout) };
156        remove(size);
157    }
158
159    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
160        let ret = unsafe { base_alloc().alloc_zeroed(layout) };
161        if !ret.is_null() {
162            let size = unsafe { base_alloc_size(ret, layout) };
163            add(size);
164        }
165        ret
166    }
167
168    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
169        let old_size = unsafe { base_alloc_size(ptr, layout) };
170        let ret = unsafe { base_alloc().realloc(ptr, layout, new_size) };
171        if !ret.is_null() {
172            // SAFETY: the caller must ensure that the `new_size` does not overflow.
173            // `layout.align()` comes from a `Layout` and is thus guaranteed to be valid.
174            let new_layout = unsafe { Layout::from_size_align_unchecked(new_size, layout.align()) };
175            let new_size = unsafe { base_alloc_size(ret, new_layout) };
176            update(old_size, new_size);
177        }
178        ret
179    }
180}
turbo_tasks_malloc/lib.rs

turbo_tasks_malloc/
lib.rs
