Simplify jemalloc setup

Cargo.lock

@@ -600,6 +600,7 @@ dependencies = [
  "serde_json",
  "tempfile",
  "termize",
+ "tikv-jemalloc-sys",
  "toml 0.9.7",
  "ui_test",
  "walkdir",
@@ -4806,6 +4807,7 @@ dependencies = [
  "stringdex",
  "tempfile",
  "threadpool",
+ "tikv-jemalloc-sys",
  "tracing",
  "tracing-subscriber",
  "tracing-tree",
@@ -5537,9 +5539,9 @@ version = "0.1.0"
 
 [[package]]
 name = "tikv-jemalloc-sys"
-version = "0.6.0+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7"
+version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd3c60906412afa9c2b5b5a48ca6a5abe5736aec9eb48ad05037a677e52e4e2d"
+checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b"
 dependencies = [
  "cc",
  "libc",

compiler/rustc/Cargo.toml

@@ -21,9 +21,9 @@ rustc_public_bridge = { path = "../rustc_public_bridge" }
 # tidy-alphabetical-end
 
 [dependencies.tikv-jemalloc-sys]
-version = "0.6.0"
+version = "0.6.1"
 optional = true
-features = ['unprefixed_malloc_on_supported_platforms']
+features = ['override_allocator_on_supported_platforms']
 
 [features]
 # tidy-alphabetical-start

compiler/rustc/src/main.rs

@@ -7,26 +7,25 @@
 // 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 must use a consistent
-// allocator across the rustc <-> llvm boundary, and `#[global_allocator]`
-// wouldn't provide that.
+// 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.
 //
-// Instead, we use a lower-level mechanism. rustc is linked with jemalloc in a
-// way such that jemalloc's implementation of `malloc`, `free`, etc., override
-// the libc allocator's implementation. This means that Rust's `System`
-// allocator, which calls `libc::malloc()` et al., is actually calling into
-// jemalloc.
+// Instead, we use a lower-level mechanism, namely the
+// `"override_allocator_on_supported_platforms"` Cargo feature of jemalloc-sys.
+//
+// This makes jemalloc-sys override the libc/system allocator's implementation
+// of `malloc`, `free`, etc.. This means that Rust's `System` allocator, which
+// calls `libc::malloc()` et al., is actually calling into jemalloc.
 //
 // 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.
 //
-// As for the symbol overrides in `main` below: we're pulling in a static copy
-// of jemalloc. We need to actually reference its symbols for it to get linked.
-// The two crates we link to here, `std` and `rustc_driver`, are both dynamic
-// libraries. So we must reference jemalloc symbols one way or another, because
-// this file is the only object code in the rustc executable.
+// 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:
@@ -36,43 +35,9 @@
 // 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 _;
 
 fn main() {
-    // See the comment at the top of this file for an explanation of this.
-    #[cfg(feature = "jemalloc")]
-    {
-        use std::os::raw::{c_int, c_void};
-
-        use tikv_jemalloc_sys as jemalloc_sys;
-
-        #[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")]
-        {
-            unsafe extern "C" {
-                fn _rjem_je_zone_register();
-            }
-
-            #[used]
-            static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
-        }
-    }
-
     rustc_driver::main()
 }

src/librustdoc/Cargo.toml

@@ -24,6 +24,7 @@ smallvec = "1.8.1"
 stringdex = "=0.0.2"
 tempfile = "3"
 threadpool = "1.8.1"
+tikv-jemalloc-sys = { version = "0.6.1", optional = true, features = ['override_allocator_on_supported_platforms'] }
 tracing = "0.1"
 tracing-tree = "0.3.0"
 unicode-segmentation = "1.9"
@@ -42,7 +43,7 @@ minifier = { version = "0.3.2", default-features = false }
 expect-test = "1.4.0"
 
 [features]
-jemalloc = []
+jemalloc = ["dep:tikv-jemalloc-sys"]
 
 [package.metadata.rust-analyzer]
 rustc_private = true

src/librustdoc/lib.rs

@@ -61,11 +61,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
-// about jemalloc.
-#[cfg(feature = "jemalloc")]
-extern crate tikv_jemalloc_sys as jemalloc_sys;
-
 use std::env::{self, VarError};
 use std::io::{self, IsTerminal};
 use std::path::Path;
@@ -77,6 +72,10 @@ use rustc_interface::interface;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::{ErrorOutputType, RustcOptGroup, make_crate_type_option};
 use rustc_session::{EarlyDiagCtxt, getopts};
+/// 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 `use` statement.
+#[cfg(feature = "jemalloc")]
+use tikv_jemalloc_sys as _;
 use tracing::info;
 
 use crate::clean::utils::DOC_RUST_LANG_ORG_VERSION;
@@ -124,37 +123,6 @@ mod visit_ast;
 mod visit_lib;
 
 pub fn main() {
-    // See docs in https://github.com/rust-lang/rust/blob/HEAD/compiler/rustc/src/main.rs
-    // about jemalloc.
-    #[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;
-
-        #[cfg(target_os = "macos")]
-        {
-            unsafe extern "C" {
-                fn _rjem_je_zone_register();
-            }
-
-            #[used]
-            static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
-        }
-    }
-
     let mut early_dcx = EarlyDiagCtxt::new(ErrorOutputType::default());
 
     rustc_driver::install_ice_hook(

src/tools/clippy/Cargo.toml

@@ -31,6 +31,7 @@ tempfile = { version = "3.20", optional = true }
 termize = "0.2"
 color-print = "0.3.4"
 anstream = "0.6.18"
+tikv-jemalloc-sys = { version = "0.6.1", optional = true, features = ['override_allocator_on_supported_platforms'] }
 
 [dev-dependencies]
 cargo_metadata = "0.18.1"
@@ -56,7 +57,7 @@ rustc_tools_util = { path = "rustc_tools_util", version = "0.4.2" }
 [features]
 integration = ["dep:tempfile"]
 internal = ["dep:clippy_lints_internal", "dep:tempfile"]
-jemalloc = []
+jemalloc = ["dep:tikv-jemalloc-sys"]
 
 [package.metadata.rust-analyzer]
 # This package uses #[feature(rustc_private)]

src/tools/clippy/src/driver.rs

@@ -13,10 +13,10 @@ 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
-// about jemalloc.
+/// 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 `use` statement.
 #[cfg(feature = "jemalloc")]
-extern crate tikv_jemalloc_sys as jemalloc_sys;
+use tikv_jemalloc_sys as _;
 
 use clippy_utils::sym;
 use declare_clippy_lint::LintListBuilder;
@@ -189,36 +189,6 @@ const BUG_REPORT_URL: &str = "https://github.com/rust-lang/rust-clippy/issues/ne
 
 #[expect(clippy::too_many_lines)]
 pub fn main() {
-    // See docs in https://github.com/rust-lang/rust/blob/HEAD/compiler/rustc/src/main.rs
-    // about jemalloc.
-    #[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;
-
-        #[cfg(target_os = "macos")]
-        {
-            unsafe extern "C" {
-                fn _rjem_je_zone_register();
-            }
-
-            #[used]
-            static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
-        }
-    }
-
     let early_dcx = EarlyDiagCtxt::new(ErrorOutputType::default());
 
     rustc_driver::init_rustc_env_logger(&early_dcx);

src/tools/miri/Cargo.toml

@@ -33,8 +33,8 @@ serde_json = { version = "1.0", optional = true }
 # But only for some targets, it fails for others. Rustc configures this in its CI, but we can't
 # easily use that since we support of-tree builds.
 [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies.tikv-jemalloc-sys]
-version = "0.6.0"
-features = ['unprefixed_malloc_on_supported_platforms']
+version = "0.6.1"
+features = ['override_allocator_on_supported_platforms']
 
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"

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

@@ -20,6 +20,11 @@ extern crate rustc_middle;
 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 `use` statement.
+#[cfg(any(target_os = "linux", target_os = "macos"))]
+use tikv_jemalloc_sys as _;
+
 mod log;
 
 use std::env;
@@ -392,48 +397,7 @@ fn parse_range(val: &str) -> Result<Range<u32>, &'static str> {
     Ok(from..to)
 }
 
-#[cfg(any(target_os = "linux", target_os = "macos"))]
-fn jemalloc_magic() {
-    // These magic runes are copied from
-    // <https://github.com/rust-lang/rust/blob/e89bd9428f621545c979c0ec686addc6563a394e/compiler/rustc/src/main.rs#L39>.
-    // See there for further comments.
-    use std::os::raw::{c_int, c_void};
-
-    use tikv_jemalloc_sys as jemalloc_sys;
-
-    #[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")]
-    {
-        unsafe extern "C" {
-            fn _rjem_je_zone_register();
-        }
-
-        #[used]
-        static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
-    }
-}
-
 fn main() {
-    #[cfg(any(target_os = "linux", target_os = "macos"))]
-    jemalloc_magic();
-
     let early_dcx = EarlyDiagCtxt::new(ErrorOutputType::default());
 
     // Snapshot a copy of the environment before `rustc` starts messing with it.