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            pub mod caml {
111                use super::*;
112
113                #[derive(Debug, ocaml::IntoValue, ocaml::FromValue, ocaml_gen::Struct)]
114                pub struct CamlTraces(String);
115
116                impl From<Traces> for CamlTraces {
117                    fn from(t: Traces) -> Self {
118                        Self(t.to_string())
119                    }
120                }
121            }
122        }
123    };
124}
125/// Noop. Internal tracing not enabled!
126#[cfg(not(feature = "enabled"))]
127#[macro_export]
128macro_rules! decl_traces {
129    ($($_ignored:tt)+) => {};
130}
131
132/// Capture the trace/checkpoint.
133#[cfg(feature = "enabled")]
134#[macro_export]
135macro_rules! checkpoint {
136    ($name:ident; $checkpoint:ident) => {
137        $crate::checkpoint!(|$name; $checkpoint, $crate::now_micros(), $crate::json!(null));
138    };
139    ($name:ident; $checkpoint:ident, {$($metadata:tt)+}) => {
140        $crate::checkpoint!(|$name; $checkpoint, $crate::now_micros(), $crate::json!({$($metadata)+}));
141    };
142    ($name:ident; $checkpoint:ident, $time:expr) => {
143        $crate::checkpoint!(|$name; $checkpoint, $crate::TimeInput::from($time).micros(), $crate::json!(null));
144    };
145    ($name:ident; $checkpoint:ident, $time:expr, {$($metadata:tt)+}) => {
146        $crate::checkpoint!(|$name; $checkpoint, $crate::TimeInput::from($time).micros(), $crate::json!({$($metadata)+}));
147    };
148    (|$name:ident; $checkpoint:ident, $time:expr, $metadata:expr) => {
149        $name::TRACES.with(|traces| traces.borrow_mut().$checkpoint = ($time, $metadata));
150    };
151}
152/// Noop. Internal tracing not enabled!
153#[cfg(not(feature = "enabled"))]
154#[macro_export]
155macro_rules! checkpoint {
156    ($($_ignored:tt)+) => {};
157}
158
159#[cfg(feature = "enabled")]
160#[cfg(test)]
161mod tests {
162    use super::*;
163
164    decl_traces!(test_fn; c1, c2, c3, c4);
165
166    #[test]
167    fn test_fn() {
168        test_fn::start_tracing();
169
170        checkpoint!(test_fn; c1);
171        checkpoint!(test_fn; c2, 2);
172        checkpoint!(test_fn; c3, { "arg": 1 });
173        checkpoint!(test_fn; c4, 3, { "arg": 2 });
174
175        let traces = test_fn::take_traces();
176
177        assert_ne!(traces.c1.0, 0);
178        assert_eq!(traces.c1.1, serde_json::Value::Null);
179
180        assert_eq!(traces.c2.0, 2);
181        assert_eq!(traces.c2.1, serde_json::Value::Null);
182
183        assert_ne!(traces.c3.0, 0);
184        assert_eq!(traces.c3.1, serde_json::json!({ "arg": 1 }));
185
186        assert_eq!(traces.c4.0, 3);
187        assert_eq!(traces.c4.1, serde_json::json!({ "arg": 2 }));
188    }
189}
internal_tracing/lib.rs

internal_tracing/
lib.rs
