From a71de77ae9459e4f05ceb8695fce793b0ac00319 Mon Sep 17 00:00:00 2001
From: Pietro Albini <pietro.albini@ferrous-systems.com>
Date: Mon, 17 Feb 2025 11:29:40 +0100
Subject: [PATCH] allow configuring jemalloc per target

---
 config.example.toml                           |  7 ++++++-
 src/bootstrap/src/core/build_steps/compile.rs |  2 +-
 src/bootstrap/src/core/build_steps/tool.rs    |  2 +-
 src/bootstrap/src/core/config/config.rs       | 10 ++++++++++
 src/bootstrap/src/lib.rs                      |  2 +-
 src/bootstrap/src/utils/change_tracker.rs     |  5 +++++
 6 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/config.example.toml b/config.example.toml
index f5395375afe4c..fd27b24e46191 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -729,7 +729,8 @@
 #remap-debuginfo = false
 
 # Link the compiler and LLVM against `jemalloc` instead of the default libc allocator.
-# This option is only tested on Linux and OSX.
+# This option is only tested on Linux and OSX. It can also be configured per-target in the
+# [target.<tuple>] section.
 #jemalloc = false
 
 # Run tests in various test suites with the "nll compare mode" in addition to
@@ -927,6 +928,10 @@
 # order to run `x check`.
 #optimized-compiler-builtins = build.optimized-compiler-builtins (bool)
 
+# Link the compiler and LLVM against `jemalloc` instead of the default libc allocator.
+# This overrides the global `rust.jemalloc` option. See that option for more info.
+#jemalloc = rust.jemalloc (bool)
+
 # =============================================================================
 # Distribution options
 #
diff --git a/src/bootstrap/src/core/build_steps/compile.rs b/src/bootstrap/src/core/build_steps/compile.rs
index 479327d63695c..64a9b51ed5d69 100644
--- a/src/bootstrap/src/core/build_steps/compile.rs
+++ b/src/bootstrap/src/core/build_steps/compile.rs
@@ -1260,7 +1260,7 @@ pub fn rustc_cargo_env(
 
     // Build jemalloc on AArch64 with support for page sizes up to 64K
     // See: https://github.com/rust-lang/rust/pull/135081
-    if builder.config.jemalloc
+    if builder.config.jemalloc(target)
         && target.starts_with("aarch64")
         && env::var_os("JEMALLOC_SYS_WITH_LG_PAGE").is_none()
     {
diff --git a/src/bootstrap/src/core/build_steps/tool.rs b/src/bootstrap/src/core/build_steps/tool.rs
index 1291a634a6f6f..9d08adaae6bd5 100644
--- a/src/bootstrap/src/core/build_steps/tool.rs
+++ b/src/bootstrap/src/core/build_steps/tool.rs
@@ -643,7 +643,7 @@ impl Step for Rustdoc {
         // to build rustdoc.
         //
         let mut features = Vec::new();
-        if builder.config.jemalloc {
+        if builder.config.jemalloc(target) {
             features.push("jemalloc".to_string());
         }
 
diff --git a/src/bootstrap/src/core/config/config.rs b/src/bootstrap/src/core/config/config.rs
index 62625fc3660f9..9389bed8e550d 100644
--- a/src/bootstrap/src/core/config/config.rs
+++ b/src/bootstrap/src/core/config/config.rs
@@ -325,6 +325,9 @@ pub struct Config {
     pub hosts: Vec<TargetSelection>,
     pub targets: Vec<TargetSelection>,
     pub local_rebuild: bool,
+    #[cfg(not(test))]
+    jemalloc: bool,
+    #[cfg(test)]
     pub jemalloc: bool,
     pub control_flow_guard: bool,
     pub ehcont_guard: bool,
@@ -643,6 +646,7 @@ pub struct Target {
     pub no_std: bool,
     pub codegen_backends: Option<Vec<String>>,
     pub optimized_compiler_builtins: Option<bool>,
+    pub jemalloc: Option<bool>,
 }
 
 impl Target {
@@ -1234,6 +1238,7 @@ define_config! {
         codegen_backends: Option<Vec<String>> = "codegen-backends",
         runner: Option<String> = "runner",
         optimized_compiler_builtins: Option<bool> = "optimized-compiler-builtins",
+        jemalloc: Option<bool> = "jemalloc",
     }
 }
 
@@ -2161,6 +2166,7 @@ impl Config {
                 target.profiler = cfg.profiler;
                 target.rpath = cfg.rpath;
                 target.optimized_compiler_builtins = cfg.optimized_compiler_builtins;
+                target.jemalloc = cfg.jemalloc;
 
                 if let Some(ref backends) = cfg.codegen_backends {
                     let available_backends = ["llvm", "cranelift", "gcc"];
@@ -2726,6 +2732,10 @@ impl Config {
             .unwrap_or(&self.rust_codegen_backends)
     }
 
+    pub fn jemalloc(&self, target: TargetSelection) -> bool {
+        self.target_config.get(&target).and_then(|cfg| cfg.jemalloc).unwrap_or(self.jemalloc)
+    }
+
     pub fn default_codegen_backend(&self, target: TargetSelection) -> Option<String> {
         self.codegen_backends(target).first().cloned()
     }
diff --git a/src/bootstrap/src/lib.rs b/src/bootstrap/src/lib.rs
index 7cd8aacf0d6c8..2adb5993b4f06 100644
--- a/src/bootstrap/src/lib.rs
+++ b/src/bootstrap/src/lib.rs
@@ -677,7 +677,7 @@ impl Build {
             crates.is_empty() || possible_features_by_crates.contains(feature)
         };
         let mut features = vec![];
-        if self.config.jemalloc && check("jemalloc") {
+        if self.config.jemalloc(target) && check("jemalloc") {
             features.push("jemalloc");
         }
         if (self.config.llvm_enabled(target) || kind == Kind::Check) && check("llvm") {
diff --git a/src/bootstrap/src/utils/change_tracker.rs b/src/bootstrap/src/utils/change_tracker.rs
index 9b23cf1843ef5..f215c3f6d0b39 100644
--- a/src/bootstrap/src/utils/change_tracker.rs
+++ b/src/bootstrap/src/utils/change_tracker.rs
@@ -350,4 +350,9 @@ pub const CONFIG_CHANGE_HISTORY: &[ChangeInfo] = &[
         severity: ChangeSeverity::Info,
         summary: "The llvm.ccache option has moved to build.ccache. llvm.ccache is now deprecated.",
     },
+    ChangeInfo {
+        change_id: 137170,
+        severity: ChangeSeverity::Info,
+        summary: "It is now possible to configure `jemalloc` for each target",
+    },
 ];