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

Cargo.lock

@@ -1732,17 +1732,6 @@ version = "0.4.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dc6f3ad7b9d11a0c00842ff8de1b60ee58661048eb8049ed33c73594f359d7e6"
 
-[[package]]
-name = "jemalloc-sys"
-version = "0.3.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0d3b9f3f5c9b31aa0f5ed3260385ac205db665baa41d49bb8338008ae94ede45"
-dependencies = [
- "cc",
- "fs_extra",
- "libc",
-]
-
 [[package]]
 name = "jobserver"
 version = "0.1.21"
@@ -3586,9 +3575,10 @@ checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 name = "rustc-main"
 version = "0.0.0"
 dependencies = [
- "jemalloc-sys",
  "rustc_codegen_ssa",
  "rustc_driver",
+ "tikv-jemalloc-sys",
+ "tikv-jemallocator",
 ]
 
 [[package]]
@@ -5318,6 +5308,27 @@ dependencies = [
 name = "tier-check"
 version = "0.1.0"
 
+[[package]]
+name = "tikv-jemalloc-sys"
+version = "0.4.1+5.2.1-patched"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a26331b05179d4cb505c8d6814a7e18d298972f0a551b0e3cefccff927f86d3"
+dependencies = [
+ "cc",
+ "fs_extra",
+ "libc",
+]
+
+[[package]]
+name = "tikv-jemallocator"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3c14a5a604eb8715bc5785018a37d00739b180bcf609916ddf4393d33d49ccdf"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
 [[package]]
 name = "time"
 version = "0.1.43"

compiler/rustc/Cargo.toml

@@ -11,12 +11,16 @@ rustc_driver = { path = "../rustc_driver" }
 # crate is intended to be used by codegen backends, which may not be in-tree.
 rustc_codegen_ssa = { path = "../rustc_codegen_ssa" }
 
-[dependencies.jemalloc-sys]
-version = '0.3.0'
+[dependencies.tikv-jemalloc-sys]
+version = '0.4.0'
 optional = true
 features = ['unprefixed_malloc_on_supported_platforms']
 
+[dependencies.tikv-jemallocator]
+version = '0.4.0'
+optional = true
+
 [features]
-jemalloc = ['jemalloc-sys']
+jemalloc = ['tikv-jemalloc-sys', 'tikv-jemallocator']
 llvm = ['rustc_driver/llvm']
 max_level_info = ['rustc_driver/max_level_info']

compiler/rustc/src/main.rs

@@ -1,3 +1,16 @@
+// Configure jemalloc as the `global_allocator` when configured. This is
+// so that we use the sized deallocation apis jemalloc provides
+// (namely `sdallocx`).
+//
+// The symbol overrides documented below are also performed so that we can
+// ensure that we use a consistent allocator across the rustc <-> llvm boundary
+#[cfg(feature = "jemalloc")]
+#[global_allocator]
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+#[cfg(feature = "tikv-jemalloc-sys")]
+use tikv_jemalloc_sys as jemalloc_sys;
+
 fn main() {
     // Pull in jemalloc when enabled.
     //
@@ -7,7 +20,7 @@ fn main() {
     // dynamic libraries. That means to pull in jemalloc we actually need to
     // reference allocation symbols one way or another (as this file is the only
     // object code in the rustc executable).
-    #[cfg(feature = "jemalloc-sys")]
+    #[cfg(feature = "tikv-jemalloc-sys")]
     {
         use std::os::raw::{c_int, c_void};
 

src/bootstrap/tool.rs

@@ -519,6 +519,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,
@@ -528,7 +533,7 @@ impl Step for Rustdoc {
             "build",
             "src/tools/rustdoc",
             SourceType::InTree,
-            &[],
+            features.as_slice(),
         );
 
         builder.info(&format!(

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

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();
 

src/tools/rustdoc/Cargo.toml

@@ -13,3 +13,6 @@ path = "main.rs"
 
 [dependencies]
 rustdoc = { path = "../../librustdoc" }
+
+[features]
+jemalloc = ['rustdoc/jemalloc']