Derive the allocator used by tools from rustc's allocator

compiler/rustc/Cargo.toml

@@ -35,7 +35,7 @@ features = ['override_allocator_on_supported_platforms']
 [features]
 # tidy-alphabetical-start
 check_only = ['rustc_driver_impl/check_only']
-jemalloc = ['dep:tikv-jemalloc-sys']
+jemalloc = ['dep:tikv-jemalloc-sys', 'rustc_driver_impl/jemalloc']
 llvm = ['rustc_driver_impl/llvm']
 llvm_offload = ['rustc_driver_impl/llvm_offload']
 max_level_info = ['rustc_driver_impl/max_level_info']

compiler/rustc/src/main.rs

@@ -7,11 +7,8 @@ use std::process::ExitCode;
 
 // A note about jemalloc: rustc uses jemalloc when built for CI and
 // distribution. The obvious way to do this is with the `#[global_allocator]`
-// mechanism. However, for complicated reasons (see
-// https://github.com/rust-lang/rust/pull/81782#issuecomment-784438001 for some
-// details) that mechanism doesn't work here. Also, we'd like to use a
-// consistent allocator across the rustc <-> llvm boundary, and
-// `#[global_allocator]` wouldn't provide that.
+// mechanism. However, that would not affect LLVM's C / C++ allocations and we also want
+// to use a single allocator in the process to reduce memory usage.
 //
 // Instead, we use a lower-level mechanism, namely the
 // `"override_allocator_on_supported_platforms"` Cargo feature of jemalloc-sys.
@@ -20,15 +17,14 @@ use std::process::ExitCode;
 // of `malloc`, `free`, etc.. This means that Rust's `System` allocator, which
 // calls `libc::malloc()` et al., is actually calling into jemalloc.
 //
+// This override happens for the entire process, ensuring that there's no mixup
+// of C allocators across dylibs / binaries, notably the rustc <-> llvm boundary.
+//
 // A consequence of not using `GlobalAlloc` (and the `tikv-jemallocator` crate
 // provides an impl of that trait, which is called `Jemalloc`) is that we
 // cannot use the sized deallocation APIs (`sdallocx`) that jemalloc provides.
 // It's unclear how much performance is lost because of this.
 //
-// NOTE: Even though Cargo passes `--extern` with `tikv_jemalloc_sys`, we still need to `use` the
-// crate for the compiler to see the `#[used]`, see https://github.com/rust-lang/rust/issues/64402.
-// This is similarly required if we used a crate with `#[global_allocator]`.
-//
 // NOTE: if you are reading this comment because you want to set a custom `global_allocator` for
 // benchmarking, consider using the benchmarks in the `rustc-perf` collector suite instead:
 // https://github.com/rust-lang/rustc-perf/blob/master/collector/README.md#profiling
@@ -37,8 +33,7 @@ use std::process::ExitCode;
 // to compare their performance, see
 // https://github.com/rust-lang/rust/commit/b90cfc887c31c3e7a9e6d462e2464db1fe506175#diff-43914724af6e464c1da2171e4a9b6c7e607d5bc1203fa95c0ab85be4122605ef
 // for an example of how to do so.
