snmalloc_rs/lib.rs
1#![no_std]
2//! `snmalloc-rs` provides a wrapper for [`microsoft/snmalloc`](https://github.com/microsoft/snmalloc) to make it usable as a global allocator for rust.
3//! snmalloc is a research allocator. Its key design features are:
4//! - Memory that is freed by the same thread that allocated it does not require any synchronising operations.
5//! - Freeing memory in a different thread to initially allocated it, does not take any locks and instead uses a novel message passing scheme to return the memory to the original allocator, where it is recycled.
6//! - The allocator uses large ranges of pages to reduce the amount of meta-data required.
7//!
8//! The benchmark is available at the [paper](https://github.com/microsoft/snmalloc/blob/master/snmalloc.pdf) of `snmalloc`
9//! There are three features defined in this crate:
10//! - `debug`: Enable the `Debug` mode in `snmalloc`.
11//! - `1mib`: Use the `1mib` chunk configuration.
12//! - `cache-friendly`: Make the allocator more cache friendly (setting `CACHE_FRIENDLY_OFFSET` to `64` in building the library).
13//!
14//! The whole library supports `no_std`.
15//!
16//! To use `snmalloc-rs` add it as a dependency:
17//! ```toml
18//! # Cargo.toml
19//! [dependencies]
20//! snmalloc-rs = "0.1.0"
21//! ```
22//!
23//! To set `SnMalloc` as the global allocator add this to your project:
24//! ```rust
25//! #[global_allocator]
26//! static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
27//! ```
28extern crate snmalloc_sys as ffi;
29
30use core::{
31 alloc::{GlobalAlloc, Layout},
32 ptr::NonNull,
33};
34
35/// Memory usage statistics from the snmalloc backend.
36///
37/// These are range-level figures (slab/chunk granularity) reflecting bytes
38/// reserved from the OS, not the count of live individual allocations.
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40pub struct AllocStats {
41 /// Bytes currently reserved from the OS.
42 pub current_memory_usage: usize,
43 /// High-water mark of `current_memory_usage`.
44 pub peak_memory_usage: usize,
45}
46
47#[derive(Debug, Copy, Clone)]
48#[repr(C)]
49pub struct SnMalloc;
50
51unsafe impl Send for SnMalloc {}
52unsafe impl Sync for SnMalloc {}
53
54impl SnMalloc {
55 #[inline(always)]
56 pub const fn new() -> Self {
57 Self
58 }
59
60 /// Returns the available bytes in a memory block.
61 #[inline(always)]
62 pub fn usable_size(&self, ptr: *const u8) -> Option<usize> {
63 match ptr.is_null() {
64 true => None,
65 false => Some(unsafe { ffi::sn_rust_usable_size(ptr.cast()) })
66 }
67 }
68
69 /// Returns current and peak OS-level memory reservation statistics.
70 /// See [`AllocStats`] for what the values measure.
71 pub fn memory_stats() -> AllocStats {
72 let mut current = 0usize;
73 let mut peak = 0usize;
74 unsafe { ffi::sn_rust_statistics(&mut current, &mut peak) };
75 AllocStats { current_memory_usage: current, peak_memory_usage: peak }
76 }
77
78 /// Allocates memory with the given layout, returning a non-null pointer on success
79 #[inline(always)]
80 pub fn alloc_aligned(&self, layout: Layout) -> Option<NonNull<u8>> {
81 match layout.size() {
82 0 => NonNull::new(layout.align() as *mut u8),
83 size => NonNull::new(unsafe { ffi::sn_rust_alloc(layout.align(), size) }.cast())
84 }
85 }
86}
87
88unsafe impl GlobalAlloc for SnMalloc {
89 /// Allocate the memory with the given alignment and size.
90 /// On success, it returns a pointer pointing to the required memory address.
91 /// On failure, it returns a null pointer.
92 /// The client must assure the following things:
93 /// - `alignment` is greater than zero
94 /// - Other constrains are the same as the rust standard library.
95 ///
96 /// The program may be forced to abort if the constrains are not full-filled.
97 #[inline(always)]
98 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
99 match layout.size() {
100 0 => layout.align() as *mut u8,
101 size => ffi::sn_rust_alloc(layout.align(), size).cast()
102 }
103 }
104
105 /// De-allocate the memory at the given address with the given alignment and size.
106 /// The client must assure the following things:
107 /// - the memory is acquired using the same allocator and the pointer points to the start position.
108 /// - Other constrains are the same as the rust standard library.
109 ///
110 /// The program may be forced to abort if the constrains are not full-filled.
111 #[inline(always)]
112 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
113 if layout.size() != 0 {
114 ffi::sn_rust_dealloc(ptr as _, layout.align(), layout.size());
115 }
116 }
117
118 /// Behaves like alloc, but also ensures that the contents are set to zero before being returned.
119 #[inline(always)]
120 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
121 match layout.size() {
122 0 => layout.align() as *mut u8,
123 size => ffi::sn_rust_alloc_zeroed(layout.align(), size).cast()
124 }
125 }
126
127 /// Re-allocate the memory at the given address with the given alignment and size.
128 /// On success, it returns a pointer pointing to the required memory address.
129 /// The memory content within the `new_size` will remains the same as previous.
130 /// On failure, it returns a null pointer. In this situation, the previous memory is not returned to the allocator.
131 /// The client must assure the following things:
132 /// - the memory is acquired using the same allocator and the pointer points to the start position
133 /// - `alignment` fulfills all the requirements as `rust_alloc`
134 /// - Other constrains are the same as the rust standard library.
135 ///
136 /// The program may be forced to abort if the constrains are not full-filled.
137 #[inline(always)]
138 unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
139 match new_size {
140 0 => {
141 self.dealloc(ptr, layout);
142 layout.align() as *mut u8
143 }
144 new_size if layout.size() == 0 => {
145 self.alloc(Layout::from_size_align_unchecked(new_size, layout.align()))
146 }
147 _ => ffi::sn_rust_realloc(ptr.cast(), layout.align(), layout.size(), new_size).cast()
148 }
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 #[test]
156 fn allocation_lifecycle() {
157 let alloc = SnMalloc::new();
158 unsafe {
159 let layout = Layout::from_size_align(8, 8).unwrap();
160
161 // Test regular allocation
162 let ptr = alloc.alloc(layout);
163 alloc.dealloc(ptr, layout);
164
165 // Test zeroed allocation
166 let ptr = alloc.alloc_zeroed(layout);
167 alloc.dealloc(ptr, layout);
168
169 // Test reallocation
170 let ptr = alloc.alloc(layout);
171 let ptr = alloc.realloc(ptr, layout, 16);
172 alloc.dealloc(ptr, layout);
173
174 // Test large allocation
175 let large_layout = Layout::from_size_align(1 << 20, 32).unwrap();
176 let ptr = alloc.alloc(large_layout);
177 alloc.dealloc(ptr, large_layout);
178 }
179 }
180 #[test]
181 fn it_frees_allocated_memory() {
182 unsafe {
183 let layout = Layout::from_size_align(8, 8).unwrap();
184 let alloc = SnMalloc;
185
186 let ptr = alloc.alloc(layout);
187 alloc.dealloc(ptr, layout);
188 }
189 }
190
191 #[test]
192 fn it_frees_zero_allocated_memory() {
193 unsafe {
194 let layout = Layout::from_size_align(8, 8).unwrap();
195 let alloc = SnMalloc;
196
197 let ptr = alloc.alloc_zeroed(layout);
198 alloc.dealloc(ptr, layout);
199 }
200 }
201
202 #[test]
203 fn it_frees_reallocated_memory() {
204 unsafe {
205 let layout = Layout::from_size_align(8, 8).unwrap();
206 let alloc = SnMalloc;
207
208 let ptr = alloc.alloc(layout);
209 let ptr = alloc.realloc(ptr, layout, 16);
210 alloc.dealloc(ptr, layout);
211 }
212 }
213
214 #[test]
215 fn it_frees_large_alloc() {
216 unsafe {
217 let layout = Layout::from_size_align(1 << 20, 32).unwrap();
218 let alloc = SnMalloc;
219
220 let ptr = alloc.alloc(layout);
221 alloc.dealloc(ptr, layout);
222 }
223 }
224
225 #[test]
226 fn test_usable_size() {
227 let alloc = SnMalloc::new();
228 unsafe {
229 let layout = Layout::from_size_align(8, 8).unwrap();
230 let ptr = alloc.alloc(layout);
231 let usz = alloc.usable_size(ptr).expect("usable_size returned None");
232 alloc.dealloc(ptr, layout);
233 assert!(usz >= 8);
234 }
235 }
236}