oni_trace/
lib.rs

1#![feature(decl_macro)]
2
3#[macro_use] extern crate serde_derive;
4#[macro_use] extern crate lazy_static;
5
6use time::precise_time_ns;
7use deflate::{
8    Compression,
9    write::GzEncoder,
10};
11use std::{
12    borrow::Cow,
13    time::Duration,
14    fs::{remove_file, OpenOptions},
15    io::{BufWriter, Write},
16    thread::{spawn, sleep, JoinHandle},
17    sync::mpsc::{channel, Sender},
18    sync::atomic::{AtomicUsize, Ordering},
19};
20
21mod scope;
22mod local;
23mod global;
24mod trace;
25pub mod colors;
26
27pub use self::scope::ScopeComplete;
28pub use self::local::{Local, LOCAL};
29pub use self::global::{Global, GLOBAL};
30pub use self::trace::{Event, Base, Instant, Async, Args, Flow};
31
32pub const ENABLED: bool = cfg!(feature = "trace");
33pub const TRACE_LOC: bool = cfg!(feature = "trace_location");
34
35pub fn generate_id() -> usize {
36    static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
37    NEXT_ID.fetch_add(1, Ordering::Relaxed)
38}
39
40use log::{Log, Metadata, Record};
41
42pub struct Logger;
43
44impl Log for Logger {
45    fn enabled(&self, _metadata: &Metadata) -> bool {
46        //metadata.level() <= Level::Info
47        true
48    }
49
50    fn log(&self, record: &Record) {
51        let ts = precise_time_ns();
52        if self.enabled(record.metadata()) {
53            LOCAL.with(|profiler| match *profiler.borrow() {
54                Some(ref profiler) => profiler.log(ts, record),
55                None => println!("ERROR: push_log on unregistered thread!"),
56            });
57        }
58    }
59
60    fn flush(&self) {}
61}
62
63pub struct AppendWorker {
64    handle: Option<JoinHandle<()>>,
65    tx: Sender<()>,
66}
67
68impl AppendWorker {
69    pub fn new(filename: &str, duration: Duration) -> Self {
70        let _ = remove_file(filename);
71        let w = OpenOptions::new()
72            .create_new(true)
73            .append(true)
74            .open(filename)
75            .unwrap();
76
77        let (tx, rx) = channel();
78        let handle = spawn(move || {
79            let encoder = GzEncoder::new(w, Compression::Default);
80            let mut buf = BufWriter::new(encoder);
81            buf.write(b"[\n").ok();
82
83            loop {
84                sleep(duration);
85                GLOBAL.lock().unwrap().write_profile(&mut buf);
86                buf.flush().ok();
87                if rx.try_recv().is_ok() {
88                    break;
89                }
90            }
91
92            write_global_instant(&mut buf, "EOF");
93
94            buf.write(b"]\n").ok();
95            buf.flush().ok();
96        });
97
98        Self { handle: Some(handle), tx }
99    }
100
101    pub fn end(&mut self) {
102        self.tx.send(()).unwrap();
103        self.handle.take().unwrap().join().unwrap();
104    }
105}
106
107fn write_global_instant<W: Write>(w: &mut W, name: &'static str) {
108    let ts = precise_time_ns();
109    serde_json::to_writer(w, &Event::Instant {
110        s: "g",
111        ts: ts / 1000,
112        base: Base {
113            name: name.into(),
114            cat: None,
115            pid: 0,
116            tid: 0,
117            args: Args::Empty,
118            cname: Some(colors::WHITE),
119        },
120    }).ok();
121}
122
123/// Registers the current thread with the global profiler.
124pub fn register_thread(pid: Option<usize>, sort_index: Option<usize>) {
125    GLOBAL.lock().unwrap().register_thread(pid.unwrap_or(0), sort_index);
126}
127
128pub macro location() {
129    if $crate::TRACE_LOC {
130        $crate::Args::Location { module: module_path!(), file: file!(), line: line!() }
131    } else {
132        $crate::Args::Empty
133    }
134}
135
136pub macro instant {
137    (json $name:expr, $value:expr) => {
138        if $crate::ENABLED {
139            $crate::instant_thread($name, "", $crate::Args::Custom {
140                value: $value,
141            });
142        }
143    },
144    ($name:expr) => {
145        $crate::instant!($name => "");
146    },
147    ([ $($cat:ident)+ ] $name:expr) => {
148        $crate::instant!($name => stringify!($($cat,)+));
149    },
150    ($cat:expr => $name:expr) => {
151        if $crate::ENABLED {
152            $crate::instant_thread($name, $cat, $crate::location!());
153        }
154    }
155}
156
157pub macro scope($($name:tt)+) {
158    let _profile_scope = if $crate::ENABLED {
159        Some($crate::ScopeComplete::new(stringify!($($name)+), location!()))
160    } else {
161        None
162    };
163}
164
165#[doc(hidden)]
166pub fn instant_thread(name: &'static str, cat: &'static str, args: Args) {
167    let ts = precise_time_ns();
168    LOCAL.with(|profiler| match *profiler.borrow() {
169        Some(ref profiler) => profiler.instant_thread(ts, name, cat, args),
170        None => println!("ERROR: instant_thread on unregistered thread!"),
171    });
172}
173
174pub macro async_event {
175    ($kind:ident, $name:expr, $cat:expr, $id:expr, $cname:expr) => {
176        if $crate::ENABLED {
177            let cat: Option<&'static str> = $cat;
178            $crate::push_async($id, $name, $cat, $crate::Async::$kind, location!(), $cname);
179        }
180    }
181}
182
183pub macro async_start {
184    ($name:expr, $cat:expr, $id:expr) =>              { $crate::async_event!(Start, $name, $cat, $id, None); },
185    ($name:expr, $cat:expr, $id:expr, $cname:expr) => { $crate::async_event!(Start, $name, $cat, $id, Some($cname)); }
186}
187
188pub macro async_instant {
189    ($name:expr, $cat:expr, $id:expr) =>              { $crate::async_event!(Start, $name, $cat, $id, None); },
190    ($name:expr, $cat:expr, $id:expr, $cname:expr) => { $crate::async_event!(Start, $name, $cat, $id, Some($cname)); }
191}
192
193pub macro async_end {
194    ($name:expr, $cat:expr, $id:expr) =>              { $crate::async_event!(Start, $name, $cat, $id, None); },
195    ($name:expr, $cat:expr, $id:expr, $cname:expr) => { $crate::async_event!(Start, $name, $cat, $id, Some($cname)); }
196}
197
198#[doc(hidden)]
199pub fn push_async<N, C>(
200    id: usize,
201    name: N,
202    cat: Option<C>,
203    kind: Async,
204    args: Args,
205    cname: Option<&'static str>,
206)
207    where
208        N: Into<Cow<'static, str>>,
209        C: Into<Cow<'static, str>>,
210{
211    let ts = precise_time_ns();
212    LOCAL.with(|profiler| match *profiler.borrow() {
213        Some(ref profiler) => profiler.async_event(kind, ts, id, name.into(), cat.map(Into::into), None, args, cname),
214        None => println!("ERROR: push_async on unregistered thread!"),
215    });
216}
217
218pub macro flow_event {
219    ($kind:ident, $name:expr, $cat:expr, $id:expr, $cname:expr) => {
220        if $crate::ENABLED {
221            let cat: Option<&'static str> = $cat;
222            $crate::push_flow($id, $name, cat, $crate::Flow::$kind, location!(), $cname);
223        }
224    }
225}
226
227pub macro flow_start {
228    ($name:expr, $id:expr) =>              { $crate::flow_event!(Start, $name, None, $id, None); },
229    ($name:expr, $id:expr, $cname:expr) => { $crate::flow_event!(Start, $name, None, $id, Some($cname)); }
230}
231
232pub macro flow_step {
233    ($name:expr, $id:expr) =>              { $crate::flow_event!(Step, $name, None, $id, None); },
234    ($name:expr, $id:expr, $cname:expr) => { $crate::flow_event!(Step, $name, None, $id, Some($cname)); }
235}
236
237pub macro flow_end {
238    ($name:expr, $id:expr) =>              { $crate::flow_event!(End, $name, None, $id, None); },
239    ($name:expr, $id:expr, $cname:expr) => { $crate::flow_event!(End, $name, None, $id, Some($cname)); }
240}
241
242#[doc(hidden)]
243pub fn push_flow<N, C>(
244    id: usize,
245    name: N,
246    cat: Option<C>,
247    kind: Flow,
248    args: Args,
249    cname: Option<&'static str>,
250)
251    where
252        N: Into<Cow<'static, str>>,
253        C: Into<Cow<'static, str>>,
254{
255    let ts = precise_time_ns();
256    LOCAL.with(|profiler| match *profiler.borrow() {
257        Some(ref profiler) => profiler.flow_event(kind, ts, id, name.into(), cat.map(Into::into), args, cname),
258        None => println!("ERROR: push_flow on unregistered thread!"),
259    });
260}
oni_trace/lib.rs

oni_trace/
lib.rs
