tikv_jemallocator/
lib.rs

1// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
2// file at the top-level directory of this distribution and at
3// http://rust-lang.org/COPYRIGHT.
4//
5// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8// option. This file may not be copied, modified, or distributed
9// except according to those terms.
10
11//! Bindings for jemalloc as an allocator
12//!
13//! This crate provides bindings to jemalloc as a memory allocator for Rust.
14//! This crate mainly exports, one type, `Jemalloc`, which implements the
15//! `GlobalAlloc` trait and optionally the `Alloc` trait,
16//! and is suitable both as a memory allocator and as a global allocator.
17
18#![cfg_attr(feature = "alloc_trait", feature(allocator_api))]
19// TODO: rename the following lint on next minor bump
20#![allow(renamed_and_removed_lints)]
21#![deny(missing_docs, broken_intra_doc_links)]
22#![no_std]
23
24#[cfg(feature = "alloc_trait")]
25use core::alloc::{Alloc, AllocErr, CannotReallocInPlace, Excess};
26use core::alloc::{GlobalAlloc, Layout};
27#[cfg(feature = "alloc_trait")]
28use core::ptr::NonNull;
29
30use libc::{c_int, c_void};
31
32// This constant equals _Alignof(max_align_t) and is platform-specific. It
33// contains the _maximum_ alignment that the memory allocations returned by the
34// C standard library memory allocation APIs (e.g. `malloc`) are guaranteed to
35// have.
36//
37// The memory allocation APIs are required to return memory that can fit any
38// object whose fundamental aligment is <= _Alignof(max_align_t).
39//
40// In C, there are no ZSTs, and the size of all types is a multiple of their
41// alignment (size >= align). So for allocations with size <=
42// _Alignof(max_align_t), the malloc-APIs return memory whose alignment is
43// either the requested size if its a power-of-two, or the next smaller
44// power-of-two.
45#[cfg(any(target_arch = "arm", target_arch = "mips", target_arch = "powerpc"))]
46const ALIGNOF_MAX_ALIGN_T: usize = 8;
47#[cfg(any(
48    target_arch = "x86",
49    target_arch = "x86_64",
50    target_arch = "aarch64",
51    target_arch = "powerpc64",
52    target_arch = "loongarch64",
53    target_arch = "mips64",
54    target_arch = "riscv64",
55    target_arch = "s390x",
56    target_arch = "sparc64"
57))]
58const ALIGNOF_MAX_ALIGN_T: usize = 16;
59
60/// If `align` is less than `_Alignof(max_align_t)`, and if the requested
61/// allocation `size` is larger than the alignment, we are guaranteed to get a
62/// suitably aligned allocation by default, without passing extra flags, and
63/// this function returns `0`.
64///
65/// Otherwise, it returns the alignment flag to pass to the jemalloc APIs.
66fn layout_to_flags(align: usize, size: usize) -> c_int {
67    if align <= ALIGNOF_MAX_ALIGN_T && align <= size {
68        0
69    } else {
70        ffi::MALLOCX_ALIGN(align)
71    }
72}
73
74// Assumes a condition that always must hold.
75macro_rules! assume {
76    ($e:expr) => {
77        debug_assert!($e);
78        if !($e) {
79            core::hint::unreachable_unchecked();
80        }
81    };
82}
83
84/// Handle to the jemalloc allocator
85///
86/// This type implements the `GlobalAllocAlloc` trait, allowing usage a global allocator.
87///
88/// When the `alloc_trait` feature of this crate is enabled, it also implements the `Alloc` trait,
89/// allowing usage in collections.
90#[derive(Copy, Clone, Default, Debug)]
91pub struct Jemalloc;
92
93unsafe impl GlobalAlloc for Jemalloc {
94    #[inline]
95    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
96        assume!(layout.size() != 0);
97        let flags = layout_to_flags(layout.align(), layout.size());
98        let ptr = if flags == 0 {
99            ffi::malloc(layout.size())
100        } else {
101            ffi::mallocx(layout.size(), flags)
102        };
103        ptr as *mut u8
104    }
105
106    #[inline]
107    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
108        assume!(layout.size() != 0);
109        let flags = layout_to_flags(layout.align(), layout.size());
110        let ptr = if flags == 0 {
111            ffi::calloc(1, layout.size())
112        } else {
113            ffi::mallocx(layout.size(), flags | ffi::MALLOCX_ZERO)
114        };
115        ptr as *mut u8
116    }
117
118    #[inline]
119    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
120        assume!(!ptr.is_null());
121        assume!(layout.size() != 0);
122        let flags = layout_to_flags(layout.align(), layout.size());
123        ffi::sdallocx(ptr as *mut c_void, layout.size(), flags)
124    }
125
126    #[inline]
127    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
128        assume!(layout.size() != 0);
129        assume!(new_size != 0);
130        let flags = layout_to_flags(layout.align(), new_size);
131        let ptr = if flags == 0 {
132            ffi::realloc(ptr as *mut c_void, new_size)
133        } else {
134            ffi::rallocx(ptr as *mut c_void, new_size, flags)
135        };
136        ptr as *mut u8
137    }
138}
139
140#[cfg(feature = "alloc_trait")]
141unsafe impl Alloc for Jemalloc {
142    #[inline]
143    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
144        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
145    }
146
147    #[inline]
148    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
149        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
150    }
151
152    #[inline]
153    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
154        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
155    }
156
157    #[inline]
158    unsafe fn realloc(
159        &mut self,
160        ptr: NonNull<u8>,
161        layout: Layout,
162        new_size: usize,
163    ) -> Result<NonNull<u8>, AllocErr> {
164        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
165    }
166
167    #[inline]
168    unsafe fn alloc_excess(&mut self, layout: Layout) -> Result<Excess, AllocErr> {
169        let flags = layout_to_flags(layout.align(), layout.size());
170        let ptr = ffi::mallocx(layout.size(), flags);
171        if let Some(nonnull) = NonNull::new(ptr as *mut u8) {
172            let excess = ffi::nallocx(layout.size(), flags);
173            Ok(Excess(nonnull, excess))
174        } else {
175            Err(AllocErr)
176        }
177    }
178
179    #[inline]
180    unsafe fn realloc_excess(
181        &mut self,
182        ptr: NonNull<u8>,
183        layout: Layout,
184        new_size: usize,
185    ) -> Result<Excess, AllocErr> {
186        let flags = layout_to_flags(layout.align(), new_size);
187        let ptr = ffi::rallocx(ptr.cast().as_ptr(), new_size, flags);
188        if let Some(nonnull) = NonNull::new(ptr as *mut u8) {
189            let excess = ffi::nallocx(new_size, flags);
190            Ok(Excess(nonnull, excess))
191        } else {
192            Err(AllocErr)
193        }
194    }
195
196    #[inline]
197    fn usable_size(&self, layout: &Layout) -> (usize, usize) {
198        let flags = layout_to_flags(layout.align(), layout.size());
199        unsafe {
200            let max = ffi::nallocx(layout.size(), flags);
201            (layout.size(), max)
202        }
203    }
204
205    #[inline]
206    unsafe fn grow_in_place(
207        &mut self,
208        ptr: NonNull<u8>,
209        layout: Layout,
210        new_size: usize,
211    ) -> Result<(), CannotReallocInPlace> {
212        let flags = layout_to_flags(layout.align(), new_size);
213        let usable_size = ffi::xallocx(ptr.cast().as_ptr(), new_size, 0, flags);
214        if usable_size >= new_size {
215            Ok(())
216        } else {
217            // `xallocx` returns a size smaller than the requested one to
218            // indicate that the allocation could not be grown in place
219            //
220            // the old allocation remains unaltered
221            Err(CannotReallocInPlace)
222        }
223    }
224
225    #[inline]
226    unsafe fn shrink_in_place(
227        &mut self,
228        ptr: NonNull<u8>,
229        layout: Layout,
230        new_size: usize,
231    ) -> Result<(), CannotReallocInPlace> {
232        if new_size == layout.size() {
233            return Ok(());
234        }
235        let flags = layout_to_flags(layout.align(), new_size);
236        let usable_size = ffi::xallocx(ptr.cast().as_ptr(), new_size, 0, flags);
237
238        if usable_size < layout.size() {
239            // If `usable_size` is smaller than the original size, the
240            // size-class of the allocation was shrunk to the size-class of
241            // `new_size`, and it is safe to deallocate the allocation with
242            // `new_size`:
243            Ok(())
244        } else if usable_size == ffi::nallocx(new_size, flags) {
245            // If the allocation was not shrunk and the size class of `new_size`
246            // is the same as the size-class of `layout.size()`, then the
247            // allocation can be properly deallocated using `new_size` (and also
248            // using `layout.size()` because the allocation did not change)
249
250            // note: when the allocation is not shrunk, `xallocx` returns the
251            // usable size of the original allocation, which in this case matches
252            // that of the requested allocation:
253            debug_assert_eq!(
254                ffi::nallocx(new_size, flags),
255                ffi::nallocx(layout.size(), flags)
256            );
257            Ok(())
258        } else {
259            // If the allocation was not shrunk, but the size-class of
260            // `new_size` is not the same as that of the original allocation,
261            // then shrinking the allocation failed:
262            Err(CannotReallocInPlace)
263        }
264    }
265}
266
267/// Return the usable size of the allocation pointed to by ptr.
268///
269/// The return value may be larger than the size that was requested during allocation.
270/// This function is not a mechanism for in-place `realloc()`;
271/// rather it is provided solely as a tool for introspection purposes.
272/// Any discrepancy between the requested allocation size
273/// and the size reported by this function should not be depended on,
274/// since such behavior is entirely implementation-dependent.
275///
276/// # Safety
277///
278/// `ptr` must have been allocated by `Jemalloc` and must not have been freed yet.
279pub unsafe fn usable_size<T>(ptr: *const T) -> usize {
280    ffi::malloc_usable_size(ptr as *const c_void)
281}
282
283/// Raw bindings to jemalloc
284mod ffi {
285    pub use tikv_jemalloc_sys::*;
286}
tikv_jemallocator/lib.rs

tikv_jemallocator/
lib.rs
