metrics_utils_macros/
lib.rs1use 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#[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#[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}