1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::{parse_macro_input, Expr, ItemFn, Lit, LitStr, Meta, Result as SynResult, Token};

struct FnMetricsAttributes {
    registry: Expr,
    name: Option<Lit>,
}

impl Parse for FnMetricsAttributes {
    /// to parse #[timed(name="...", registry="...")]
    fn parse(input: ParseStream) -> SynResult<Self> {
        let metas = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
        let mut result = FnMetricsAttributes {
            registry: syn::parse_str("metriki_core::global::global_registry()")?,
            name: None,
        };

        // convert attribute metas to key-value map
        for i in metas {
            if let Meta::NameValue(mnv) = i {
                if mnv.path.is_ident("name") {
                    if let Lit::Str(ref litstr) = mnv.lit {
                        result.name = Some(Lit::Str(litstr.clone()));
                    }
                }

                if mnv.path.is_ident("registry") {
                    if let Lit::Str(ref litstr) = mnv.lit {
                        result.registry = syn::parse_str(&litstr.value())?;
                    }
                }
            }
        }

        Ok(result)
    }
}

/// `timed` macro is design as an attibute for function.
///
/// The macro adds timer metric for the function execution.
///
/// Available options:
///
/// * `registry`: the code to access `MetricsRegistry`. `global_registry()` by default.
/// * `name`: name for the timer metric. Function name is used by default,
///   be careful with name conflict in different module because the macro cannot detect the module name.
#[proc_macro_attribute]
pub fn timed(attrs: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);
    let timer_data = parse_macro_input!(attrs as FnMetricsAttributes);

    // function data
    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = f;
    let stmts = &block.stmts;

    let registry = timer_data.registry;
    // use function name by default
    let name = timer_data
        .name
        .unwrap_or_else(|| Lit::Str(LitStr::new(&sig.ident.to_string(), Span::call_site())));

    // generated code
    let tokens = quote! {
        #(#attrs)*
        #vis #sig {
            let __timer = #registry.timer(#name);
            let __timer_ctx = __timer.start();

            #(#stmts)*
        }
    };

    tokens.into()
}

/// `metered` attribute adds a meter to current function.
///
/// Available options:
///
/// * `registry`: the code to access `MetricsRegistry`. `global_registry()` by default.
/// * `name`: name for the timer metric. Function name is used by default,
///   be careful with name conflict in different module because the macro cannot detect the module name.
#[proc_macro_attribute]
pub fn metered(attrs: TokenStream, input: TokenStream) -> TokenStream {
    let f = parse_macro_input!(input as ItemFn);
    let timer_data = parse_macro_input!(attrs as FnMetricsAttributes);

    // function data
    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = f;
    let stmts = &block.stmts;

    let registry = timer_data.registry;
    // use function name by default
    let name = timer_data
        .name
        .unwrap_or_else(|| Lit::Str(LitStr::new(&sig.ident.to_string(), Span::call_site())));

    // generated code
    let tokens = quote! {
        #(#attrs)*
        #vis #sig {
            #registry.meter(#name).mark();

            #(#stmts)*
        }
    };

    tokens.into()
}