thread_profiler/
lib.rs

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        // Stop reading samples that are written after
102        // write_profile() is called.
103        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    /// When the ProfileScope is dropped it records the
151    /// length of time it was alive for and records it
152    /// against the Profiler.
153    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
167/// Writes the global profile to a specific file.
168pub fn write_profile(filename: &str) {
169    GLOBAL_PROFILER.lock().unwrap().write_profile(filename);
170}
171
172/// Registers the current thread with the global profiler.
173pub fn register_thread_with_profiler() {
174    GLOBAL_PROFILER.lock().unwrap().register_thread();
175}