3 files changed, 135 insertions, 30 deletions

diff --git a/crates/profile/src/google_cpu_profiler.rs b/crates/profile/src/google_cpu_profiler.rs
index db865c65b..cae6caeaa 100644
--- a/crates/profile/src/google_cpu_profiler.rs
+++ b/crates/profile/src/google_cpu_profiler.rs
@@ -14,26 +14,31 @@ extern "C" {
     fn ProfilerStop();
 }
 
-static PROFILER_STATE: AtomicUsize = AtomicUsize::new(OFF);
 const OFF: usize = 0;
 const ON: usize = 1;
 const PENDING: usize = 2;
 
-pub fn start(path: &Path) {
-    if PROFILER_STATE.compare_and_swap(OFF, PENDING, Ordering::SeqCst) != OFF {
+fn transition(current: usize, new: usize) -> bool {
+    static STATE: AtomicUsize = AtomicUsize::new(OFF);
+
+    STATE.compare_exchange(current, new, Ordering::SeqCst, Ordering::SeqCst).is_ok()
+}
+
+pub(crate) fn start(path: &Path) {
+    if !transition(OFF, PENDING) {
         panic!("profiler already started");
     }
     let path = CString::new(path.display().to_string()).unwrap();
     if unsafe { ProfilerStart(path.as_ptr()) } == 0 {
         panic!("profiler failed to start")
     }
-    assert!(PROFILER_STATE.compare_and_swap(PENDING, ON, Ordering::SeqCst) == PENDING);
+    assert!(transition(PENDING, ON));
 }
 
-pub fn stop() {
-    if PROFILER_STATE.compare_and_swap(ON, PENDING, Ordering::SeqCst) != ON {
+pub(crate) fn stop() {
+    if !transition(ON, PENDING) {
         panic!("profiler is not started")
     }
     unsafe { ProfilerStop() };
-    assert!(PROFILER_STATE.compare_and_swap(PENDING, OFF, Ordering::SeqCst) == PENDING);
+    assert!(transition(PENDING, OFF));
 }
diff --git a/crates/profile/src/hprof.rs b/crates/profile/src/hprof.rs
index 29d2ed518..5fdb37206 100644
--- a/crates/profile/src/hprof.rs
+++ b/crates/profile/src/hprof.rs
@@ -1,5 +1,4 @@
 //! Simple hierarchical profiler
-use once_cell::sync::Lazy;
 use std::{
     cell::RefCell,
     collections::{BTreeMap, HashSet},
@@ -12,6 +11,8 @@ use std::{
     time::{Duration, Instant},
 };
 
+use once_cell::sync::Lazy;
+
 use crate::tree::{Idx, Tree};
 
 /// Filtering syntax
@@ -56,18 +57,32 @@ type Label = &'static str;
 ///  0ms - profile
 ///      0ms - profile2
 /// ```
+#[inline]
 pub fn span(label: Label) -> ProfileSpan {
-    assert!(!label.is_empty());
+    debug_assert!(!label.is_empty());
 
-    if PROFILING_ENABLED.load(Ordering::Relaxed)
-        && PROFILE_STACK.with(|stack| stack.borrow_mut().push(label))
-    {
+    let enabled = PROFILING_ENABLED.load(Ordering::Relaxed);
+    if enabled && with_profile_stack(|stack| stack.push(label)) {
         ProfileSpan(Some(ProfilerImpl { label, detail: None }))
     } else {
         ProfileSpan(None)
     }
 }
 
+#[inline]
+pub fn heartbeat_span() -> HeartbeatSpan {
+    let enabled = PROFILING_ENABLED.load(Ordering::Relaxed);
+    HeartbeatSpan::new(enabled)
+}
+
+#[inline]
+pub fn heartbeat() {
+    let enabled = PROFILING_ENABLED.load(Ordering::Relaxed);
+    if enabled {
+        with_profile_stack(|it| it.heartbeat(1));
+    }
+}
+
 pub struct ProfileSpan(Option<ProfilerImpl>);
 
 struct ProfilerImpl {
@@ -85,20 +100,48 @@ impl ProfileSpan {
 }
 
 impl Drop for ProfilerImpl {
+    #[inline]
+    fn drop(&mut self) {
+        with_profile_stack(|it| it.pop(self.label, self.detail.take()));
+    }
+}
+
+pub struct HeartbeatSpan {
+    enabled: bool,
+}
+
+impl HeartbeatSpan {
+    #[inline]
+    pub fn new(enabled: bool) -> Self {
+        if enabled {
+            with_profile_stack(|it| it.heartbeats(true))
+        }
+        Self { enabled }
+    }
+}
+
+impl Drop for HeartbeatSpan {
     fn drop(&mut self) {
-        PROFILE_STACK.with(|it| it.borrow_mut().pop(self.label, self.detail.take()));
+        if self.enabled {
+            with_profile_stack(|it| it.heartbeats(false))
+        }
     }
 }
 
 static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);
 static FILTER: Lazy<RwLock<Filter>> = Lazy::new(Default::default);
-thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));
+
+fn with_profile_stack<T>(f: impl FnOnce(&mut ProfileStack) -> T) -> T {
+    thread_local!(static STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));
+    STACK.with(|it| f(&mut *it.borrow_mut()))
+}
 
 #[derive(Default, Clone, Debug)]
 struct Filter {
     depth: usize,
     allowed: HashSet<String>,
     longer_than: Duration,
+    heartbeat_longer_than: Duration,
     version: usize,
 }
 
@@ -115,6 +158,7 @@ impl Filter {
         } else {
             Duration::new(0, 0)
         };
+        let heartbeat_longer_than = longer_than;
 
         let depth = if let Some(idx) = spec.rfind('@') {
             let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
@@ -125,7 +169,7 @@ impl Filter {
         };
         let allowed =
             if spec == "*" { HashSet::new() } else { spec.split('|').map(String::from).collect() };
-        Filter { depth, allowed, longer_than, version: 0 }
+        Filter { depth, allowed, longer_than, heartbeat_longer_than, version: 0 }
     }
 
     fn install(mut self) {
@@ -137,9 +181,15 @@ impl Filter {
 }
 
 struct ProfileStack {
-    starts: Vec<Instant>,
+    frames: Vec<Frame>,
     filter: Filter,
     messages: Tree<Message>,
+    heartbeats: bool,
+}
+
+struct Frame {
+    t: Instant,
+    heartbeats: u32,
 }
 
 #[derive(Default)]
@@ -151,35 +201,49 @@ struct Message {
 
 impl ProfileStack {
     fn new() -> ProfileStack {
-        ProfileStack { starts: Vec::new(), messages: Tree::default(), filter: Default::default() }
+        ProfileStack {
+            frames: Vec::new(),
+            messages: Tree::default(),
+            filter: Default::default(),
+            heartbeats: false,
+        }
     }
 
     fn push(&mut self, label: Label) -> bool {
-        if self.starts.is_empty() {
+        if self.frames.is_empty() {
             if let Ok(f) = FILTER.try_read() {
                 if f.version > self.filter.version {
                     self.filter = f.clone();
                 }
             };
         }
-        if self.starts.len() > self.filter.depth {
+        if self.frames.len() > self.filter.depth {
             return false;
         }
         let allowed = &self.filter.allowed;
-        if self.starts.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
+        if self.frames.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
             return false;
         }
 
-        self.starts.push(Instant::now());
+        self.frames.push(Frame { t: Instant::now(), heartbeats: 0 });
         self.messages.start();
         true
     }
 
     fn pop(&mut self, label: Label, detail: Option<String>) {
-        let start = self.starts.pop().unwrap();
-        let duration = start.elapsed();
+        let frame = self.frames.pop().unwrap();
+        let duration = frame.t.elapsed();
+
+        if self.heartbeats {
+            self.heartbeat(frame.heartbeats);
+            let avg_span = duration / (frame.heartbeats + 1);
+            if avg_span > self.filter.heartbeat_longer_than {
+                eprintln!("Too few heartbeats {} ({}/{:?})?", label, frame.heartbeats, duration)
+            }
+        }
+
         self.messages.finish(Message { duration, label, detail });
-        if self.starts.is_empty() {
+        if self.frames.is_empty() {
             let longer_than = self.filter.longer_than;
             // Convert to millis for comparison to avoid problems with rounding
             // (otherwise we could print `0ms` despite user's `>0` filter when
@@ -192,6 +256,15 @@ impl ProfileStack {
             self.messages.clear();
         }
     }
+
+    fn heartbeats(&mut self, yes: bool) {
+        self.heartbeats = yes;
+    }
+    fn heartbeat(&mut self, n: u32) {
+        if let Some(frame) = self.frames.last_mut() {
+            frame.heartbeats += n;
+        }
+    }
 }
 
 fn print(
diff --git a/crates/profile/src/lib.rs b/crates/profile/src/lib.rs
index 79dba47d5..a31fb8f43 100644
--- a/crates/profile/src/lib.rs
+++ b/crates/profile/src/lib.rs
@@ -10,7 +10,7 @@ mod tree;
 use std::cell::RefCell;
 
 pub use crate::{
-    hprof::{init, init_from, span},
+    hprof::{heartbeat, heartbeat_span, init, init_from, span},
     memory_usage::{Bytes, MemoryUsage},
     stop_watch::{StopWatch, StopWatchSpan},
 };
@@ -52,7 +52,7 @@ impl Drop for Scope {
 /// Usage:
 /// 1. Install gpref_tools (https://github.com/gperftools/gperftools), probably packaged with your Linux distro.
 /// 2. Build with `cpu_profiler` feature.
-/// 3. Tun the code, the *raw* output would be in the `./out.profile` file.
+/// 3. Run the code, the *raw* output would be in the `./out.profile` file.
 /// 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.
@@ -60,8 +60,17 @@ impl Drop for Scope {
 /// For example, here's how I run profiling on NixOS:
 ///
 /// ```bash
-/// $ nix-shell -p gperftools --run \
-///     'cargo run --release -p rust-analyzer -- parse < ~/projects/rustbench/parser.rs > /dev/null'
+/// $ bat -p shell.nix
+/// with import <nixpkgs> {};
+/// mkShell {
+///   buildInputs = [ gperftools ];
+///   shellHook = ''
+///     export LD_LIBRARY_PATH="${gperftools}/lib:"
+///   '';
+/// }
+/// $ set -x CPUPROFILE_FREQUENCY 1000
+/// $ nix-shell --run 'cargo test --release --package rust-analyzer --lib -- benchmarks::benchmark_integrated_highlighting --exact --nocapture'
+/// $ pprof -svg target/release/deps/rust_analyzer-8739592dc93d63cb crates/rust-analyzer/out.profile > profile.svg
 /// ```
 ///
 /// See this diff for how to profile completions:
@@ -81,7 +90,9 @@ pub fn cpu_span() -> CpuSpan {
 
     #[cfg(not(feature = "cpu_profiler"))]
     {
-        eprintln!("cpu_profiler feature is disabled")
+        eprintln!(
+            r#"cpu profiling is disabled, uncomment `default = [ "cpu_profiler" ]` in Cargo.toml to enable."#
+        )
     }
 
     CpuSpan { _private: () }
@@ -91,7 +102,23 @@ impl Drop for CpuSpan {
     fn drop(&mut self) {
         #[cfg(feature = "cpu_profiler")]
         {
-            google_cpu_profiler::stop()
+            google_cpu_profiler::stop();
+            let profile_data = std::env::current_dir().unwrap().join("out.profile");
+            eprintln!("Profile data saved to:\n\n    {}\n", profile_data.display());
+            let mut cmd = std::process::Command::new("pprof");
+            cmd.arg("-svg").arg(std::env::current_exe().unwrap()).arg(&profile_data);
+            let out = cmd.output();
+
+            match out {
+                Ok(out) if out.status.success() => {
+                    let svg = profile_data.with_extension("svg");
+                    std::fs::write(&svg, &out.stdout).unwrap();
+                    eprintln!("Profile rendered to:\n\n    {}\n", svg.display());
+                }
+                _ => {
+                    eprintln!("Failed to run:\n\n   {:?}\n", cmd);
+                }
+            }
         }
     }
 }