Add back jemalloc support

7 files changed, 94 insertions, 4 deletions

diff --git a/Cargo.lock b/Cargo.lock
index b6b6f0341..35713a0c4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -449,6 +449,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "fs_extra"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2022715d62ab30faffd124d40b76f4134a550a87792276512b18d63272333394"
+
+[[package]]
 name = "fsevent"
 version = "2.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -758,6 +764,38 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dd25036021b0de88a0aff6b850051563c6516d0bf53f8638938edbb9de732736"
 
 [[package]]
+name = "jemalloc-ctl"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c502a5ff9dd2924f1ed32ba96e3b65735d837b4bfd978d3161b1702e66aca4b7"
+dependencies = [
+ "jemalloc-sys",
+ "libc",
+ "paste",
+]
+
+[[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 = "jemallocator"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43ae63fcfc45e99ab3d1b29a46782ad679e98436c3169d15a167a1108a724b69"
+dependencies = [
+ "jemalloc-sys",
+ "libc",
+]
+
+[[package]]
 name = "jod-thread"
 version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1105,6 +1143,25 @@ dependencies = [
 ]
 
 [[package]]
+name = "paste"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "45ca20c77d80be666aef2b45486da86238fabe33e38306bd3118fe4af33fa880"
+dependencies = [
+ "paste-impl",
+ "proc-macro-hack",
+]
+
+[[package]]
+name = "paste-impl"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d95a7db200b97ef370c8e6de0088252f7e0dfff7d047a28528e47456c0fc98b6"
+dependencies = [
+ "proc-macro-hack",
+]
+
+[[package]]
 name = "paths"
 version = "0.0.0"
 
@@ -1165,6 +1222,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "439697af366c49a6d0a010c56a0d97685bc140ce0d377b13a2ea2aa42d64a827"
 
 [[package]]
+name = "proc-macro-hack"
+version = "0.5.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dbf0c48bc1d91375ae5c3cd81e3722dff1abcf81a30960240640d223f59fe0e5"
+
+[[package]]
 name = "proc-macro2"
 version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1212,6 +1275,7 @@ name = "profile"
 version = "0.0.0"
 dependencies = [
  "cfg-if 1.0.0",
+ "jemalloc-ctl",
  "la-arena",
  "libc",
  "once_cell",
@@ -1354,6 +1418,7 @@ dependencies = [
  "ide",
  "ide_db",
  "itertools 0.10.0",
+ "jemallocator",
  "jod-thread",
  "log",
  "lsp-server",
diff --git a/crates/profile/Cargo.toml b/crates/profile/Cargo.toml
index 9c7ce26be..f7231c2b8 100644
--- a/crates/profile/Cargo.toml
+++ b/crates/profile/Cargo.toml
@@ -14,12 +14,14 @@ once_cell = "1.3.1"
 cfg-if = "1"
 libc = "0.2.73"
 la-arena = { version = "0.2.0", path = "../../lib/arena" }
+jemalloc-ctl = { version = "0.3.3", optional = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 perf-event = "0.4"
 
 [features]
 cpu_profiler = []
+jemalloc = ["jemalloc-ctl"]
 
 # Uncomment to enable for the whole crate graph
 # default = [ "cpu_profiler" ]
diff --git a/crates/profile/src/memory_usage.rs b/crates/profile/src/memory_usage.rs
index 83390212a..cb4e54447 100644
--- a/crates/profile/src/memory_usage.rs
+++ b/crates/profile/src/memory_usage.rs
@@ -24,7 +24,12 @@ impl std::ops::Sub for MemoryUsage {
 impl MemoryUsage {
     pub fn current() -> MemoryUsage {
         cfg_if! {
-            if #[cfg(all(target_os = "linux", target_env = "gnu"))] {
+            if #[cfg(all(feature = "jemalloc", not(target_env = "msvc")))] {
+                jemalloc_ctl::epoch::advance().unwrap();
+                MemoryUsage {
+                    allocated: Bytes(jemalloc_ctl::stats::allocated::read().unwrap() as isize),
+                }
+            } else if #[cfg(all(target_os = "linux", target_env = "gnu"))] {
                 // Note: This is incredibly slow.
                 let alloc = unsafe { libc::mallinfo() }.uordblks as isize;
                 MemoryUsage { allocated: Bytes(alloc) }
diff --git a/crates/rust-analyzer/Cargo.toml b/crates/rust-analyzer/Cargo.toml
index af7b86ead..3cb45b030 100644
--- a/crates/rust-analyzer/Cargo.toml
+++ b/crates/rust-analyzer/Cargo.toml
@@ -61,8 +61,14 @@ proc_macro_srv = { path = "../proc_macro_srv", version = "0.0.0" }
 [target.'cfg(windows)'.dependencies]
 winapi = "0.3.8"
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+jemallocator = { version = "0.3.2", optional = true }
+
 [dev-dependencies]
 expect-test = "1.1"
 test_utils = { path = "../test_utils" }
 mbe = { path = "../mbe" }
 tt = { path = "../tt" }
+
+[features]
+jemalloc = ["jemallocator", "profile/jemalloc"]
diff --git a/crates/rust-analyzer/src/bin/main.rs b/crates/rust-analyzer/src/bin/main.rs
index bf42654a8..2f7f94a39 100644
--- a/crates/rust-analyzer/src/bin/main.rs
+++ b/crates/rust-analyzer/src/bin/main.rs
@@ -15,6 +15,10 @@ use vfs::AbsPathBuf;
 #[global_allocator]
 static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
+#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
 fn main() {
     if let Err(err) = try_main() {
         eprintln!("{}", err);
diff --git a/xtask/src/install.rs b/xtask/src/install.rs
index 12962bcfa..202c74426 100644
--- a/xtask/src/install.rs
+++ b/xtask/src/install.rs
@@ -67,6 +67,7 @@ pub struct ServerOpt {
 pub enum Malloc {
     System,
     Mimalloc,
+    Jemalloc,
 }
 
 impl InstallCmd {
@@ -176,6 +177,7 @@ fn install_server(opts: ServerOpt) -> Result<()> {
     let features = match opts.malloc {
         Malloc::System => &[][..],
         Malloc::Mimalloc => &["--features", "mimalloc"],
+        Malloc::Jemalloc => &["--features", "jemalloc"],
     };
 
     let cmd = cmd!("cargo install --path crates/rust-analyzer --locked --force {features...}");
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index dec48629c..c3e5c7326 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -49,7 +49,8 @@ FLAGS:
         --client[=CLIENT] Install only VS Code plugin.
                           CLIENT is one of 'code', 'code-exploration', 'code-insiders', 'codium', or 'code-oss'
         --server          Install only the language server
-        --mimalloc        Use mimalloc for server
+        --mimalloc        Use mimalloc allocator for server
+        --jemalloc        Use jemalloc allocator for server
     -h, --help            Prints help information
         "
                 );
@@ -65,8 +66,13 @@ FLAGS:
                 return Ok(());
             }
 
-            let malloc =
-                if args.contains("--mimalloc") { Malloc::Mimalloc } else { Malloc::System };
+            let malloc = if args.contains("--mimalloc") {
+                Malloc::Mimalloc
+            } else if args.contains("--jemalloc") {
+                Malloc::Jemalloc
+            } else {
+                Malloc::System
+            };
 
             let client_opt = args.opt_value_from_str("--client")?;