[TEST] Use global Rust jemalloc

Cargo.lock

@@ -756,7 +756,7 @@ checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681"
 dependencies = [
  "serde",
  "termcolor",
- "unicode-width 0.1.14",
+ "unicode-width 0.2.2",
 ]
 
 [[package]]
@@ -3837,6 +3837,8 @@ dependencies = [
  "rustc_target",
  "serde_json",
  "shlex",
+ "tikv-jemalloc-sys",
+ "tikv-jemallocator",
  "tracing",
  "windows 0.61.3",
 ]
@@ -5551,6 +5553,16 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "tikv-jemallocator"
+version = "0.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a"
+dependencies = [
+ "libc",
+ "tikv-jemalloc-sys",
+]
+
 [[package]]
 name = "tinystr"
 version = "0.8.2"

compiler/rustc/Cargo.toml

@@ -30,12 +30,11 @@ wasi = "=0.14.2"
 [dependencies.tikv-jemalloc-sys]
 version = "0.6.1"
 optional = true
-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_enzyme = ['rustc_driver_impl/llvm_enzyme']
 llvm_offload = ['rustc_driver_impl/llvm_offload']

compiler/rustc/src/main.rs

@@ -37,8 +37,68 @@ 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 _;
+mod c_alloc {
+    use std::os::raw::{c_int, c_void};
+    #[used]
+    static _F1: unsafe extern "C" fn(usize, usize) -> *mut c_void = calloc;
+    #[used]
+    static _F2: unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int = posix_memalign;
+    #[used]
+    static _F3: unsafe extern "C" fn(usize, usize) -> *mut c_void = aligned_alloc;
+    #[used]
+    static _F4: unsafe extern "C" fn(usize) -> *mut c_void = malloc;
+    #[used]
+    static _F5: unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void = realloc;
+    #[used]
+    static _F6: unsafe extern "C" fn(*mut c_void) = 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]
+    #[cfg(target_os = "macos")]
+    static _F7: unsafe extern "C" fn() = _rjem_je_zone_register;
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn calloc(items: usize, size: usize) -> *mut c_void {
+        unsafe { rustc_driver::jemalloc::calloc(items, size) }
+    }
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn posix_memalign(ptr: *mut *mut c_void, size: usize, align: usize) -> c_int {
+        unsafe { rustc_driver::jemalloc::posix_memalign(ptr, size, align) }
+    }
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn aligned_alloc(size: usize, align: usize) -> *mut c_void {
+        unsafe { rustc_driver::jemalloc::aligned_alloc(size, align) }
+    }
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
+        unsafe { rustc_driver::jemalloc::malloc(size) }
+    }
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
+        unsafe { rustc_driver::jemalloc::realloc(ptr, size) }
+    }
+
+    #[unsafe(no_mangle)]
+    unsafe extern "C" fn free(ptr: *mut c_void) {
+        unsafe {
+            rustc_driver::jemalloc::free(ptr);
+        }
+    }
+}
 
 fn main() -> ExitCode {
     rustc_driver::main()

compiler/rustc_driver_impl/Cargo.toml

@@ -55,9 +55,18 @@ features = [
 ctrlc = "3.4.4"
 # tidy-alphabetical-end
 
+[dependencies.tikv-jemallocator]
+version = "0.6.1"
+optional = true
+
+[dependencies.tikv-jemalloc-sys]
+version = "0.6.1"
+optional = true
+
 [features]
 # tidy-alphabetical-start
 check_only = ['rustc_interface/check_only']
+jemalloc = ['dep:tikv-jemallocator', 'dep:tikv-jemalloc-sys']
 llvm = ['rustc_interface/llvm']
 llvm_enzyme = ['rustc_interface/llvm_enzyme']
 llvm_offload = ['rustc_interface/llvm_offload']

compiler/rustc_driver_impl/src/lib.rs

@@ -63,6 +63,44 @@ use rustc_target::json::ToJson;
 use rustc_target::spec::{Target, TargetTuple};
 use tracing::trace;
 
+#[cfg(feature = "jemalloc")]
+pub mod jemalloc {
+    use std::os::raw::{c_int, c_void};
+
+    #[global_allocator]
+    static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+    pub unsafe extern "C" fn calloc(items: usize, size: usize) -> *mut c_void {
+        unsafe { tikv_jemalloc_sys::calloc(items, size) }
+    }
+
+    pub unsafe extern "C" fn posix_memalign(
+        ptr: *mut *mut c_void,
+        size: usize,
+        align: usize,
+    ) -> c_int {
+        unsafe { tikv_jemalloc_sys::posix_memalign(ptr, size, align) }
+    }
+
+    pub unsafe extern "C" fn aligned_alloc(size: usize, align: usize) -> *mut c_void {
+        unsafe { tikv_jemalloc_sys::aligned_alloc(size, align) }
+    }
+
+    pub unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
+        unsafe { tikv_jemalloc_sys::malloc(size) }
+    }
+
+    pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
+        unsafe { tikv_jemalloc_sys::realloc(ptr, size) }
+    }
+
+    pub unsafe extern "C" fn free(ptr: *mut c_void) {
+        unsafe {
+            tikv_jemalloc_sys::free(ptr);
+        }
+    }
+}
+
 #[allow(unused_macros)]
 macro do_not_use_print($($t:tt)*) {
     std::compile_error!(

src/tools/tidy/src/deps.rs

@@ -442,6 +442,7 @@ const PERMITTED_RUSTC_DEPENDENCIES: &[&str] = &[
     "thorin-dwp",
     "thread_local",
     "tikv-jemalloc-sys",
+    "tikv-jemallocator",
     "tinystr",
     "tinyvec",
     "tinyvec_macros",