From cc91a7fe72f39977a9ec8ccef17da0653617fc6f Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Fri, 10 Jan 2025 17:39:34 -0500
Subject: [PATCH] feat: add snmalloc support

---
 Cargo.lock                       | 29 +++++++++++++++++++++++++++++
 Cargo.toml                       |  2 ++
 bin/reth/Cargo.toml              |  6 ++++++
 crates/cli/util/Cargo.toml       | 11 +++++++++++
 crates/cli/util/src/allocator.rs | 15 ++++++++++++++-
 5 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1211d0466f9d..be630932f91d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1859,6 +1859,15 @@ version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f46ad14479a25103f283c0f10005961cf086d8dc42205bb44c46ac563475dca6"
 
+[[package]]
+name = "cmake"
+version = "0.1.52"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c682c223677e0e5b6b7f63a64b9351844c3f1b1678a68b7ee617e30fb082620e"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "codspeed"
 version = "2.7.2"
@@ -6680,6 +6689,7 @@ dependencies = [
  "reth-fs-util",
  "secp256k1",
  "serde",
+ "snmalloc-rs",
  "thiserror 2.0.9",
  "tikv-jemallocator",
  "tracy-client",
@@ -10416,6 +10426,25 @@ version = "1.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b"
 
+[[package]]
+name = "snmalloc-rs"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d43ff92911d7d9705d1c0203300a3edfd00d16c8b8b0c27c56f9407a3f31e7a6"
+dependencies = [
+ "snmalloc-sys",
+]
+
+[[package]]
+name = "snmalloc-sys"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "954e1f984860770475196be81a547ed1517d34fcb8a15cb87bdb37cff3353230"
+dependencies = [
+ "cc",
+ "cmake",
+]
+
 [[package]]
 name = "socket2"
 version = "0.5.8"
diff --git a/Cargo.toml b/Cargo.toml
index a49d3052a557..0e6a39084b2e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -604,9 +604,11 @@ tempfile = "3.8"
 test-fuzz = "6"
 rstest = "0.23.0"
 
+# allocators
 tikv-jemalloc-ctl = "0.6"
 tikv-jemallocator = "0.6"
 tracy-client = "0.17.3"
+snmalloc-rs = { version = "0.3.7", features = ["build_cc"] }
 
 # [patch.crates-io]
 # alloy-consensus = { git = "https://github.com/alloy-rs/alloy", rev = "5492e40" }
diff --git a/bin/reth/Cargo.toml b/bin/reth/Cargo.toml
index f7bdfd8ceed2..8487a08ea976 100644
--- a/bin/reth/Cargo.toml
+++ b/bin/reth/Cargo.toml
@@ -118,6 +118,12 @@ jemalloc-prof = [
 ]
 tracy-allocator = ["reth-cli-util/tracy-allocator"]
 
+# Because jemalloc is default and preferred over snmalloc when both features are
+# enabled, `--no-default-features` should be used when enabling snmalloc or
+# snmalloc-native.
+snmalloc = ["reth-cli-util/snmalloc"]
+snmalloc-native = ["reth-cli-util/snmalloc-native"]
+
 min-error-logs = ["tracing/release_max_level_error"]
 min-warn-logs = ["tracing/release_max_level_warn"]
 min-info-logs = ["tracing/release_max_level_info"]
diff --git a/crates/cli/util/Cargo.toml b/crates/cli/util/Cargo.toml
index 70515f83b4b7..b7f0f11ff214 100644
--- a/crates/cli/util/Cargo.toml
+++ b/crates/cli/util/Cargo.toml
@@ -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"]
diff --git a/crates/cli/util/src/allocator.rs b/crates/cli/util/src/allocator.rs
index ee13e7c61cb5..753c987324d7 100644
--- a/crates/cli/util/src/allocator.rs
+++ b/crates/cli/util/src/allocator.rs
@@ -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>;