From 552f0b1c8d903e95296c91d487e5a857c564de06 Mon Sep 17 00:00:00 2001 From: Austin Seipp Date: Wed, 1 Nov 2023 22:48:38 -0500 Subject: [PATCH] cbits: implement `MiMalloc` for use with `#[global_allocator]` Summary: This module is 100% abstract except for this single export. No other API is supported at this time, by design. Signed-off-by: Austin Seipp Change-Id: I0316dcf3ec85972336056e8fdca86f33 --- cbits/src/lib.rs | 2 + cbits/src/mimalloc.rs | 289 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 291 insertions(+) create mode 100644 cbits/src/mimalloc.rs diff --git a/cbits/src/lib.rs b/cbits/src/lib.rs index 00b5be4cf0..f65c68f7f3 100644 --- a/cbits/src/lib.rs +++ b/cbits/src/lib.rs @@ -56,3 +56,5 @@ // Disable std support to avoid recursive dependencies, since this library // contains an allocator. #![no_std] + +pub mod mimalloc; diff --git a/cbits/src/mimalloc.rs b/cbits/src/mimalloc.rs new file mode 100644 index 0000000000..d5e74c18ca --- /dev/null +++ b/cbits/src/mimalloc.rs @@ -0,0 +1,289 @@ +// Copyright 2024 The Jujutsu Authors +// Copyright 2019 Octavian Oncescu +// +// Taken from mimalloc_rust library. mimalloc_rust contains the following +// license: +// +// Permission is hereby granted, free of charge, to any person obtaining a copy +// of this software and associated documentation files (the "Software"), to deal +// in the Software without restriction, including without limitation the rights +// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +// copies of the Software, and to permit persons to whom the Software is +// furnished to do so, subject to the following conditions: +// +// The above copyright notice and this permission notice shall be included in +// all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +//! A module for using the **[mimalloc](https://github.com/microsoft/mimalloc)** +//! memory allocator in Rust programs. mimalloc is a small, easy-to-understand, +//! scalable, cache-and-thread friendly memory allocator. It is written in C, +//! has no external dependencies, and is linked and exists entirely inside of +//! this crate. +//! +//! By design, this module is nearly opaque. The only supported APIs are using +//! [`MiMalloc`] with the [`#[global_allocator]` attribute][module@core::alloc], +//! and some tools for diagnostics. In the future this may expand to features +//! like dedicated heaps or arenas, if needed. +//! +//! ## Motivation +//! +//! This is effectively a soft fork of the [mimalloc](https://docs.rs/mimalloc) +//! and [mimalloc-sys](https://docs.rs/mimalloc-sys) crates, merging them into +//! one module for our own needs with some extra cherries on top. There are also +//! an array of third party crates providing mimalloc support. +//! +//! The motivation for our own crates include, but are not limited to: +//! +//! - Usage of mimalloc 2.x, while many upstream crates use seem to still use +//! the 1.x series, +//! - Reduce "crate bloat", as we have no need for separate `foo`/`foo-sys` +//! designs, +//! - Have a space for _other_ C code in the future, +//! - Have better control over the exported API for deeper integration. +//! +//! Please read the top-level documentation for the [jj_cbits](../index.html) +//! crate for more information. +//! +//! In the future, it is possible these changes may be integrated back into a +//! more general, widely-usable crate. +//! +//! ## Basic usage +//! +//! ```rust,ignore +//! #[global_allocator] +//! static ALLOC: jj_cbits::mimalloc::MiMalloc = jj_cbits::mimalloc::MiMalloc; +//! ``` + +use core::alloc::GlobalAlloc; +use core::ffi::{c_char, c_ulonglong, c_void, CStr}; + +// Full set of external C APIs. +// +// NOTE: These should under no circumstances ever be made `pub`. +extern "C" { + /// Allocate a block of memory of at least `size` bytes, aligned to the + /// given alignment, that is completely zeroed out. + fn mi_zalloc_aligned(size: usize, alignment: usize) -> *mut c_void; + + /// Allocate a block of memory of at least `size` bytes, aligned to the + /// given alignment. The initial contents of the memory is unspecified, and + /// accessing it before initialization should be considered undefined + /// behavior. + fn mi_malloc_aligned(size: usize, alignment: usize) -> *mut c_void; + + /// Reallocate a block of memory to a new size. + fn mi_realloc_aligned(p: *mut c_void, newsize: usize, alignment: usize) -> *mut c_void; + + /// Free any previously allocated memory. + fn mi_free(p: *mut c_void); + + /// Merge all thread-local statistics with the main statistics buffer, and + /// reset. + fn mi_stats_merge(); + + /// Reset the internal statistics counters. + fn mi_stats_reset(); + + /// Print global mimalloc statistics to a given output. The output is + /// currently fixed to `stderr`; both the first and second arguments must be + /// the NULL pointer. + fn mi_stats_print_out(out: unsafe extern "C" fn(*const c_char, *mut c_void), ctx: *mut c_void); + + /// Register a "deferred free" function which will be called + /// deterministically over the course of the program. + fn mi_register_deferred_free( + f: unsafe extern "C" fn(bool, c_ulonglong, *mut c_void), + ctx: *mut c_void, + ); +} + +/// Provide global mimalloc heap statistics to a user-provided callback. +/// +/// The provided callback function will be called multiple times in sequence, +/// each time with a single argument, which is a single null-terminated line, +/// represented as a [CStr]. Collectively, these lines will represent a summary +/// of the global heap statistics for the entire program, meant to be written to +/// the terminal or a log file. +/// +/// Therefore, the simplest way to use this function is to simply provide a +/// closure that prints the given log messages to `stderr` immediately: +/// +/// ```rust,ignore +/// eprintln!("========================================"); +/// eprintln!("mimalloc memory allocation statistics:\n"); +/// jj_cbits::mimalloc::stats_print(&|l| { +/// eprint!("{}", l.to_string_lossy()); +/// }); +/// ``` +/// +/// Note that this merges all thread-local statistics into the main statistics +/// summary before printing to `stderr`, so while it will give a global summary +/// of the heap, it may cause some performance overhead while thread-local +/// buffers are being flushed and merged. +pub fn stats_print(f: &'static F) { + unsafe { + unsafe extern "C" fn wrapper(value: *const c_char, ctx: *mut c_void) { + (*(ctx as *const F))(CStr::from_ptr(value)); + } + mi_stats_merge(); + mi_stats_print_out(wrapper::, f as *const F as *mut c_void) + } +} + +/// Reset heap statistics counters and histograms. +/// +/// Primarily useful to clear out any existing statistics, so that a subsequent +/// call to `mimalloc_stats_print` will only show statistics since the last +/// reset. +/// +/// This should also reset and merge all thread-local statistics, too. +pub fn stats_reset() { + unsafe { + mi_stats_reset(); + } +} + +/// Register a "deferred free" function, which will be called by the memory +/// allocator after some (deterministic) number of calls to +/// [`dealloc`](core::alloc::GlobalAlloc::dealloc) in the heap. +/// +/// Typically, the callback function is provided as a simple closure with static +/// lifetime, as it may be called at any point in the program's lifetime. The +/// result of the closure is ignored and has no meaning. +/// +/// The provided closure will be invoked at an unspecified future point with the +/// following arguments: +/// +/// * `force`, type `bool`: If `true`, the deferred free function should free +/// any memory it has allocated, or that may be possible to free to reduce +/// heap pressure. +/// * `count`, type `c_ulonglong`: A monotonically increasing "heartbeat +/// counter." May be assigned any semantic meaning to your program that you +/// desire. This counter MUST NOT be assumed to have any relation to the +/// structure of the heap, in any way. +/// +/// These two parameters are completely independent from each other; that is, +/// any combination of `force` and `count` may be provided to the callback and +/// should not be assumed to influence each other in any meaningful way. +/// +/// Note that this function is called _deterministically_ based on heap +/// allocations. Therefore, assuming the program itself exhibits deterministic +/// allocation behavior the resulting deferred free callback will also be called +/// deterministically over the program's lifetime. The number of allocations +/// between invocations is unspecified. +/// +/// Despite the name, this registered callback does not need to free any extra +/// memory in any way, and can be used purely as a "heartbeat" mechanism to +/// implement other functionality, such as periodic state logging or timeouts +/// that are not tied to the wall clock. +/// +/// There may be only a single deferred free function registered at any given +/// time. If this function is called multiple times, the last registered +/// function will be used. +/// +/// Reference: +/// +/// - Section 2.3 _The Local Free List_; Leijen 2019, "[Mimalloc: Free List +/// Sharding in Action][mimalloc-pdf]", MSR-TR 2019-18. +/// +/// [mimalloc-pdf]: +/// https://www.microsoft.com/en-us/research/uploads/prod/2019/06/mimalloc-tr-v1.pdf +pub fn register_deferred_free(f: &'static F) { + unsafe { + unsafe extern "C" fn wrapper( + force: bool, + count: c_ulonglong, + ctx: *mut c_void, + ) { + (*(ctx as *const F))(force, count); + } + mi_register_deferred_free(wrapper::, f as *const F as *mut c_void) + } +} + +/// Global memory allocator, based on the mimalloc library. +/// +/// ## Usage +/// +/// Inside of the `main.rs` for any binary: +/// +/// ```rust,ignore +/// #[global_allocator] +/// static ALLOC: jj_cbits::mimalloc::MiMalloc = jj_cbits::mimalloc::MiMalloc; +/// ``` +pub struct MiMalloc; + +unsafe impl GlobalAlloc for MiMalloc { + #[inline] + unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 { + mi_malloc_aligned(layout.size(), layout.align()) as *mut u8 + } + + #[inline] + unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 { + mi_zalloc_aligned(layout.size(), layout.align()) as *mut u8 + } + + #[inline] + unsafe fn dealloc(&self, ptr: *mut u8, _layout: core::alloc::Layout) { + mi_free(ptr as *mut c_void) + } + + #[inline] + unsafe fn realloc( + &self, + ptr: *mut u8, + layout: core::alloc::Layout, + new_size: usize, + ) -> *mut u8 { + mi_realloc_aligned(ptr as *mut c_void, new_size, layout.align()) as *mut u8 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ok_free_malloc() { + let ptr = unsafe { mi_malloc_aligned(8, 8) } as *mut u8; + unsafe { mi_free(ptr as *mut c_void) }; + } + + #[test] + fn ok_free_zalloc() { + let ptr = unsafe { mi_zalloc_aligned(8, 8) } as *mut u8; + unsafe { mi_free(ptr as *mut c_void) }; + } + + #[test] + fn ok_free_realloc() { + let ptr = unsafe { mi_malloc_aligned(8, 8) } as *mut u8; + let ptr = unsafe { mi_realloc_aligned(ptr as *mut c_void, 8, 8) } as *mut u8; + unsafe { mi_free(ptr as *mut c_void) }; + } + + #[test] + fn ok_usable_size() { + extern "C" { + // not used elsewhere, so scope it here to the tests to avoid + // spurious dead_code warnings + fn mi_usable_size(p: *const c_void) -> usize; + } + + let ptr = unsafe { mi_malloc_aligned(32, 64) } as *mut u8; + let usable_size = unsafe { mi_usable_size(ptr as *mut c_void) }; + assert!( + usable_size >= 32, + "usable_size should at least equal to the allocated size" + ); + } +}