feat: add snmalloc support (#13771)

This commit is contained in:
Dan Cline
2025-01-11 21:43:24 -05:00
committed by GitHub
parent cc84f83b6c
commit 567d5c60e6
5 changed files with 62 additions and 1 deletions

View File

@ -30,10 +30,21 @@ tracy-client = { workspace = true, optional = true, features = ["demangle"] }
[target.'cfg(unix)'.dependencies]
tikv-jemallocator = { workspace = true, optional = true }
snmalloc-rs = { workspace = true, optional = true }
libc = "0.2"
[features]
jemalloc = ["dep:tikv-jemallocator"]
# Enables jemalloc profiling features
jemalloc-prof = ["jemalloc", "tikv-jemallocator?/profiling"]
# Wraps the selected allocator in the tracy profiling allocator
tracy-allocator = ["dep:tracy-client"]
snmalloc = ["dep:snmalloc-rs"]
# Enables the snmalloc-rs `native-cpu` feature, which optimizes snmalloc for the
# native CPU of the host machine. Not sure why this feature is not derived from
# RUSTFLAGS or enabled when `target-cpu=native`.
snmalloc-native = ["snmalloc", "snmalloc-rs/native-cpu"]

View File

@ -1,14 +1,27 @@
//! Custom allocator implementation.
//!
//! We provide support for jemalloc and snmalloc on unix systems, and prefer jemalloc if both are
//! enabled.
// We use jemalloc for performance reasons.
// We provide jemalloc allocator support, alongside snmalloc. If both features are enabled, jemalloc
// is prioritized.
cfg_if::cfg_if! {
if #[cfg(all(feature = "jemalloc", unix))] {
type AllocatorInner = tikv_jemallocator::Jemalloc;
} else if #[cfg(all(feature = "snmalloc", unix))] {
type AllocatorInner = snmalloc_rs::SnMalloc;
} else {
type AllocatorInner = std::alloc::System;
}
}
// This is to prevent clippy unused warnings when we do `--all-features`
cfg_if::cfg_if! {
if #[cfg(all(feature = "snmalloc", feature = "jemalloc", unix))] {
use snmalloc_rs as _;
}
}
cfg_if::cfg_if! {
if #[cfg(feature = "tracy-allocator")] {
type AllocatorWrapper = tracy_client::ProfiledAllocator<AllocatorInner>;