internal_tracing/
lib.rs

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

internal_tracing/
lib.rs
