3 files changed, 128 insertions, 27 deletions

diff --git a/crates/ra_prof/src/lib.rs b/crates/ra_prof/src/lib.rs
index 89df7f04b..eb50965ae 100644
--- a/crates/ra_prof/src/lib.rs
+++ b/crates/ra_prof/src/lib.rs
@@ -1,5 +1,6 @@
 //! A collection of tools for profiling rust-analyzer.
 
+mod stop_watch;
 mod memory_usage;
 #[cfg(feature = "cpu_profiler")]
 mod google_cpu_profiler;
@@ -11,14 +12,9 @@ use std::cell::RefCell;
 pub use crate::{
     hprof::{init, init_from, profile},
     memory_usage::{Bytes, MemoryUsage},
+    stop_watch::{StopWatch, StopWatchSpan},
 };
 
-// We use jemalloc mainly to get heap usage statistics, actual performance
-// difference is not measures.
-#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
-#[global_allocator]
-static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
-
 /// Prints backtrace to stderr, useful for debugging.
 #[cfg(feature = "backtrace")]
 pub fn print_backtrace() {
@@ -43,6 +39,7 @@ pub struct Scope {
 }
 
 impl Scope {
+    #[must_use]
     pub fn enter() -> Scope {
         let prev = IN_SCOPE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), true));
         Scope { prev }
