1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
use std::cell::RefCell; use std::env; use std::fmt; use std::io::{stdout, StdoutLock, Write}; use std::iter::repeat; use std::mem; use std::time; thread_local!(static PROFILE_STACK: RefCell<Vec<time::Instant>> = RefCell::new(Vec::new())); thread_local!(static MESSAGES: RefCell<Vec<Message>> = RefCell::new(Vec::new())); type Message = (usize, u64, String); pub struct Profiler { desc: String, } fn enabled_level() -> Option<usize> { env::var("CARGO_PROFILE").ok().and_then(|s| s.parse().ok()) } pub fn start<T: fmt::Display>(desc: T) -> Profiler { if enabled_level().is_none() { return Profiler { desc: String::new(), }; } PROFILE_STACK.with(|stack| stack.borrow_mut().push(time::Instant::now())); Profiler { desc: desc.to_string(), } } impl Drop for Profiler { fn drop(&mut self) { let enabled = match enabled_level() { Some(i) => i, None => return, }; let (start, stack_len) = PROFILE_STACK.with(|stack| { let mut stack = stack.borrow_mut(); let start = stack.pop().unwrap(); (start, stack.len()) }); let duration = start.elapsed(); let duration_ms = duration.as_secs() * 1000 + u64::from(duration.subsec_millis()); let msg = ( stack_len, duration_ms, mem::replace(&mut self.desc, String::new()), ); MESSAGES.with(|msgs| msgs.borrow_mut().push(msg)); if stack_len == 0 { fn print(lvl: usize, msgs: &[Message], enabled: usize, stdout: &mut StdoutLock<'_>) { if lvl > enabled { return; } let mut last = 0; for (i, &(l, time, ref msg)) in msgs.iter().enumerate() { if l != lvl { continue; } writeln!( stdout, "{} {:6}ms - {}", repeat(" ").take(lvl + 1).collect::<String>(), time, msg ) .expect("printing profiling info to stdout"); print(lvl + 1, &msgs[last..i], enabled, stdout); last = i; } } let stdout = stdout(); MESSAGES.with(|msgs| { let mut msgs = msgs.borrow_mut(); print(0, &msgs, enabled, &mut stdout.lock()); msgs.clear(); }); } } }