metrics_utils_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse::Parse, parse::ParseStream, parse_macro_input, ItemFn, LitStr};
4
5struct MacroArgs {
6    custom_name: Option<LitStr>,
7}
8
9impl Parse for MacroArgs {
10    fn parse(input: ParseStream) -> syn::Result<Self> {
11        if input.is_empty() {
12            return Ok(MacroArgs { custom_name: None });
13        }
14
15        let custom_name = input.parse()?;
16        Ok(MacroArgs {
17            custom_name: Some(custom_name),
18        })
19    }
20}
21
22/// A procedural macro attribute that measures the execution time of an async function.
23///
24/// This macro wraps an async function to record its execution duration as a histogram metric.
25/// The duration is recorded using the `FunctionDurationSeconds` metric with a "function" label
26/// containing either the function name or a custom name if provided.
27///
28/// # Arguments
29///
30/// * `attr` - Optional custom name for the metric label
31/// * `item` - The async function to be measured
32///
33/// # Examples
34///
35/// Basic usage with default function name as label:
36/// ```ignore
37/// use metrics_macros::measured_async_function;
38///
39/// #[measured_async_function]
40/// async fn process_data() {
41///     // Function implementation
42/// }
43/// ```
44///
45/// Using a custom name for the metric label:
46/// ```ignore
47/// use metrics_macros::measured_async_function;
48///
49/// #[measured_async_function("custom_process_name")]
50/// async fn process_data() {
51///     // Function implementation
52/// }
53/// ```
54///
55/// The macro will record timing metrics that can be queried like:
56/// `function_duration_seconds{function="process_data"}` or
57/// `function_duration_seconds{function="custom_process_name"}`
58///
59#[proc_macro_attribute]
60pub fn measured_async_function(attr: TokenStream, item: TokenStream) -> TokenStream {
61    let args = parse_macro_input!(attr as MacroArgs);
62    let input_fn = parse_macro_input!(item as ItemFn);
63    let fn_name = &input_fn.sig.ident;
64    let fn_block = &input_fn.block;
65    let vis = &input_fn.vis;
66    let attrs = &input_fn.attrs;
67    let sig = &input_fn.sig;
68
69    let metric_name = match args.custom_name {
70        Some(name) => quote! { #name },
71        None => quote! { stringify!(#fn_name) },
72    };
73
74    let output = quote! {
75        #(#attrs)*
76        #vis #sig {
77            let __measured_async = async move #fn_block;
78            let __measured_start = std::time::Instant::now();
79            let __measured_result = __measured_async.await;
80            let __measured_duration = __measured_start.elapsed().as_millis() as f64;
81            metrics::histogram!(
82                "async_function_duration_milliseconds",
83                &[("function", #metric_name)]
84            ).record(__measured_duration);
85            __measured_result
86        }
87    };
88
89    output.into()
90}
91
92/// Same as measured_async_function but for sync functions
93#[proc_macro_attribute]
94pub fn measured_function(attr: TokenStream, item: TokenStream) -> TokenStream {
95    let args = parse_macro_input!(attr as MacroArgs);
96    let input_fn = parse_macro_input!(item as ItemFn);
97    let fn_name = &input_fn.sig.ident;
98    let fn_block = &input_fn.block;
99    let vis = &input_fn.vis;
100    let attrs = &input_fn.attrs;
101    let sig = &input_fn.sig;
102
103    let metric_name = match args.custom_name {
104        Some(name) => quote! { #name },
105        None => quote! { stringify!(#fn_name) },
106    };
107
108    let output = quote! {
109        #(#attrs)*
110        #vis #sig {
111            let start = std::time::Instant::now();
112            let result = (|| #fn_block)();
113            let duration = start.elapsed().as_millis() as f64;
114
115            metrics::histogram!(
116                "function_duration_milliseconds",
117                &[("function", #metric_name)]
118            ).record(duration);
119
120            result
121        }
122    };
123
124    output.into()
125}