hotpath_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::parse::Parser;
4use syn::{parse_macro_input, ItemFn, LitInt, LitStr};
5
6#[derive(Clone, Copy)]
7enum Format {
8    Table,
9    Json,
10    JsonPretty,
11}
12
13impl Format {
14    fn to_tokens(self) -> proc_macro2::TokenStream {
15        match self {
16            Format::Table => quote!(hotpath::Format::Table),
17            Format::Json => quote!(hotpath::Format::Json),
18            Format::JsonPretty => quote!(hotpath::Format::JsonPretty),
19        }
20    }
21}
22
23/// Initializes the hotpath profiling system and generates a performance report on program exit.
24///
25/// This attribute macro should be applied to your program's main (or other entry point) function to enable profiling.
26/// It creates a guard that initializes the background measurement processing thread and
27/// automatically displays a performance summary when the program exits.
28/// Additionally it creates a measurement guard that will be used to measure the wrapper function itself.
29///
30/// # Parameters
31///
32/// * `percentiles` - Array of percentile values (0-100) to display in the report. Default: `[95]`
33/// * `format` - Output format as a string: `"table"` (default), `"json"`, or `"json-pretty"`
34///
35/// # Examples
36///
37/// Basic usage with default settings (P95 percentile, table format):
38///
39/// ```rust,no_run
40/// #[cfg_attr(feature = "hotpath", hotpath::main)]
41/// fn main() {
42///     // Your code here
43/// }
44/// ```
45///
46/// Custom percentiles:
47///
48/// ```rust,no_run
49/// #[tokio::main]
50/// #[cfg_attr(feature = "hotpath", hotpath::main(percentiles = [50, 90, 95, 99]))]
51/// async fn main() {
52///     // Your code here
53/// }
54/// ```
55///
56/// JSON output format:
57///
58/// ```rust,no_run
59/// #[cfg_attr(feature = "hotpath", hotpath::main(format = "json-pretty"))]
60/// fn main() {
61///     // Your code here
62/// }
63/// ```
64///
65/// Combined parameters:
66///
67/// ```rust,no_run
68/// #[cfg_attr(feature = "hotpath", hotpath::main(percentiles = [50, 99], format = "json"))]
69/// fn main() {
70///     // Your code here
71/// }
72/// ```
73///
74/// # Usage with Tokio
75///
76/// When using with tokio, place `#[tokio::main]` before `#[hotpath::main]`:
77///
78/// ```rust,no_run
79/// #[tokio::main]
80/// #[cfg_attr(feature = "hotpath", hotpath::main)]
81/// async fn main() {
82///     // Your code here
83/// }
84/// ```
85///
86/// # Limitations
87///
88/// Only one hotpath guard can be active at a time. Creating a second guard (either via this
89/// macro or via [`GuardBuilder`](../hotpath/struct.GuardBuilder.html)) will cause a panic.
90///
91/// # See Also
92///
93/// * [`measure`](macro@measure) - Attribute macro for instrumenting functions
94/// * [`measure_block!`](../hotpath/macro.measure_block.html) - Macro for measuring code blocks
95/// * [`GuardBuilder`](../hotpath/struct.GuardBuilder.html) - Manual control over profiling lifecycle
96#[proc_macro_attribute]
97pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
98    let input = parse_macro_input!(item as ItemFn);
99    let vis = &input.vis;
100    let sig = &input.sig;
101    let block = &input.block;
102
103    // Defaults
104    let mut percentiles: Vec<u8> = vec![95];
105    let mut format = Format::Table;
106
107    // Parse named args like: percentiles=[..], format=".."
108    if !attr.is_empty() {
109        let parser = syn::meta::parser(|meta| {
110            if meta.path.is_ident("percentiles") {
111                meta.input.parse::<syn::Token![=]>()?;
112                let content;
113                syn::bracketed!(content in meta.input);
114                let mut vals = Vec::new();
115                while !content.is_empty() {
116                    let li: LitInt = content.parse()?;
117                    let v: u8 = li.base10_parse()?;
118                    if !(0..=100).contains(&v) {
119                        return Err(
120                            meta.error(format!("Invalid percentile {} (must be 0..=100)", v))
121                        );
122                    }
123                    vals.push(v);
124                    if !content.is_empty() {
125                        content.parse::<syn::Token![,]>()?;
126                    }
127                }
128                if vals.is_empty() {
129                    return Err(meta.error("At least one percentile must be specified"));
130                }
131                percentiles = vals;
132                return Ok(());
133            }
134
135            if meta.path.is_ident("format") {
136                meta.input.parse::<syn::Token![=]>()?;
137                let lit: LitStr = meta.input.parse()?;
138                format =
139                    match lit.value().as_str() {
140                        "table" => Format::Table,
141                        "json" => Format::Json,
142                        "json-pretty" => Format::JsonPretty,
143                        other => return Err(meta.error(format!(
144                            "Unknown format {:?}. Expected one of: \"table\", \"json\", \"json-pretty\"",
145                            other
146                        ))),
147                    };
148                return Ok(());
149            }
150
151            Err(meta.error("Unknown parameter. Supported: percentiles=[..], format=\"..\""))
152        });
153
154        if let Err(e) = parser.parse2(proc_macro2::TokenStream::from(attr)) {
155            return e.to_compile_error().into();
156        }
157    }
158
159    let percentiles_array = quote! { &[#(#percentiles),*] };
160    let format_token = format.to_tokens();
161
162    let asyncness = sig.asyncness.is_some();
163    let fn_name = &sig.ident;
164    let measurement_name = quote! { concat!(module_path!(), "::", stringify!(#fn_name)) };
165
166    let output = if asyncness {
167        quote! {
168            #vis #sig {
169                async {
170                    let _hotpath = {
171                        fn __caller_fn() {}
172                        let caller_name = std::any::type_name_of_val(&__caller_fn)
173                            .strip_suffix("::__caller_fn")
174                            .unwrap_or(std::any::type_name_of_val(&__caller_fn))
175                            .replace("::{{closure}}", "");
176
177                        hotpath::GuardBuilder::new(caller_name.to_string())
178                            .percentiles(#percentiles_array)
179                            .format(#format_token)
180                            .build()
181                    };
182
183                    hotpath::cfg_if! {
184                        if #[cfg(feature = "hotpath-off")] {
185                            // No-op when hotpath-off is enabled
186                        } else if #[cfg(any(
187                            feature = "hotpath-alloc-bytes-total",
188                            feature = "hotpath-alloc-bytes-max",
189                            feature = "hotpath-alloc-count-total",
190                            feature = "hotpath-alloc-count-max"
191                        ))] {
192                            use hotpath::{Handle, RuntimeFlavor};
193                            let runtime_flavor = Handle::try_current().ok().map(|h| h.runtime_flavor());
194
195                            let _measure_guard = match runtime_flavor {
196                                Some(RuntimeFlavor::CurrentThread) => {
197                                    hotpath::AllocGuardType::AllocGuard(hotpath::AllocGuard::new(#measurement_name))
198                                }
199                                _ => {
200                                    hotpath::AllocGuardType::NoopAsyncAllocGuard(hotpath::NoopAsyncAllocGuard::new(#measurement_name))
201                                }
202                            };
203                        } else {
204                            let _measure_guard = hotpath::TimeGuard::new(#measurement_name);
205                        }
206                    }
207
208                    #block
209                }.await
210            }
211        }
212    } else {
213        quote! {
214            #vis #sig {
215                let _hotpath = {
216                    fn __caller_fn() {}
217                    let caller_name = std::any::type_name_of_val(&__caller_fn)
218                        .strip_suffix("::__caller_fn")
219                        .unwrap_or(std::any::type_name_of_val(&__caller_fn))
220                        .replace("::{{closure}}", "");
221
222                    hotpath::GuardBuilder::new(caller_name.to_string())
223                        .percentiles(#percentiles_array)
224                        .format(#format_token)
225                        .build()
226                };
227
228                hotpath::cfg_if! {
229                    if #[cfg(feature = "hotpath-off")] {
230                        // No-op when hotpath-off is enabled
231                    } else if #[cfg(any(
232                        feature = "hotpath-alloc-bytes-total",
233                        feature = "hotpath-alloc-bytes-max",
234                        feature = "hotpath-alloc-count-total",
235                        feature = "hotpath-alloc-count-max"
236                    ))] {
237                        let _measure_guard = hotpath::AllocGuard::new(#measurement_name);
238                    } else {
239                        let _measure_guard = hotpath::TimeGuard::new(#measurement_name);
240                    }
241                }
242
243                #block
244            }
245        }
246    };
247
248    output.into()
249}
250
251/// Instruments a function to send performance measurements to the hotpath profiler.
252///
253/// This attribute macro wraps functions with profiling code that measures execution time
254/// or memory allocations (depending on enabled feature flags). The measurements are sent
255/// to a background processing thread for aggregation.
256///
257/// # Behavior
258///
259/// The macro automatically detects whether the function is sync or async and instruments
260/// it appropriately. Measurements include:
261///
262/// * **Time profiling** (default): Execution duration using high-precision timers
263/// * **Allocation profiling**: Memory allocations when allocation features are enabled
264///   - `hotpath-alloc-bytes-total` - Total bytes allocated
265///   - `hotpath-alloc-bytes-max` - Peak memory usage
266///   - `hotpath-alloc-count-total` - Total allocation count
267///   - `hotpath-alloc-count-max` - Peak allocation count
268///
269/// # Async Function Limitations
270///
271/// When using allocation profiling features with async functions, you must use the
272/// `tokio` runtime in `current_thread` mode:
273///
274/// ```rust,no_run
275/// #[tokio::main(flavor = "current_thread")]
276/// async fn main() {
277///     // Your async code here
278/// }
279/// ```
280///
281/// This limitation exists because allocation tracking uses thread-local storage. In multi-threaded
282/// runtimes, async tasks can migrate between threads, making it impossible to accurately
283/// attribute allocations to specific function calls. Time-based profiling works with any runtime flavor.
284///
285/// When the `hotpath` feature is disabled, this macro compiles to zero overhead (no instrumentation).
286///
287/// # See Also
288///
289/// * [`main`](macro@main) - Attribute macro that initializes profiling
290/// * [`measure_block!`](../hotpath/macro.measure_block.html) - Macro for measuring code blocks
291#[proc_macro_attribute]
292pub fn measure(_attr: TokenStream, item: TokenStream) -> TokenStream {
293    let input = parse_macro_input!(item as ItemFn);
294    let vis = &input.vis;
295    let sig = &input.sig;
296    let block = &input.block;
297
298    let name = sig.ident.to_string();
299    let asyncness = sig.asyncness.is_some();
300
301    let output = if asyncness {
302        quote! {
303            #vis #sig {
304                async {
305                    hotpath::cfg_if! {
306                        if #[cfg(feature = "hotpath-off")] {
307                            // No-op when hotpath-off is enabled
308                        } else if #[cfg(any(
309                            feature = "hotpath-alloc-bytes-total",
310                            feature = "hotpath-alloc-bytes-max",
311                            feature = "hotpath-alloc-count-total",
312                            feature = "hotpath-alloc-count-max"
313                        ))] {
314                            use hotpath::{Handle, RuntimeFlavor};
315                            let runtime_flavor = Handle::try_current().ok().map(|h| h.runtime_flavor());
316
317                            let _guard = match runtime_flavor {
318                                Some(RuntimeFlavor::CurrentThread) => {
319                                    hotpath::AllocGuardType::AllocGuard(hotpath::AllocGuard::new(concat!(module_path!(), "::", #name)))
320                                }
321                                _ => {
322                                    hotpath::AllocGuardType::NoopAsyncAllocGuard(hotpath::NoopAsyncAllocGuard::new(concat!(module_path!(), "::", #name)))
323                                }
324                            };
325                        } else {
326                            let _guard = hotpath::TimeGuard::new(concat!(module_path!(), "::", #name));
327                        }
328                    }
329
330                    #block
331                }.await
332            }
333        }
334    } else {
335        quote! {
336            #vis #sig {
337                hotpath::cfg_if! {
338                    if #[cfg(feature = "hotpath-off")] {
339                        // No-op when hotpath-off is enabled
340                    } else if #[cfg(any(
341                        feature = "hotpath-alloc-bytes-total",
342                        feature = "hotpath-alloc-bytes-max",
343                        feature = "hotpath-alloc-count-total",
344                        feature = "hotpath-alloc-count-max"
345                    ))] {
346                        let _guard = hotpath::AllocGuard::new(concat!(module_path!(), "::", #name));
347                    } else {
348                        let _guard = hotpath::TimeGuard::new(concat!(module_path!(), "::", #name));
349                    }
350                }
351
352                #block
353            }
354        }
355    };
356
357    output.into()
358}