crates/profile/src/hprof.rs

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