1extern crate time;
2#[macro_use]
3extern crate serde_json;
4#[macro_use]
5extern crate lazy_static;
6
7use std::cell::RefCell;
8use std::fs::File;
9use std::io::BufWriter;
10use std::string::String;
11use std::sync::Mutex;
12use std::sync::mpsc::{channel, Receiver, Sender};
13use std::thread;
14use time::precise_time_ns;
15
16lazy_static! {
17 static ref GLOBAL_PROFILER: Mutex<Profiler> = Mutex::new(Profiler::new());
18}
19
20#[macro_export]
21macro_rules! profile_scope {
22 ($string:expr) => {
23 #[cfg(feature = "thread_profiler")]
24 let _profile_scope =
25 $crate::ProfileScope::new(format!("{}: {}", module_path!(), $string));
26 };
27}
28
29thread_local!(static THREAD_PROFILER: RefCell<Option<ThreadProfiler>> = RefCell::new(None));
30
31#[derive(Copy, Clone)]
32struct ThreadId(usize);
33
34struct ThreadInfo {
35 name: String,
36}
37
38struct Sample {
39 tid: ThreadId,
40 name: String,
41 t0: u64,
42 t1: u64,
43}
44
45struct ThreadProfiler {
46 id: ThreadId,
47 tx: Sender<Sample>,
48}
49
50impl ThreadProfiler {
51 fn push_sample(&self, name: String, t0: u64, t1: u64) {
52 let sample = Sample {
53 tid: self.id,
54 name: name,
55 t0: t0,
56 t1: t1,
57 };
58 self.tx.send(sample).ok();
59 }
60}
61
62struct Profiler {
63 rx: Receiver<Sample>,
64 tx: Sender<Sample>,
65 threads: Vec<ThreadInfo>,
66}
67
68impl Profiler {
69 fn new() -> Profiler {
70 let (tx, rx) = channel();
71
72 Profiler {
73 rx: rx,
74 tx: tx,
75 threads: Vec::new(),
76 }
77 }
78
79 fn register_thread(&mut self) {
80 let id = ThreadId(self.threads.len());
81 let name = match thread::current().name() {
82 Some(s) => s.to_string(),
83 None => format!("<unnamed-{}>", id.0),
84 };
85
86 self.threads.push(ThreadInfo { name });
87
88 THREAD_PROFILER.with(|profiler| {
89 assert!(profiler.borrow().is_none());
90
91 let thread_profiler = ThreadProfiler {
92 id: id,
93 tx: self.tx.clone(),
94 };
95
96 *profiler.borrow_mut() = Some(thread_profiler);
97 });
98 }
99
100 fn write_profile(&self, filename: &str) {
101 let start_time = precise_time_ns();
104 let mut data = Vec::new();
105
106 while let Ok(sample) = self.rx.try_recv() {
107 if sample.t0 > start_time {
108 break;
109 }
110
111 let thread_id = self.threads[sample.tid.0].name.as_str();
112 let t0 = sample.t0 / 1000;
113 let t1 = sample.t1 / 1000;
114
115 data.push(json!({
116 "pid": 0,
117 "tid": thread_id,
118 "name": sample.name,
119 "ph": "B",
120 "ts": t0
121 }));
122
123 data.push(json!({
124 "pid": 0,
125 "tid": thread_id,
126 "ph": "E",
127 "ts": t1
128 }));
129 }
130
131 let f = BufWriter::new(File::create(filename).unwrap());
132 serde_json::to_writer(f, &data).unwrap();
133 }
134}
135
136#[doc(hidden)]
137pub struct ProfileScope {
138 name: String,
139 t0: u64,
140}
141
142impl ProfileScope {
143 pub fn new(name: String) -> ProfileScope {
144 let t0 = precise_time_ns();
145 ProfileScope { name: name, t0: t0 }
146 }
147}
148
149impl Drop for ProfileScope {
150 fn drop(&mut self) {
154 let t1 = precise_time_ns();
155
156 THREAD_PROFILER.with(|profiler| match *profiler.borrow() {
157 Some(ref profiler) => {
158 profiler.push_sample(self.name.clone(), self.t0, t1);
159 }
160 None => {
161 println!("ERROR: ProfileScope {} on unregistered thread!", self.name);
162 }
163 });
164 }
165}
166
167pub fn write_profile(filename: &str) {
169 GLOBAL_PROFILER.lock().unwrap().write_profile(filename);
170}
171
172pub fn register_thread_with_profiler() {
174 GLOBAL_PROFILER.lock().unwrap().register_thread();
175}