tcmalloc_better/
lib.rs

1#![no_std]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3
4//! A drop-in global allocator wrapper around the [TCMalloc](https://github.com/google/tcmalloc) allocator.
5//! TCMalloc is a general-purpose, performance-oriented allocator built by Google.
6//!
7//! ## Usage
8//! ```rust,ignore
9//! use tcmalloc_better::TCMalloc;
10//!
11//! #[global_allocator]
12//! static GLOBAL: TCMalloc = TCMalloc;
13//!
14//! fn main() {
15//!     TCMalloc::process_background_actions_thread();
16//!
17//!     // Rest of main
18//! }
19//! ```
20//!
21//! ## Caveats
22//!
23//! This library is intended for use on Linux (x86_64, aarch64).
24//! For applications requiring support on a broader range of platforms,
25//! consider using [malloc-best-effort](https://crates.io/crates/malloc-best-effort), which automatically selects
26//! the best allocator based on the target platform.
27//!
28//! ## Feature flags
29#![doc = document_features::document_features!()]
30
31#[cfg(feature = "extension")]
32#[cfg_attr(docsrs, doc(cfg(feature = "extension")))]
33mod extension;
34
35use core::alloc::{GlobalAlloc, Layout};
36
37/// A memory allocator that can be registered as the standard library’s default
38/// through the `#[global_allocator]` attribute.
39pub struct TCMalloc;
40
41unsafe impl GlobalAlloc for TCMalloc {
42    #[inline]
43    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
44        let (size, alignment) = (layout.size(), layout.align());
45        let ptr =
46            unsafe { libtcmalloc_sys::BridgeTCMallocInternalNewAlignedNothrow(size, alignment) };
47        ptr as *mut u8
48    }
49
50    #[cfg(not(feature = "realloc"))]
51    #[inline]
52    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
53        let ptr = ptr as *mut core::ffi::c_void;
54        let (size, alignment) = (layout.size(), layout.align());
55        unsafe {
56            libtcmalloc_sys::TCMallocInternalDeleteSizedAligned(ptr, size, alignment);
57        }
58    }
59
60    #[cfg(feature = "realloc")]
61    #[inline]
62    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
63        let ptr = ptr as *mut core::ffi::c_void;
64        let alignment = layout.align();
65        unsafe {
66            libtcmalloc_sys::TCMallocInternalDeleteAligned(ptr, alignment);
67        }
68    }
69
70    #[cfg(feature = "realloc")]
71    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
72        let alignment = layout.align();
73        let ptr = ptr as *mut core::ffi::c_void;
74        let mut old_size_to_free = Default::default();
75        let new_ptr = {
76            let old_size_to_free = &mut old_size_to_free;
77            unsafe {
78                libtcmalloc_sys::BridgePrepareReallocAligned(
79                    ptr,
80                    new_size,
81                    alignment,
82                    old_size_to_free,
83                )
84            }
85        };
86        if !new_ptr.is_null() && new_ptr != ptr {
87            let size_to_copy = layout.size().min(new_size);
88            unsafe { core::ptr::copy_nonoverlapping(ptr, new_ptr, size_to_copy) };
89            unsafe {
90                libtcmalloc_sys::TCMallocInternalDeleteSizedAligned(
91                    ptr,
92                    old_size_to_free,
93                    alignment,
94                )
95            };
96        }
97        new_ptr as *mut u8
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn it_frees_allocated_memory() {
107        unsafe {
108            let layout = Layout::from_size_align(8, 16).unwrap();
109            let alloc = TCMalloc;
110
111            let ptr = alloc.alloc(layout);
112            alloc.dealloc(ptr, layout);
113        }
114    }
115
116    #[test]
117    fn it_frees_allocated_big_memory() {
118        unsafe {
119            let layout = Layout::from_size_align(1 << 20, 32).unwrap();
120            let alloc = TCMalloc;
121
122            let ptr = alloc.alloc(layout);
123            alloc.dealloc(ptr, layout);
124        }
125    }
126
127    #[test]
128    fn it_frees_zero_allocated_memory() {
129        unsafe {
130            let layout = Layout::from_size_align(8, 16).unwrap();
131            let alloc = TCMalloc;
132
133            let ptr = alloc.alloc_zeroed(layout);
134            alloc.dealloc(ptr, layout);
135        }
136    }
137
138    #[test]
139    fn it_frees_zero_allocated_big_memory() {
140        unsafe {
141            let layout = Layout::from_size_align(1 << 20, 32).unwrap();
142            let alloc = TCMalloc;
143
144            let ptr = alloc.alloc_zeroed(layout);
145            alloc.dealloc(ptr, layout);
146        }
147    }
148
149    #[test]
150    fn it_frees_reallocated_memory() {
151        unsafe {
152            let layout = Layout::from_size_align(8, 16).unwrap();
153            let new_size = 16;
154            let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap();
155            let alloc = TCMalloc;
156
157            let ptr = alloc.alloc(layout);
158            let ptr = alloc.realloc(ptr, layout, new_size);
159            alloc.dealloc(ptr, new_layout);
160        }
161    }
162
163    #[test]
164    fn it_frees_reallocated_big_memory() {
165        unsafe {
166            let layout = Layout::from_size_align(1 << 20, 32).unwrap();
167            let new_size = 2 << 20;
168            let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap();
169            let alloc = TCMalloc;
170
171            let ptr = alloc.alloc(layout);
172            let ptr = alloc.realloc(ptr, layout, new_size);
173            alloc.dealloc(ptr, new_layout);
174        }
175    }
176}