metriki_macros/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::quote;
4use syn::parse::{Parse, ParseStream};
5use syn::punctuated::Punctuated;
6use syn::{parse_macro_input, Expr, ItemFn, Lit, LitStr, Meta, Result as SynResult, Token};
7
8struct FnMetricsAttributes {
9    registry: Expr,
10    name: Option<Lit>,
11}
12
13impl Parse for FnMetricsAttributes {
14    /// to parse #[timed(name="...", registry="...")]
15    fn parse(input: ParseStream) -> SynResult<Self> {
16        let metas = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
17        let mut result = FnMetricsAttributes {
18            registry: syn::parse_str("metriki_core::global::global_registry()")?,
19            name: None,
20        };
21
22        // convert attribute metas to key-value map
23        for i in metas {
24            if let Meta::NameValue(mnv) = i {
25                if mnv.path.is_ident("name") {
26                    if let Lit::Str(ref litstr) = mnv.lit {
27                        result.name = Some(Lit::Str(litstr.clone()));
28                    }
29                }
30
31                if mnv.path.is_ident("registry") {
32                    if let Lit::Str(ref litstr) = mnv.lit {
33                        result.registry = syn::parse_str(&litstr.value())?;
34                    }
35                }
36            }
37        }
38
39        Ok(result)
40    }
41}
42
43/// `timed` macro is design as an attibute for function.
44///
45/// The macro adds timer metric for the function execution.
46///
47/// Available options:
48///
49/// * `registry`: the code to access `MetricsRegistry`. `global_registry()` by default.
50/// * `name`: name for the timer metric. Function name is used by default,
51///   be careful with name conflict in different module because the macro cannot detect the module name.
52#[proc_macro_attribute]
53pub fn timed(attrs: TokenStream, input: TokenStream) -> TokenStream {
54    let f = parse_macro_input!(input as ItemFn);
55    let timer_data = parse_macro_input!(attrs as FnMetricsAttributes);
56
57    // function data
58    let ItemFn {
59        attrs,
60        vis,
61        sig,
62        block,
63    } = f;
64    let stmts = &block.stmts;
65
66    let registry = timer_data.registry;
67    // use function name by default
68    let name = timer_data
69        .name
70        .unwrap_or_else(|| Lit::Str(LitStr::new(&sig.ident.to_string(), Span::call_site())));
71
72    // generated code
73    let tokens = quote! {
74        #(#attrs)*
75        #vis #sig {
76            let __timer = #registry.timer(#name);
77            let __timer_ctx = __timer.start();
78
79            #(#stmts)*
80        }
81    };
82
83    tokens.into()
84}
85
86/// `metered` attribute adds a meter to current function.
87///
88/// Available options:
89///
90/// * `registry`: the code to access `MetricsRegistry`. `global_registry()` by default.
91/// * `name`: name for the timer metric. Function name is used by default,
92///   be careful with name conflict in different module because the macro cannot detect the module name.
93#[proc_macro_attribute]
94pub fn metered(attrs: TokenStream, input: TokenStream) -> TokenStream {
95    let f = parse_macro_input!(input as ItemFn);
96    let timer_data = parse_macro_input!(attrs as FnMetricsAttributes);
97
98    // function data
99    let ItemFn {
100        attrs,
101        vis,
102        sig,
103        block,
104    } = f;
105    let stmts = &block.stmts;
106
107    let registry = timer_data.registry;
108    // use function name by default
109    let name = timer_data
110        .name
111        .unwrap_or_else(|| Lit::Str(LitStr::new(&sig.ident.to_string(), Span::call_site())));
112
113    // generated code
114    let tokens = quote! {
115        #(#attrs)*
116        #vis #sig {
117            #registry.meter(#name).mark();
118
119            #(#stmts)*
120        }
121    };
122
123    tokens.into()
124}
metriki_macros/lib.rs

metriki_macros/
lib.rs