@@ -65,7 +62,8 @@ impl Drop for Scope {
 /// 2. Build with `cpu_profiler` feature.
 /// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
 /// 4. Install pprof for visualization (https://github.com/google/pprof).
-/// 5. Use something like `pprof -svg target/release/rust-analyzer ./out.profile` to see the results.
+/// 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.
 ///
 /// For example, here's how I run profiling on NixOS:
 ///
@@ -73,11 +71,16 @@ impl Drop for Scope {
 /// $ nix-shell -p gperftools --run \
 ///     'cargo run --release -p rust-analyzer -- parse < ~/projects/rustbench/parser.rs > /dev/null'
 /// ```
+///
+/// See this diff for how to profile completions:
+///
+/// https://github.com/rust-analyzer/rust-analyzer/pull/5306
 #[derive(Debug)]
 pub struct CpuProfiler {
     _private: (),
 }
 
+#[must_use]
 pub fn cpu_profiler() -> CpuProfiler {
     #[cfg(feature = "cpu_profiler")]
     {
diff --git a/crates/ra_prof/src/memory_usage.rs b/crates/ra_prof/src/memory_usage.rs
index 9768f656c..c2ecbd33c 100644
--- a/crates/ra_prof/src/memory_usage.rs
+++ b/crates/ra_prof/src/memory_usage.rs
@@ -1,46 +1,58 @@
 //! FIXME: write short doc here
-
 use std::fmt;
 
+use cfg_if::cfg_if;
+
+#[derive(Copy, Clone)]
 pub struct MemoryUsage {
     pub allocated: Bytes,
-    pub resident: Bytes,
 }
 
-impl MemoryUsage {
-    #[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
-    pub fn current() -> MemoryUsage {
-        jemalloc_ctl::epoch::advance().unwrap();
-        MemoryUsage {
-            allocated: Bytes(jemalloc_ctl::stats::allocated::read().unwrap()),
-            resident: Bytes(jemalloc_ctl::stats::resident::read().unwrap()),
-        }
+impl fmt::Display for MemoryUsage {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "{}", self.allocated)
     }
+}
 
-    #[cfg(any(not(feature = "jemalloc"), target_env = "msvc"))]
-    pub fn current() -> MemoryUsage {
-        MemoryUsage { allocated: Bytes(0), resident: Bytes(0) }
+impl std::ops::Sub for MemoryUsage {
+    type Output = MemoryUsage;
+    fn sub(self, rhs: MemoryUsage) -> MemoryUsage {
+        MemoryUsage { allocated: self.allocated - rhs.allocated }
     }
 }
 
-impl fmt::Display for MemoryUsage {
-    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        write!(fmt, "{} allocated {} resident", self.allocated, self.resident,)
+impl MemoryUsage {
+    pub fn current() -> MemoryUsage {
+        cfg_if! {
+            if #[cfg(target_os = "linux")] {
+                // Note: This is incredibly slow.
+                let alloc = unsafe { libc::mallinfo() }.uordblks as isize;
+                MemoryUsage { allocated: Bytes(alloc) }
+            } else {
+                MemoryUsage { allocated: Bytes(0) }
+            }
+        }
     }
 }
 
 #[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
-pub struct Bytes(usize);
+pub struct Bytes(isize);
+
+impl Bytes {
+    pub fn megabytes(self) -> isize {
+        self.0 / 1024 / 1024
+    }
+}
 
 impl fmt::Display for Bytes {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let bytes = self.0;
         let mut value = bytes;
         let mut suffix = "b";
-        if value > 4096 {
+        if value.abs() > 4096 {
             value /= 1024;
             suffix = "kb";
-            if value > 4096 {
+            if value.abs() > 4096 {
                 value /= 1024;
                 suffix = "mb";
             }
@@ -51,7 +63,7 @@ impl fmt::Display for Bytes {
 
 impl std::ops::AddAssign<usize> for Bytes {
     fn add_assign(&mut self, x: usize) {
-        self.0 += x;
+        self.0 += x as isize;
     }
 }
 
diff --git a/crates/ra_prof/src/stop_watch.rs b/crates/ra_prof/src/stop_watch.rs
new file mode 100644
index 000000000..5e276190e
--- /dev/null
+++ b/crates/ra_prof/src/stop_watch.rs
@@ -0,0 +1,86 @@
+//! Like `std::time::Instant`, but also measures memory & CPU cycles.
+use std::{
+    fmt,
+    time::{Duration, Instant},
+};
+
+use crate::MemoryUsage;
+
+pub struct StopWatch {
+    time: Instant,
+    #[cfg(target_os = "linux")]
+    counter: Option<perf_event::Counter>,
+    memory: Option<MemoryUsage>,
+}
+
+pub struct StopWatchSpan {
+    pub time: Duration,
+    pub instructions: Option<u64>,
+    pub memory: Option<MemoryUsage>,
+}
+
+impl StopWatch {
+    pub fn start() -> StopWatch {
+        #[cfg(target_os = "linux")]
+        let counter = {
+            let mut counter = perf_event::Builder::new()
+                .build()
+                .map_err(|err| eprintln!("Failed to create perf counter: {}", err))
+                .ok();
+            if let Some(counter) = &mut counter {
+                if let Err(err) = counter.enable() {
+                    eprintln!("Failed to start perf counter: {}", err)
+                }
+            }
+            counter
+        };
+        let time = Instant::now();
+        StopWatch {
+            time,
+            #[cfg(target_os = "linux")]
+            counter,
+            memory: None,
+        }
+    }
+    pub fn memory(mut self, yes: bool) -> StopWatch {
+        if yes {
+            self.memory = Some(MemoryUsage::current());
+        }
+        self
+    }
+    pub fn elapsed(&mut self) -> StopWatchSpan {
+        let time = self.time.elapsed();
+
+        #[cfg(target_os = "linux")]
+        let instructions = self.counter.as_mut().and_then(|it| {
+            it.read().map_err(|err| eprintln!("Failed to read perf counter: {}", err)).ok()
+        });
+        #[cfg(not(target_os = "linux"))]
+        let instructions = None;
+
+        let memory = self.memory.map(|it| MemoryUsage::current() - it);
+        StopWatchSpan { time, instructions, memory }
+    }
+}
+
+impl fmt::Display for StopWatchSpan {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:.2?}", self.time)?;
+        if let Some(mut instructions) = self.instructions {
+            let mut prefix = "";
+            if instructions > 10000 {
+                instructions /= 1000;
+                prefix = "k"
+            }
+            if instructions > 10000 {
+                instructions /= 1000;
+                prefix = "m"
+            }
+            write!(f, ", {}{}i", instructions, prefix)?;
+        }
+        if let Some(memory) = self.memory {
+            write!(f, ", {}", memory)?;
+        }
+        Ok(())
+    }
+}