2 files changed, 21 insertions, 4 deletions
diff --git a/compiler/rustc/Cargo.toml b/compiler/rustc/Cargo.toml
index 6e6c0c71a1f..ca6055c46a6 100644
--- a/compiler/rustc/Cargo.toml
+++ b/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']
diff --git a/compiler/rustc/src/main.rs b/compiler/rustc/src/main.rs
index 859028957db..c80fab99496 100644
--- a/compiler/rustc/src/main.rs
+++ b/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};