Auto merge of #83152 - guswynn:jemallocator_part2, r=Mark-Simulacrum

Use tikv-jemallocator in rustc/rustdoc in addition to jemalloc-sys when enabled.

In https://github.com/rust-lang/rust/pull/81782 it was mentioned that one reason rustc may benefit from minimalloc is it doesn't use the `sdallocx` api from jemalloc.

Currently, on unix, rust uses jemalloc by importing its symbols to use them with the default, System (libc) global allocator.
This PR switches its global alloc to `tikv-jemallocator`, which correctly uses sized deallocation (https://docs.rs/tikv-jemallocator/0.4.1/src/tikv_jemallocator/lib.rs.html#121-126). `tikv-jemallocator`, as far as I can tell, is a more up-to-date set of bindings to jemalloc than `jemallocator`

The perf results of this pr are in large part due to the version upgrade of jemalloc, but sized deallocation has a non-trivial improvement, particularly to rustdoc.

This pr also includes changes to bootstrap to correctly pass the jemalloc feature through to the rustdoc build

4 files changed, 63 insertions, 1 deletions

diff --git a/src/bootstrap/tool.rs b/src/bootstrap/tool.rs
index f9be35d7c5e..e85f4628fb0 100644
--- a/src/bootstrap/tool.rs
+++ b/src/bootstrap/tool.rs
@@ -532,6 +532,11 @@ impl Step 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 features = Vec::new();
+        if builder.config.jemalloc {
+            features.push("jemalloc".to_string());
+        }
 
         let cargo = prepare_tool_cargo(
             builder,
@@ -541,7 +546,7 @@ impl Step for Rustdoc {
             "build",
             "src/tools/rustdoc",
             SourceType::InTree,
-            &[],
+            features.as_slice(),
         );
 
         builder.info(&format!(
diff --git a/src/librustdoc/Cargo.toml b/src/librustdoc/Cargo.toml
index eb38adb0b73..d778a507425 100644
--- a/src/librustdoc/Cargo.toml
+++ b/src/librustdoc/Cargo.toml
@@ -30,5 +30,8 @@ features = ["fmt", "env-filter", "smallvec", "parking_lot", "ansi"]
 [dev-dependencies]
 expect-test = "1.0"
 
+[features]
+jemalloc = []
+
 [package.metadata.rust-analyzer]
 rustc_private = true
diff --git a/src/librustdoc/lib.rs b/src/librustdoc/lib.rs
index c6262f5873e..fc5b7a4f4f3 100644
--- a/src/librustdoc/lib.rs
+++ b/src/librustdoc/lib.rs
@@ -30,6 +30,7 @@ extern crate tracing;
 // So if `rustc` was specified in Cargo.toml, this would spuriously rebuild crates.
 //
 // Dependencies listed in Cargo.toml do not need `extern crate`.
+
 extern crate rustc_ast;
 extern crate rustc_ast_lowering;
 extern crate rustc_ast_pretty;
@@ -60,6 +61,15 @@ extern crate rustc_trait_selection;
 extern crate rustc_typeck;
 extern crate test as testing;
 
+#[cfg(feature = "jemalloc")]
+extern crate tikv_jemalloc_sys;
+#[cfg(feature = "jemalloc")]
+use tikv_jemalloc_sys as jemalloc_sys;
+#[cfg(feature = "jemalloc")]
+extern crate tikv_jemallocator;
+#[cfg(feature = "jemalloc")]
+use tikv_jemallocator as jemallocator;
+
 use std::default::Default;
 use std::env;
 use std::process;
@@ -113,7 +123,48 @@ mod theme;
 mod visit_ast;
 mod visit_lib;
 
+// See docs in https://github.com/rust-lang/rust/blob/master/compiler/rustc/src/main.rs
+// about jemallocator
+#[cfg(feature = "jemalloc")]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
 pub fn main() {
+    // See docs in https://github.com/rust-lang/rust/blob/master/compiler/rustc/src/main.rs
+    // about jemalloc-sys
+    #[cfg(feature = "jemalloc")]
+    {
+        use std::os::raw::{c_int, c_void};
+
+        #[used]
+        static _F1: unsafe extern "C" fn(usize, usize) -> *mut c_void = jemalloc_sys::calloc;
+        #[used]
+        static _F2: unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int =
+            jemalloc_sys::posix_memalign;
+        #[used]
+        static _F3: unsafe extern "C" fn(usize, usize) -> *mut c_void = jemalloc_sys::aligned_alloc;
+        #[used]
+        static _F4: unsafe extern "C" fn(usize) -> *mut c_void = jemalloc_sys::malloc;
+        #[used]
+        static _F5: unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void = jemalloc_sys::realloc;
+        #[used]
+        static _F6: unsafe extern "C" fn(*mut c_void) = jemalloc_sys::free;
+
+        // On OSX, jemalloc doesn't directly override malloc/free, but instead
+        // registers itself with the allocator's zone APIs in a ctor. However,
+        // the linker doesn't seem to consider ctors as "used" when statically
+        // linking, so we need to explicitly depend on the function.
+        #[cfg(target_os = "macos")]
+        {
+            extern "C" {
+                fn _rjem_je_zone_register();
+            }
+
+            #[used]
+            static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
+        }
+    }
+
     rustc_driver::set_sigpipe_handler();
     rustc_driver::install_ice_hook();
 
diff --git a/src/tools/rustdoc/Cargo.toml b/src/tools/rustdoc/Cargo.toml
index 36aa5916da7..d0c047ad6df 100644
--- a/src/tools/rustdoc/Cargo.toml
+++ b/src/tools/rustdoc/Cargo.toml
@@ -13,3 +13,6 @@ path = "main.rs"
 
 [dependencies]
 rustdoc = { path = "../../librustdoc" }
+
+[features]
+jemalloc = ['rustdoc/jemalloc']