crates/profile/src/hprof.rs

   1 //! Simple hierarchical profiler
   2 use once_cell::sync::Lazy;
   3 use std::{
   4     cell::RefCell,
   5     collections::{BTreeMap, HashSet},
   6     io::{stderr, Write},
   7     sync::{
   8         atomic::{AtomicBool, Ordering},
   9         RwLock,
  10     },
  11     time::{Duration, Instant},
  12 };
  13
  14 use crate::tree::{Idx, Tree};
  15
  16 /// Filtering syntax
  17 /// env RA_PROFILE=*             // dump everything
  18 /// env RA_PROFILE=foo|bar|baz   // enabled only selected entries
  19 /// env RA_PROFILE=*@3>10        // dump everything, up to depth 3, if it takes more than 10 ms
  20 pub fn init() {
  21     let spec = std::env::var("RA_PROFILE").unwrap_or_default();
  22     init_from(&spec);
  23 }
  24
  25 pub fn init_from(spec: &str) {
  26     let filter = if spec.is_empty() { Filter::disabled() } else { Filter::from_spec(spec) };
  27     filter.install();
  28 }
  29
  30 type Label = &'static str;
  31
  32 /// This function starts a profiling scope in the current execution stack with a given description.
  33 /// It returns a `Profile` struct that measures elapsed time between this method invocation and `Profile` struct drop.
  34 /// It supports nested profiling scopes in case when this function is invoked multiple times at the execution stack.
  35 /// In this case the profiling information will be nested at the output.
  36 /// Profiling information is being printed in the stderr.
  37 ///
  38 /// # Example
  39 /// ```
  40 /// profile::init_from("profile1|profile2@2");
  41 /// profiling_function1();
  42 ///
  43 /// fn profiling_function1() {
  44 ///     let _p = profile::span("profile1");
  45 ///     profiling_function2();
  46 /// }
  47 ///
  48 /// fn profiling_function2() {
  49 ///     let _p = profile::span("profile2");
  50 /// }
  51 /// ```
  52 /// This will print in the stderr the following:
  53 /// ```text
  54 ///  0ms - profile
  55 ///      0ms - profile2
  56 /// ```
  57 pub fn span(label: Label) -> ProfileSpan {
  58     assert!(!label.is_empty());
  59
  60     if PROFILING_ENABLED.load(Ordering::Relaxed)
  61         && PROFILE_STACK.with(|stack| stack.borrow_mut().push(label))
  62     {
  63         ProfileSpan(Some(ProfilerImpl { label, detail: None }))
  64     } else {
  65         ProfileSpan(None)
  66     }
  67 }
  68
  69 pub struct ProfileSpan(Option<ProfilerImpl>);
  70
  71 struct ProfilerImpl {
  72     label: Label,
  73     detail: Option<String>,
  74 }
  75
  76 impl ProfileSpan {
  77     pub fn detail(mut self, detail: impl FnOnce() -> String) -> ProfileSpan {
  78         if let Some(profiler) = &mut self.0 {
  79             profiler.detail = Some(detail())
  80         }
  81         self
  82     }
  83 }
  84
  85 impl Drop for ProfilerImpl {
  86     fn drop(&mut self) {
  87         PROFILE_STACK.with(|it| it.borrow_mut().pop(self.label, self.detail.take()));
  88     }
  89 }
  90
  91 static PROFILING_ENABLED: AtomicBool = AtomicBool::new(false);
  92 static FILTER: Lazy<RwLock<Filter>> = Lazy::new(Default::default);
  93 thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(ProfileStack::new()));
  94
  95 #[derive(Default, Clone, Debug)]
  96 struct Filter {
  97     depth: usize,
  98     allowed: HashSet<String>,
  99     longer_than: Duration,
 100     version: usize,
 101 }
 102
 103 impl Filter {
 104     fn disabled() -> Filter {
 105         Filter::default()
 106     }
 107
 108     fn from_spec(mut spec: &str) -> Filter {
 109         let longer_than = if let Some(idx) = spec.rfind('>') {
 110             let longer_than = spec[idx + 1..].parse().expect("invalid profile longer_than");
 111             spec = &spec[..idx];
 112             Duration::from_millis(longer_than)
 113         } else {
 114             Duration::new(0, 0)
 115         };
 116
 117         let depth = if let Some(idx) = spec.rfind('@') {
 118             let depth: usize = spec[idx + 1..].parse().expect("invalid profile depth");
 119             spec = &spec[..idx];
 120             depth
 121         } else {
 122             999
 123         };
 124         let allowed =
 125             if spec == "*" { HashSet::new() } else { spec.split('|').map(String::from).collect() };
 126         Filter { depth, allowed, longer_than, version: 0 }
 127     }
 128
 129     fn install(mut self) {
 130         PROFILING_ENABLED.store(self.depth > 0, Ordering::SeqCst);
 131         let mut old = FILTER.write().unwrap();
 132         self.version = old.version + 1;
 133         *old = self;
 134     }
 135 }
 136
 137 struct ProfileStack {
 138     starts: Vec<Instant>,
 139     filter: Filter,
 140     messages: Tree<Message>,
 141 }
 142
 143 #[derive(Default)]
 144 struct Message {
 145     duration: Duration,
 146     label: Label,
 147     detail: Option<String>,
 148 }
 149
 150 impl ProfileStack {
 151     fn new() -> ProfileStack {
 152         ProfileStack { starts: Vec::new(), messages: Tree::default(), filter: Default::default() }
 153     }
 154
 155     fn push(&mut self, label: Label) -> bool {
 156         if self.starts.is_empty() {
 157             if let Ok(f) = FILTER.try_read() {
 158                 if f.version > self.filter.version {
 159                     self.filter = f.clone();
 160                 }
 161             };
 162         }
 163         if self.starts.len() > self.filter.depth {
 164             return false;
 165         }
 166         let allowed = &self.filter.allowed;
 167         if self.starts.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
 168             return false;
 169         }
 170
 171         self.starts.push(Instant::now());
 172         self.messages.start();
 173         true
 174     }
 175
 176     fn pop(&mut self, label: Label, detail: Option<String>) {
 177         let start = self.starts.pop().unwrap();
 178         let duration = start.elapsed();
 179         self.messages.finish(Message { duration, label, detail });
 180         if self.starts.is_empty() {
 181             let longer_than = self.filter.longer_than;
 182             // Convert to millis for comparison to avoid problems with rounding
 183             // (otherwise we could print `0ms` despite user's `>0` filter when
 184             // `duration` is just a few nanos).
 185             if duration.as_millis() > longer_than.as_millis() {
 186                 if let Some(root) = self.messages.root() {
 187                     print(&self.messages, root, 0, longer_than, &mut stderr().lock());
 188                 }
 189             }
 190             self.messages.clear();
 191         }
 192     }
 193 }
 194
 195 fn print(
 196     tree: &Tree<Message>,
 197     curr: Idx<Message>,
 198     level: u32,
 199     longer_than: Duration,
 200     out: &mut impl Write,
 201 ) {
 202     let current_indent = "    ".repeat(level as usize);
 203     let detail = tree[curr].detail.as_ref().map(|it| format!(" @ {}", it)).unwrap_or_default();
 204     writeln!(
 205         out,
 206         "{}{:5}ms - {}{}",
 207         current_indent,
 208         tree[curr].duration.as_millis(),
 209         tree[curr].label,
 210         detail,
 211     )
 212     .expect("printing profiling info");
 213
 214     let mut accounted_for = Duration::default();
 215     let mut short_children = BTreeMap::new(); // Use `BTreeMap` to get deterministic output.
 216     for child in tree.children(curr) {
 217         accounted_for += tree[child].duration;
 218
 219         if tree[child].duration.as_millis() > longer_than.as_millis() {
 220             print(tree, child, level + 1, longer_than, out)
 221         } else {
 222             let (total_duration, cnt) =
 223                 short_children.entry(tree[child].label).or_insert((Duration::default(), 0));
 224             *total_duration += tree[child].duration;
 225             *cnt += 1;
 226         }
 227     }
 228
 229     for (child_msg, (duration, count)) in short_children.iter() {
 230         let millis = duration.as_millis();
 231         writeln!(out, "    {}{:5}ms - {} ({} calls)", current_indent, millis, child_msg, count)
 232             .expect("printing profiling info");
 233     }
 234
 235     let unaccounted = tree[curr].duration - accounted_for;
 236     if tree.children(curr).next().is_some() && unaccounted > longer_than {
 237         writeln!(out, "    {}{:5}ms - ???", current_indent, unaccounted.as_millis())
 238             .expect("printing profiling info");
 239     }
 240 }