turbo_tasks_malloc/
lib.rs

1mod counter;
2
3use std::{
4    alloc::{GlobalAlloc, Layout},
5    marker::PhantomData,
6};
7
8use self::counter::{add, flush, get, remove, update};
9
10#[derive(Default, Clone, Debug)]
11pub struct AllocationInfo {
12    pub allocations: usize,
13    pub deallocations: usize,
14    pub allocation_count: usize,
15    pub deallocation_count: usize,
16}
17
18impl AllocationInfo {
19    pub fn is_empty(&self) -> bool {
20        self.allocations == 0
21            && self.deallocations == 0
22            && self.allocation_count == 0
23            && self.deallocation_count == 0
24    }
25}
26
27#[derive(Default, Clone, Debug)]
28pub struct AllocationCounters {
29    pub allocations: usize,
30    pub deallocations: usize,
31    pub allocation_count: usize,
32    pub deallocation_count: usize,
33    _not_send: PhantomData<*mut ()>,
34}
35
36impl AllocationCounters {
37    const fn new() -> Self {
38        Self {
39            allocation_count: 0,
40            deallocation_count: 0,
41            allocations: 0,
42            deallocations: 0,
43            _not_send: PhantomData {},
44        }
45    }
46    pub fn until_now(&self) -> AllocationInfo {
47        let new = TurboMalloc::allocation_counters();
48        AllocationInfo {
49            allocations: new.allocations - self.allocations,
50            deallocations: new.deallocations - self.deallocations,
51            allocation_count: new.allocation_count - self.allocation_count,
52            deallocation_count: new.deallocation_count - self.deallocation_count,
53        }
54    }
55}
56
57/// Turbo's preferred global allocator. This is a new type instead of a type
58/// alias because you can't use type aliases to instantiate unit types (E0423).
59pub struct TurboMalloc;
60
61impl TurboMalloc {
62    // Returns the current amount of memory
63    pub fn memory_usage() -> usize {
64        get()
65    }
66
67    pub fn thread_stop() {
68        flush();
69    }
70
71    pub fn allocation_counters() -> AllocationCounters {
72        self::counter::allocation_counters()
73    }
74
75    pub fn reset_allocation_counters(start: AllocationCounters) {
76        self::counter::reset_allocation_counters(start);
77    }
78}
79
80/// Get the allocator for this platform that we should wrap with TurboMalloc.
81#[inline]
82fn base_alloc() -> &'static impl GlobalAlloc {
83    #[cfg(all(
84        feature = "custom_allocator",
85        not(any(target_family = "wasm", target_env = "musl"))
86    ))]
87    return &mimalloc::MiMalloc;
88    #[cfg(any(
89        not(feature = "custom_allocator"),
90        any(target_family = "wasm", target_env = "musl")
91    ))]
92    return &std::alloc::System;
93}
94
95unsafe impl GlobalAlloc for TurboMalloc {
96    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
97        let ret = unsafe { base_alloc().alloc(layout) };
98        if !ret.is_null() {
99            add(layout.size());
100        }
101        ret
102    }
103
104    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
105        unsafe { base_alloc().dealloc(ptr, layout) };
106        remove(layout.size());
107    }
108
109    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
110        let ret = unsafe { base_alloc().alloc_zeroed(layout) };
111        if !ret.is_null() {
112            add(layout.size());
113        }
114        ret
115    }
116
117    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
118        let ret = unsafe { base_alloc().realloc(ptr, layout, new_size) };
119        if !ret.is_null() {
120            let old_size = layout.size();
121            update(old_size, new_size);
122        }
123        ret
124    }
125}
turbo_tasks_malloc/lib.rs

turbo_tasks_malloc/
lib.rs
