From abff373305572e6d706fb296959b4936867b289d Mon Sep 17 00:00:00 2001
From: Andy Leiserson <aleiserson@mozilla.com>
Date: Mon, 16 Dec 2024 13:37:37 -0800
Subject: [PATCH] Upgrade jemalloc to 0.6; add jemalloc feature

jemalloc is used by default on linux. The feature enables it on
non-linux platforms.
---
 ipa-core/Cargo.toml             | 11 ++++++++---
 ipa-core/benches/oneshot/ipa.rs |  8 ++------
 ipa-core/build.rs               | 11 +++++++++++
 ipa-core/src/bin/helper.rs      |  8 ++++++--
 4 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/ipa-core/Cargo.toml b/ipa-core/Cargo.toml
index 8bb49f189..101416e4f 100644
--- a/ipa-core/Cargo.toml
+++ b/ipa-core/Cargo.toml
@@ -63,7 +63,9 @@ enable-benches = ["cli", "in-memory-infra", "test-fixture", "criterion", "iai"]
 # of unit tests use it. Real world infra uses HTTP implementation and is suitable for integration/e2e tests
 in-memory-infra = []
 real-world-infra = []
-dhat-heap = ["cli", "test-fixture"]
+# Force use of jemalloc on non-Linux platforms. jemalloc is used by default on Linux.
+jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"]
+dhat-heap = ["cli", "dhat", "test-fixture"]
 # Enable this feature to enable our colossally weak Fp31.
 weak-field = []
 # Enable using more than one thread for protocol execution. Most of the parallelism occurs at parallel/seq_join operations
@@ -111,7 +113,7 @@ criterion = { version = "0.5.1", optional = true, default-features = false, feat
 curve25519-dalek = "4.1.1"
 dashmap = "5.4"
 delegate = "0.10.0"
-dhat = "0.3.2"
+dhat = { version = "0.3.2", optional = true }
 embed-doc-image = "0.1.4"
 futures = "0.3.28"
 futures-util = "0.3.28"
@@ -144,6 +146,8 @@ sha2 = "0.10"
 shuttle-crate = { package = "shuttle", version = "0.6.1", optional = true }
 subtle = "2.6"
 thiserror = "1.0"
+tikv-jemallocator = { version = "0.6", optional = true, features = ["profiling"] }
+tikv-jemalloc-ctl = { version = "0.6", optional = true, features = ["stats"] }
 time = { version = "0.3", optional = true }
 tokio = { version = "1.42", features = ["fs", "rt", "rt-multi-thread", "macros"] }
 tokio-rustls = { version = "0.26", optional = true }
@@ -158,7 +162,8 @@ typenum = { version = "1.17", features = ["i128"] }
 x25519-dalek = "2.0.0-rc.3"
 
 [target.'cfg(all(not(target_env = "msvc"), not(target_os = "macos")))'.dependencies]
-tikv-jemallocator = "0.5.0"
+tikv-jemallocator = { version = "0.6", features = ["profiling"] }
+tikv-jemalloc-ctl = { version = "0.6", features = ["stats"] }
 
 [build-dependencies]
 cfg_aliases = "0.1.1"
diff --git a/ipa-core/benches/oneshot/ipa.rs b/ipa-core/benches/oneshot/ipa.rs
index b24d7ea7b..1f9eec376 100644
--- a/ipa-core/benches/oneshot/ipa.rs
+++ b/ipa-core/benches/oneshot/ipa.rs
@@ -19,13 +19,9 @@ use ipa_step::StepNarrow;
 use rand::{random, rngs::StdRng, SeedableRng};
 use tokio::runtime::Builder;
 
-#[cfg(all(
-    not(target_env = "msvc"),
-    not(feature = "dhat-heap"),
-    not(target_os = "macos")
-))]
+#[cfg(jemalloc)]
 #[global_allocator]
-static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
 #[cfg(feature = "dhat-heap")]
 #[global_allocator]
diff --git a/ipa-core/build.rs b/ipa-core/build.rs
index ce1987c72..07d70ffd4 100644
--- a/ipa-core/build.rs
+++ b/ipa-core/build.rs
@@ -47,10 +47,21 @@ fn main() {
         descriptive_gate: { all(not(feature = "compact-gate"), feature = "descriptive-gate") },
         unit_test: { all(not(feature = "shuttle"), feature = "in-memory-infra", descriptive_gate) },
         web_test: { all(not(feature = "shuttle"), feature = "real-world-infra") },
+        jemalloc: { all(
+            not(feature = "dhat-heap"),
+            any(
+                feature = "jemalloc",
+                all(
+                    not(target_env = "msvc"),
+                    not(target_os = "macos")
+                )
+            )
+        ) },
     }
     println!("cargo::rustc-check-cfg=cfg(descriptive_gate)");
     println!("cargo::rustc-check-cfg=cfg(compact_gate)");
     println!("cargo::rustc-check-cfg=cfg(unit_test)");
     println!("cargo::rustc-check-cfg=cfg(web_test)");
+    println!("cargo::rustc-check-cfg=cfg(jemalloc)");
     println!("cargo::rustc-check-cfg=cfg(coverage)");
 }
diff --git a/ipa-core/src/bin/helper.rs b/ipa-core/src/bin/helper.rs
index 85b0c7110..8622ea998 100644
--- a/ipa-core/src/bin/helper.rs
+++ b/ipa-core/src/bin/helper.rs
@@ -30,9 +30,13 @@ use ipa_core::{
 use tokio::runtime::Runtime;
 use tracing::{error, info};
 
-#[cfg(all(not(target_env = "msvc"), not(target_os = "macos")))]
+#[cfg(jemalloc)]
 #[global_allocator]
-static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+#[cfg(feature = "dhat-heap")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
 
 #[derive(Debug, Parser)]
 #[clap(