3 files changed, 54 insertions, 6 deletions

diff --git a/crates/profile/Cargo.toml b/crates/profile/Cargo.toml
index 1a8c8f862..653d3d983 100644
--- a/crates/profile/Cargo.toml
+++ b/crates/profile/Cargo.toml
@@ -20,6 +20,9 @@ jemalloc-ctl = { version = "0.4.1", package = "tikv-jemalloc-ctl", optional = tr
 [target.'cfg(target_os = "linux")'.dependencies]
 perf-event = "0.4"
 
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3.8", features = ["psapi"] }
+
 [features]
 cpu_profiler = []
 jemalloc = ["jemalloc-ctl"]
diff --git a/crates/profile/src/lib.rs b/crates/profile/src/lib.rs
index 7cd01a7df..5ea5039db 100644
--- a/crates/profile/src/lib.rs
+++ b/crates/profile/src/lib.rs
@@ -50,10 +50,10 @@ impl Drop for Scope {
 /// A wrapper around google_cpu_profiler.
 ///
 /// Usage:
-/// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
+/// 1. Install gpref_tools (<https://github.com/gperftools/gperftools>), probably packaged with your Linux distro.
 /// 2. Build with `cpu_profiler` feature.
 /// 3. Run the code, the *raw* output would be in the `./out.profile` file.
-/// 4. Install pprof for visualization (https://github.com/google/pprof).
+/// 4. Install pprof for visualization (<https://github.com/google/pprof>).
 /// 5. Bump sampling frequency to once per ms: `export CPUPROFILE_FREQUENCY=1000`
 /// 6. Use something like `pprof -svg target/release/rust-analyzer ./out.profile` to see the results.
 ///
@@ -75,7 +75,7 @@ impl Drop for Scope {
 ///
 /// See this diff for how to profile completions:
 ///
-/// https://github.com/rust-analyzer/rust-analyzer/pull/5306
+/// <https://github.com/rust-analyzer/rust-analyzer/pull/5306>
 #[derive(Debug)]
 pub struct CpuSpan {
     _private: (),
diff --git a/crates/profile/src/memory_usage.rs b/crates/profile/src/memory_usage.rs
index 2917ded60..fbcb9e3c2 100644
--- a/crates/profile/src/memory_usage.rs
+++ b/crates/profile/src/memory_usage.rs
@@ -32,9 +32,23 @@ impl 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) }
+                memusage_linux()
+            } else if #[cfg(windows)] {
+                // There doesn't seem to be an API for determining heap usage, so we try to
+                // approximate that by using the Commit Charge value.
+
+                use winapi::um::processthreadsapi::*;
+                use winapi::um::psapi::*;
+                use std::mem::{MaybeUninit, size_of};
+
+                let proc = unsafe { GetCurrentProcess() };
+                let mut mem_counters = MaybeUninit::uninit();
+                let cb = size_of::<PROCESS_MEMORY_COUNTERS>();
+                let ret = unsafe { GetProcessMemoryInfo(proc, mem_counters.as_mut_ptr(), cb as u32) };
+                assert!(ret != 0);
+
+                let usage = unsafe { mem_counters.assume_init().PagefileUsage };
+                MemoryUsage { allocated: Bytes(usage as isize) }
             } else {
                 MemoryUsage { allocated: Bytes(0) }
             }
@@ -42,6 +56,37 @@ impl MemoryUsage {
     }
 }
 
+#[cfg(all(target_os = "linux", target_env = "gnu", not(feature = "jemalloc")))]
+fn memusage_linux() -> MemoryUsage {
+    // Linux/glibc has 2 APIs for allocator introspection that we can use: mallinfo and mallinfo2.
+    // mallinfo uses `int` fields and cannot handle memory usage exceeding 2 GB.
+    // mallinfo2 is very recent, so its presence needs to be detected at runtime.
+    // Both are abysmally slow.
+
+    use std::ffi::CStr;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    static MALLINFO2: AtomicUsize = AtomicUsize::new(1);
+
+    let mut mallinfo2 = MALLINFO2.load(Ordering::Relaxed);
+    if mallinfo2 == 1 {
+        let cstr = CStr::from_bytes_with_nul(b"mallinfo2\0").unwrap();
+        mallinfo2 = unsafe { libc::dlsym(libc::RTLD_DEFAULT, cstr.as_ptr()) } as usize;
+        // NB: races don't matter here, since they'll always store the same value
+        MALLINFO2.store(mallinfo2, Ordering::Relaxed);
+    }
+
+    if mallinfo2 == 0 {
+        // mallinfo2 does not exist, use mallinfo.
+        let alloc = unsafe { libc::mallinfo() }.uordblks as isize;
+        MemoryUsage { allocated: Bytes(alloc) }
+    } else {
+        let mallinfo2: fn() -> libc::mallinfo2 = unsafe { std::mem::transmute(mallinfo2) };
+        let alloc = mallinfo2().uordblks as isize;
+        MemoryUsage { allocated: Bytes(alloc) }
+    }
+}
+
 #[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
 pub struct Bytes(isize);