[DRAFT] Performance impact of jemalloc profiling

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 = ["jemalloc_pprof", "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,9 @@ 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"] }
+jemalloc_pprof = { version = "0.6", optional = true }
 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 +163,9 @@ 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"] }
+jemalloc_pprof = { version = "0.6" }
 
 [build-dependencies]
 cfg_aliases = "0.1.1"

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]

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)");
 }

ipa-core/src/bin/helper.rs

@@ -30,9 +30,18 @@ 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(jemalloc)]
+#[allow(non_upper_case_globals)]
+#[export_name = "_rjem_malloc_conf"]
+pub static _rjem_malloc_conf: &[u8] = b"prof:true,prof_active:false\0";
+
+#[cfg(feature = "dhat-heap")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
 
 #[derive(Debug, Parser)]
 #[clap(