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    pub fn until_now(&self) -> AllocationInfo {
38        let new = TurboMalloc::allocation_counters();
39        AllocationInfo {
40            allocations: new.allocations - self.allocations,
41            deallocations: new.deallocations - self.deallocations,
42            allocation_count: new.allocation_count - self.allocation_count,
43            deallocation_count: new.deallocation_count - self.deallocation_count,
44        }
45    }
46}
47
48/// Turbo's preferred global allocator. This is a new type instead of a type
49/// alias because you can't use type aliases to instantiate unit types (E0423).
50pub struct TurboMalloc;
51
52impl TurboMalloc {
53    pub fn memory_usage() -> usize {
54        get()
55    }
56
57    pub fn thread_stop() {
58        flush();
59    }
60
61    pub fn allocation_counters() -> AllocationCounters {
62        self::counter::allocation_counters()
63    }
64
65    pub fn reset_allocation_counters(start: AllocationCounters) {
66        self::counter::reset_allocation_counters(start);
67    }
68}
69
70/// Get the allocator for this platform that we should wrap with TurboMalloc.
71#[inline]
72fn base_alloc() -> &'static impl GlobalAlloc {
73    #[cfg(all(
74        feature = "custom_allocator",
75        not(any(target_family = "wasm", target_env = "musl"))
76    ))]
77    return &mimalloc::MiMalloc;
78    #[cfg(any(
79        not(feature = "custom_allocator"),
80        any(target_family = "wasm", target_env = "musl")
81    ))]
82    return &std::alloc::System;
83}
84
85unsafe impl GlobalAlloc for TurboMalloc {
86    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
87        let ret = unsafe { base_alloc().alloc(layout) };
88        if !ret.is_null() {
89            add(layout.size());
90        }
91        ret
92    }
93
94    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
95        unsafe { base_alloc().dealloc(ptr, layout) };
96        remove(layout.size());
97    }
98
99    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
100        let ret = unsafe { base_alloc().alloc_zeroed(layout) };
101        if !ret.is_null() {
102            add(layout.size());
103        }
104        ret
105    }
106
107    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
108        let ret = unsafe { base_alloc().realloc(ptr, layout, new_size) };
109        if !ret.is_null() {
110            let old_size = layout.size();
111            update(old_size, new_size);
112        }
113        ret
114    }
115}
turbo_tasks_malloc/lib.rs

turbo_tasks_malloc/
lib.rs