-#[cfg(feature = "jemalloc")]
-use tikv_jemalloc_sys as _;
+rustc_driver::override_c_allocator_in_binary!();
 
 fn main() -> ExitCode {
     rustc_driver::main()

compiler/rustc_driver_impl/Cargo.toml

@@ -57,6 +57,7 @@ ctrlc = "3.4.4"
 [features]
 # tidy-alphabetical-start
 check_only = ['rustc_interface/check_only']
+jemalloc = []
 llvm = ['rustc_interface/llvm']
 llvm_offload = ['rustc_interface/llvm_offload']
 max_level_info = ['rustc_log/max_level_info']

compiler/rustc_driver_impl/src/allocator.rs

@@ -0,0 +1,24 @@
+/// This macro overrides the C allocator (i.e., `malloc`) in final binaries by linking
+/// jemalloc with the override feature enabled. The C allocator is used by the default
+/// Rust allocator (`alloc::System`) on Unix targets but not on Windows targets.
+#[cfg(feature = "jemalloc")]
+#[macro_export]
+macro_rules! override_c_allocator_in_binary {
+    () => {
+        // NOTE: even though Cargo passes `--extern` for this in rustc-main, the crate still has
+        // to be named for the compiler to see the `#[used]` inside it, see
+        // <https://github.com/rust-lang/rust/issues/64402>.
+        //
+        // FIXME(madsmtm): for the rustc-private tools this is loaded from the sysroot that was
+        // built with the other `rustc` crates, instead of via Cargo as you'd normally do. This is
+        // currently needed for LTO due to <https://github.com/rust-lang/cc-rs/issues/1613>.
+        extern crate tikv_jemalloc_sys as _;
+    };
+}
+
+/// This macro does nothing when no allocator features are enabled.
+#[cfg(not(feature = "jemalloc"))]
+#[macro_export]
+macro_rules! override_c_allocator_in_binary {
+    () => {};
+}

compiler/rustc_driver_impl/src/lib.rs

@@ -82,6 +82,7 @@ macro do_not_use_safe_print($($t:tt)*) {
 #[allow(unused_imports)]
 use {do_not_use_print as print, do_not_use_print as println};
 
+mod allocator;
 pub mod args;
 pub mod pretty;
 #[macro_use]

src/bootstrap/src/core/build_steps/tool.rs

@@ -765,11 +765,7 @@ impl CommandLineStep for Rustdoc {
         // they'll be linked to those libraries). As such, don't explicitly `ensure` any additional
         // libraries here. The intuition here is that If we've built a compiler, we should be able
         // to build rustdoc.
-        //
         let mut extra_features = Vec::new();
-        if let Some(allocator_feature_name) = builder.config.allocator(target).feature_name() {
-            extra_features.push(allocator_feature_name.to_string());
-        }
         if !builder.config.rust_debug_logging {
             extra_features.push("max_level_info".to_string())
         }
@@ -1428,7 +1424,6 @@ macro_rules! tool_rustc_extended {
             tool_name: $tool_name:expr,
             stable: $stable:expr
             $( , add_bins_to_sysroot: $add_bins_to_sysroot:expr )?
-            $( , add_features: $add_features:expr )?
             $( , cargo_args: $cargo_args:expr )?
             $( , )?
         }
@@ -1479,7 +1474,6 @@ macro_rules! tool_rustc_extended {
                     $tool_name,
                     $path,
                     None $( .or(Some(&$add_bins_to_sysroot)) )?,
-                    None $( .or(Some($add_features)) )?,
                     None $( .or(Some($cargo_args)) )?,
                 )
             }
@@ -1524,23 +1518,17 @@ fn build_extended_rustc_tool(
     tool_name: &'static str,
     path: &'static str,
     add_bins_to_sysroot: Option<&[&str]>,
-    add_features: Option<fn(&Builder<'_>, TargetSelection, &mut Vec<String>)>,
     cargo_args: Option<&[&'static str]>,
 ) -> ToolBuildResult {
     let target = compilers.target();
-    let mut extra_features = Vec::new();
-    if let Some(func) = add_features {
-        func(builder, target, &mut extra_features);
-    }
-
     let build_compiler = compilers.build_compiler;
     let ToolBuildResult { tool_path, .. } = builder.ensure(ToolBuild {
         build_compiler,
         target,
         tool: tool_name,
         mode: Mode::ToolRustcPrivate,
         path,
-        extra_features,
+        extra_features: Vec::new(),
         source_type: SourceType::InTree,
         allow_features: "",
         cargo_args: cargo_args.unwrap_or_default().iter().map(|s| String::from(*s)).collect(),
@@ -1583,23 +1571,13 @@ tool_rustc_extended!(Clippy {
     path: "src/tools/clippy",
     tool_name: "clippy-driver",
     stable: true,
-    add_bins_to_sysroot: ["clippy-driver"],
-    add_features: |builder, target, features| {
-        if let Some(allocator_feature_name) = builder.config.allocator(target).feature_name() {
-            features.push(allocator_feature_name.to_string());
-        }
-    }
+    add_bins_to_sysroot: ["clippy-driver"]
 });
 tool_rustc_extended!(Miri {
     path: "src/tools/miri",
     tool_name: "miri",
     stable: false,
     add_bins_to_sysroot: ["miri"],
-    add_features: |builder, target, features| {
-        if let Some(allocator_feature_name) = builder.config.allocator(target).feature_name() {
-            features.push(allocator_feature_name.to_string());
-        }
-    },
     // Always compile also tests when building miri. Otherwise feature unification can cause rebuilds between building and testing miri.
     cargo_args: &["--all-targets"],
 });

src/librustdoc/Cargo.toml

@@ -43,7 +43,6 @@ minifier = { version = "0.4.0", default-features = false }
 expect-test = "1.4.0"
 
 [features]
-jemalloc = []
 max_level_info = ["tracing/max_level_info", "tracing/release_max_level_info"]
 
 [package.metadata.rust-analyzer]

src/librustdoc/lib.rs

@@ -58,15 +58,6 @@ extern crate rustc_target;
 extern crate rustc_trait_selection;
 extern crate test;
 
-/// See docs in https://github.com/rust-lang/rust/blob/HEAD/compiler/rustc/src/main.rs
-/// and https://github.com/rust-lang/rust/pull/146627 for why we need this.
-///
-/// FIXME(madsmtm): This is loaded from the sysroot that was built with the other `rustc` crates
-/// above, instead of via Cargo as you'd normally do. This is currently needed for LTO due to
-/// https://github.com/rust-lang/cc-rs/issues/1613.
-#[cfg(feature = "jemalloc")]
-extern crate tikv_jemalloc_sys as _;
-
 use std::env::{self, VarError};
 use std::io::{self, IsTerminal};
 use std::path::Path;

src/tools/clippy/Cargo.toml

@@ -56,7 +56,6 @@ rustc_tools_util = { path = "rustc_tools_util", version = "0.4.2" }
 [features]
 integration = ["dep:tempfile"]
 internal = ["dep:clippy_lints_internal", "dep:tempfile"]
-jemalloc = []
 
 [package.metadata.rust-analyzer]
 # This package uses #[feature(rustc_private)]

src/tools/clippy/src/driver.rs

@@ -11,14 +11,8 @@ extern crate rustc_interface;
 extern crate rustc_session;
 extern crate rustc_span;
 
-/// See docs in <https://github.com/rust-lang/rust/blob/HEAD/compiler/rustc/src/main.rs>
-/// and <https://github.com/rust-lang/rust/pull/146627> for why we need this.
-///
-/// FIXME(madsmtm): This is loaded from the sysroot that was built with the other `rustc` crates
-/// above, instead of via Cargo as you'd normally do. This is currently needed for LTO due to
-/// <https://github.com/rust-lang/cc-rs/issues/1613>.
-#[cfg(feature = "jemalloc")]
-extern crate tikv_jemalloc_sys as _;
+// Override the C allocator in the same way that the `rustc` binary would do.
+rustc_driver::override_c_allocator_in_binary!();
 
 use clippy_utils::sym;
 use declare_clippy_lint::LintListBuilder;

src/tools/miri/Cargo.toml

@@ -67,7 +67,6 @@ stack-cache = []
 expensive-consistency-checks = ["stack-cache"]
 tracing = ["serde_json"]
 native-lib = ["dep:libffi", "dep:libloading", "dep:capstone", "dep:ipc-channel", "dep:nix", "dep:serde"]
-jemalloc = []
 check_only = ["libffi?/check_only", "capstone?/check_only", "genmc-sys?/check_only"]
 
 [lints.rust.unexpected_cfgs]

src/tools/miri/src/bin/miri.rs

@@ -16,19 +16,8 @@ extern crate rustc_log;
 extern crate rustc_middle;
 extern crate rustc_session;
 
-/// See docs in https://github.com/rust-lang/rust/blob/HEAD/compiler/rustc/src/main.rs
-/// and https://github.com/rust-lang/rust/pull/146627 for why we need this.
-///
-/// FIXME(madsmtm): This is loaded from the sysroot that was built with the other `rustc` crates
-/// above, instead of via Cargo as you'd normally do. This is currently needed for LTO due to
-/// https://github.com/rust-lang/cc-rs/issues/1613.
-#[cfg(feature = "jemalloc")]
-// Make sure `--all-features` works: only Linux and macOS actually use jemalloc, and not on arm32.
-#[cfg(all(
-    any(target_os = "linux", target_os = "macos"),
-    any(target_arch = "x86_64", target_arch = "x86", target_arch = "aarch64"),
-))]
-extern crate tikv_jemalloc_sys as _;
+// Override the C allocator in the same way that the `rustc` binary would do.
+rustc_driver::override_c_allocator_in_binary!();
 
 mod log;
 

src/tools/rustdoc/Cargo.toml

@@ -14,5 +14,7 @@ path = "main.rs"
 rustdoc = { path = "../../librustdoc" }
 
 [features]
-jemalloc = ['rustdoc/jemalloc']
 max_level_info = ["rustdoc/max_level_info"]
+
+[package.metadata.rust-analyzer]
+rustc_private = true

src/tools/rustdoc/main.rs

@@ -1,8 +1,13 @@
 // We need this feature as it changes `dylib` linking behavior and allows us to link to `rustc_driver`.
 #![feature(rustc_private)]
 
+extern crate rustc_driver;
+
 use std::process::ExitCode;
 
+// Override the C allocator in the same way that the `rustc` binary would do.
+rustc_driver::override_c_allocator_in_binary!();
+
 fn main() -> ExitCode {
     rustdoc::main()
 }