metricrs_derive/
lib.rs

1use proc_macro::TokenStream;
2use quote::{ToTokens, quote};
3use syn::{Expr, ItemFn, LitStr, Result, meta, parse::Parser, parse_macro_input};
4
5struct InstrumentOptions {
6    kind: Option<Expr>,
7    fields: Vec<proc_macro2::TokenStream>,
8}
9
10impl InstrumentOptions {
11    pub fn parse(options: TokenStream) -> Result<Self> {
12        let mut kind = None;
13        let mut fields = vec![];
14
15        let parser = meta::parser(|meta| {
16            macro_rules! error {
17                ($($t:tt)+) => {
18                    return Err(meta.error(format_args!($($t)+)))
19                };
20            }
21
22            let Some(ident) = meta.path.get_ident() else {
23                error!("unsupported macro `instrument` option.");
24            };
25
26            match ident.to_string().as_str() {
27                "labels" => {
28                    let mut kv = vec![];
29
30                    meta.parse_nested_meta(|meta| {
31                        let Some(ident) = meta.path.get_ident() else {
32                            error!("expect label `name`.");
33                        };
34
35                        let value: LitStr = meta.value()?.parse()?;
36
37                        kv.push(quote! { (stringify!(#ident), #value) });
38
39                        Ok(())
40                    })?;
41
42                    fields.push(
43                        quote! { #ident: Some(&[("rust_module_path",module_path!()), #(#kv),*]) },
44                    );
45
46                    Ok(())
47                }
48                "kind" => {
49                    let value: proc_macro2::TokenStream = match meta.value() {
50                        Ok(value) => {
51                            if kind.is_some() {
52                                error!("repeated 'instrument' option 'kind'");
53                            }
54                            let expr: Expr = value.parse()?;
55                            kind = Some(expr.clone());
56                            quote! { #ident: Some(#expr) }
57                        }
58                        Err(_) => {
59                            quote! { #ident: None }
60                        }
61                    };
62
63                    fields.push(value);
64
65                    return Ok(());
66                }
67                _ => {
68                    let value: proc_macro2::TokenStream = match meta.value() {
69                        Ok(value) => {
70                            let expr: Expr = value.parse()?;
71                            quote! { #ident: Some(#expr) }
72                        }
73                        Err(_) => {
74                            quote! { #ident: None }
75                        }
76                    };
77
78                    fields.push(value);
79
80                    return Ok(());
81                }
82            }
83        });
84
85        parser.parse(options)?;
86
87        Ok(Self { kind, fields })
88    }
89}
90
91/// Create measuring instruments for methods via attribute
92#[proc_macro_attribute]
93pub fn instrument(options: TokenStream, item: TokenStream) -> TokenStream {
94    let InstrumentOptions { kind, fields } = match InstrumentOptions::parse(options) {
95        Ok(options) => options,
96        Err(err) => return err.into_compile_error().into(),
97    };
98
99    let ItemFn {
100        attrs,
101        vis,
102        sig,
103        block,
104    } = parse_macro_input!(item as ItemFn);
105
106    let block = if sig.asyncness.is_some() {
107        quote! {
108            async #block.await
109        }
110    } else {
111        quote!(#block)
112    };
113
114    let make_counter = || {
115        quote! {
116            #(#attrs)*
117            #vis #sig {
118                static COUNTER: std::sync::LazyLock<Option<metricrs::Counter>> = std::sync::LazyLock::new(|| {
119                    metricrs::global::get_global_registry().map(|registry| {
120                        use metricrs::*;
121                        use DeriveKind::*;
122                        registry.counter(DeriveOption {
123                            #(#fields,)*
124                            ..Default::default()
125                        }.into())
126                    })
127                });
128
129                if let Some(counter) = COUNTER.as_ref() {
130                    let r = #block;
131                    counter.increment(1);
132                    r
133                } else {
134                    #block
135                }
136            }
137        }
138    };
139
140    let make_timer = || {
141        quote! {
142            #(#attrs)*
143            #vis #sig {
144
145                static TIMER: std::sync::LazyLock<Option<metricrs::Histogram>> = std::sync::LazyLock::new(|| {
146                    metricrs::global::get_global_registry().map(|registry| {
147                        use metricrs::*;
148                        use DeriveKind::*;
149                        registry.histogam(DeriveOption {
150                              #(#fields,)*
151                            ..Default::default()
152                        }.into())
153                    })
154                });
155
156                if let Some(timer) = TIMER.as_ref() {
157                    let now = std::time::Instant::now();
158                    let r = #block;
159                    timer.record(now.elapsed().as_secs_f64());
160                    r
161                } else {
162                    #block
163                }
164            }
165        }
166    };
167
168    let make_gauge = || {
169        quote! {
170            #(#attrs)*
171            #vis #sig {
172
173                static GAUGE: std::sync::LazyLock<Option<metricrs::Gauge>> = std::sync::LazyLock::new(|| {
174                    metricrs::global::get_global_registry().map(|registry| {
175                        use metricrs::*;
176                        use DeriveKind::*;
177                        registry.gauge(DeriveOption {
178                              #(#fields,)*
179                            ..Default::default()
180                        }.into())
181                    })
182                });
183
184                if let Some(gauge) = GAUGE.as_ref() {
185                    gauge.increment(1f64);
186                    let r = #block;
187                    gauge.decrement(1f64);
188                    r
189                } else {
190                    #block
191                }
192            }
193        }
194    };
195
196    if let Some(kind) = kind {
197        match kind.to_token_stream().to_string().as_str() {
198            "Timer" => return make_timer().into(),
199            "Gauge" => return make_gauge().into(),
200            _ => return make_counter().into(),
201        }
202    }
203
204    return make_counter().into();
205}
metricrs_derive/lib.rs

metricrs_derive/
lib.rs
