internal_tracing/
lib.rs

1use std::time::SystemTime;
2
3#[cfg(feature = "enabled")]
4pub use serde_json::{json, to_writer as json_to_writer, Value as JsonValue};
5
6pub fn time_to_micros(time: SystemTime) -> u64 {
7    time.duration_since(SystemTime::UNIX_EPOCH)
8        .unwrap()
9        .as_micros() as u64
10}
11
12pub fn now_micros() -> u64 {
13    time_to_micros(SystemTime::now())
14}
15
16pub enum TimeInput {
17    Microseconds(u64),
18    SystemTime(SystemTime),
19}
20
21impl TimeInput {
22    pub fn micros(self) -> u64 {
23        match self {
24            Self::Microseconds(v) => v,
25            Self::SystemTime(v) => time_to_micros(v),
26        }
27    }
28}
29
30impl From<u64> for TimeInput {
31    fn from(value: u64) -> Self {
32        Self::Microseconds(value)
33    }
34}
35
36impl From<SystemTime> for TimeInput {
37    fn from(value: SystemTime) -> Self {
38        Self::SystemTime(value)
39    }
40}
41
42/// Declare traces group.
43///
44/// Creates a module with the passed `$name`.
45///
46/// **Note:**Traces are stored locally for each thread (`thread_local`).
47///
48/// Module exposes two main methods:
49/// - `fn start_tracing()` - Simply calls `take_traces()` and discards
50///   the result, in order to clean up old traces.
51/// - `fn take_traces()` - Take accumulated traces.
52#[cfg(feature = "enabled")]
53#[macro_export]
54macro_rules! decl_traces {
55    ($name:ident; $($checkpoint:ident),+) => {
56
57        pub mod $name {
58            use std::rc::Rc;
59            use std::cell::RefCell;
60
61            #[derive(serde::Serialize, Debug, Default, Clone)]
62            pub struct Traces {
63                $(
64                pub $checkpoint: (u64, $crate::JsonValue),
65                )*
66            }
67
68            impl std::fmt::Display for Traces {
69                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70                    let mut arr = [
71                        $(
72                        (stringify!($checkpoint), self.$checkpoint.0, &self.$checkpoint.1),
73                        )+
74                    ];
75                    arr.sort_by_key(|v| v.1);
76                    let mut buf = Vec::new();
77                    for (name, time, meta) in arr.into_iter().filter(|v| v.1 != 0) {
78                        $crate::json_to_writer(&mut buf, &(name, (time as f64) / 1_000_000.0)).unwrap();
79                        buf.push(b'\n');
80                        if !meta.is_null() {
81                            $crate::json_to_writer(&mut buf, meta).unwrap();
82                            buf.push(b'\n');
83                        }
84                    }
85                    write!(f, "{}", &String::from_utf8(buf).unwrap())
86                }
87            }
88
89            impl From<Traces> for String {
90                fn from(t: Traces) -> Self {
91                    t.to_string()
92                }
93            }
94
95            thread_local! {
96                pub static TRACES: Rc<RefCell<Traces>> = Default::default();
97            }
98
99            /// Clean up old traces and start fresh.
100            pub fn start_tracing() {
101                take_traces();
102            }
103
104            /// Take captured traces.
105            pub fn take_traces() -> Traces {
106                TRACES.with(|v| v.take())
107            }
108
109            #[cfg(feature = "ocaml_types")]
110            #[allow(dead_code, non_local_definitions)]
111            pub mod caml {
112                use super::*;
113
114                #[derive(Debug, ocaml::IntoValue, ocaml::FromValue, ocaml_gen::Struct)]
115                pub struct CamlTraces(String);
116
117                impl From<Traces> for CamlTraces {
118                    fn from(t: Traces) -> Self {
119                        Self(t.to_string())
120                    }
121                }
122            }
123        }
124    };
125}
126/// Noop. Internal tracing not enabled!
127#[cfg(not(feature = "enabled"))]
128#[macro_export]
129macro_rules! decl_traces {
130    ($($_ignored:tt)+) => {};
131}
132
133/// Capture the trace/checkpoint.
134#[cfg(feature = "enabled")]
135#[macro_export]
136macro_rules! checkpoint {
137    ($name:ident; $checkpoint:ident) => {
138        $crate::checkpoint!(|$name; $checkpoint, $crate::now_micros(), $crate::json!(null));
139    };
140    ($name:ident; $checkpoint:ident, {$($metadata:tt)+}) => {
141        $crate::checkpoint!(|$name; $checkpoint, $crate::now_micros(), $crate::json!({$($metadata)+}));
142    };
143    ($name:ident; $checkpoint:ident, $time:expr) => {
144        $crate::checkpoint!(|$name; $checkpoint, $crate::TimeInput::from($time).micros(), $crate::json!(null));
145    };
146    ($name:ident; $checkpoint:ident, $time:expr, {$($metadata:tt)+}) => {
147        $crate::checkpoint!(|$name; $checkpoint, $crate::TimeInput::from($time).micros(), $crate::json!({$($metadata)+}));
148    };
149    (|$name:ident; $checkpoint:ident, $time:expr, $metadata:expr) => {
150        $name::TRACES.with(|traces| traces.borrow_mut().$checkpoint = ($time, $metadata));
151    };
152}
153/// Noop. Internal tracing not enabled!
154#[cfg(not(feature = "enabled"))]
155#[macro_export]
156macro_rules! checkpoint {
157    ($($_ignored:tt)+) => {};
158}
159
160#[cfg(feature = "enabled")]
161#[cfg(test)]
162mod tests {
163    decl_traces!(test_fn; c1, c2, c3, c4);
164
165    #[test]
166    fn test_fn() {
167        test_fn::start_tracing();
168
169        checkpoint!(test_fn; c1);
170        checkpoint!(test_fn; c2, 2);
171        checkpoint!(test_fn; c3, { "arg": 1 });
172        checkpoint!(test_fn; c4, 3, { "arg": 2 });
173
174        let traces = test_fn::take_traces();
175
176        assert_ne!(traces.c1.0, 0);
177        assert_eq!(traces.c1.1, serde_json::Value::Null);
178
179        assert_eq!(traces.c2.0, 2);
180        assert_eq!(traces.c2.1, serde_json::Value::Null);
181
182        assert_ne!(traces.c3.0, 0);
183        assert_eq!(traces.c3.1, serde_json::json!({ "arg": 1 }));
184
185        assert_eq!(traces.c4.0, 3);
186        assert_eq!(traces.c4.1, serde_json::json!({ "arg": 2 }));
187    }
188}
internal_tracing/lib.rs

internal_tracing/
lib.rs
