veecle_telemetry_macros/
lib.rs

1// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
2// Copyright 2025 Veecle GmbH.
3//
4// This file has been modified from the original TiKV implementation.
5
6//! An attribute macro designed to eliminate boilerplate code for [`veecle_telemetry`](https://crates.io/crates/veecle_telemetry).
7
8#![recursion_limit = "256"]
9#![cfg_attr(not(feature = "enable"), allow(dead_code))]
10#![cfg_attr(not(feature = "enable"), allow(unreachable_code))]
11
12use std::collections::HashMap;
13
14use proc_macro2::{Ident, Span};
15use quote::{quote, quote_spanned};
16use syn::parse::{Parse, ParseStream};
17use syn::punctuated::Punctuated;
18use syn::*;
19
20struct Arguments {
21    name: Option<LitStr>,
22    short_name: bool,
23    properties: Vec<Property>,
24    veecle_telemetry_crate: Option<syn::Path>,
25    span: Span,
26}
27
28struct Property {
29    key: LitStr,
30    value: Lit,
31    span: Span,
32}
33
34impl Parse for Property {
35    fn parse(input: ParseStream) -> Result<Self> {
36        let key: LitStr = input.parse()?;
37        input.parse::<Token![:]>()?;
38        let value: Lit = input.parse()?;
39
40        // For some reason, `join` fails in doc macros.
41        let span = key.span().join(value.span()).unwrap_or_else(|| key.span());
42        Ok(Property { key, value, span })
43    }
44}
45
46impl Parse for Arguments {
47    fn parse(input: ParseStream) -> Result<Self> {
48        let mut name = None;
49        let mut short_name = false;
50        let mut properties = Vec::<Property>::new();
51        let mut veecle_telemetry_crate = None;
52        let mut seen = HashMap::new();
53
54        while !input.is_empty() {
55            let ident: Ident = input.parse()?;
56            if seen.contains_key(&ident.to_string()) {
57                return Err(Error::new(ident.span(), "duplicate argument"));
58            }
59            seen.insert(ident.to_string(), ());
60            input.parse::<Token![=]>()?;
61            match ident.to_string().as_str() {
62                "name" => {
63                    let parsed_name: LitStr = input.parse()?;
64                    name = Some(parsed_name);
65                }
66                "short_name" => {
67                    let parsed_short_name: LitBool = input.parse()?;
68                    short_name = parsed_short_name.value;
69                }
70                "properties" => {
71                    let content;
72                    let _brace_token = braced!(content in input);
73                    let property_list = content.parse_terminated(Property::parse, Token![,])?;
74                    for property in property_list {
75                        if properties
76                            .iter()
77                            .any(|existing| existing.key == property.key)
78                        {
79                            return Err(Error::new(Span::call_site(), "duplicate property key"));
80                        }
81                        properties.push(property);
82                    }
83                }
84                "crate" => {
85                    let crate_path: syn::Path = input.parse()?;
86                    veecle_telemetry_crate = Some(crate_path);
87                }
88                _ => return Err(Error::new(Span::call_site(), "unexpected identifier")),
89            }
90            if !input.is_empty() {
91                let _ = input.parse::<Token![,]>();
92            }
93        }
94
95        Ok(Arguments {
96            name,
97            short_name,
98            properties,
99            veecle_telemetry_crate,
100            span: input.span(),
101        })
102    }
103}
104
105/// An attribute macro designed to eliminate boilerplate code.
106///
107/// This macro automatically creates a span for the annotated function. The span name defaults to
108/// the function name but can be customized by passing a string literal as an argument using the
109/// `name` parameter.
110///
111/// The `#[trace]` attribute requires a local parent context to function correctly. Ensure that
112/// the function annotated with `#[trace]` is called within __a local context of a `Span`__, which
113/// is established by invoking the `Span::set_local_parent()` method.
114///
115/// ## Arguments
116///
117/// * `name` - The name of the span. Defaults to the full path of the function.
118/// * `short_name` - Whether to use the function name without path as the span name. Defaults to `false`.
119/// * `properties` - A list of key-value pairs to be added as properties to the span. The value can be a format string,
120///   where the function arguments are accessible. Defaults to `{}`.
121///
122/// # Examples
123///
124/// ```
125/// use veecle_telemetry::instrument;
126///
127/// #[veecle_telemetry::instrument]
128/// fn simple() {
129///     // ...
130/// }
131///
132/// #[veecle_telemetry::instrument(short_name = true)]
133/// async fn simple_async() {
134///     // ...
135/// }
136///
137/// #[veecle_telemetry::instrument(properties = { "k1": "v1", "a": 2 })]
138/// async fn properties(a: u64) {
139///     // ...
140/// }
141/// ```
142///
143/// The code snippets above will be expanded to:
144///
145/// ```
146/// # extern crate alloc;
147/// # use veecle_telemetry::Span;
148/// # use veecle_telemetry::value::KeyValue;
149///
150/// fn simple() {
151///     let __guard__ = Span::new("example::simple", &[]).entered();
152///     // ...
153/// }
154///
155/// async fn simple_async() {
156///     veecle_telemetry::future::FutureExt::with_span(
157///         async move {
158///             // ...
159///         },
160///         veecle_telemetry::Span::new("simple_async", &[]),
161///     )
162///     .await
163/// }
164///
165/// async fn properties(a: u64) {
166///     veecle_telemetry::future::FutureExt::with_span(
167///         async move {
168///             // ...
169///         },
170///         veecle_telemetry::Span::new("example::properties", &[
171///             KeyValue::new("k1", "v1"),
172///             KeyValue::new("a", 2),
173///         ]),
174///     )
175///     .await
176/// }
177/// ```
178#[proc_macro_attribute]
179pub fn instrument(
180    arguments: proc_macro::TokenStream,
181    item: proc_macro::TokenStream,
182) -> proc_macro::TokenStream {
183    #[cfg(not(feature = "enable"))]
184    {
185        let _ = parse_macro_input!(arguments as Arguments);
186        return item;
187    }
188
189    let arguments = parse_macro_input!(arguments as Arguments);
190    let input = parse_macro_input!(item as ItemFn);
191
192    let veecle_telemetry_crate = arguments
193        .veecle_telemetry_crate
194        .clone()
195        .map(Ok)
196        .unwrap_or_else(veecle_telemetry_path);
197    let veecle_telemetry_crate = match veecle_telemetry_crate {
198        Ok(path) => path,
199        Err(error) => return error.to_compile_error().into(),
200    };
201
202    let function_name = &input.sig.ident;
203
204    // Check for async_trait-like patterns in the block, and instrument the future instead of the wrapper.
205    let function_body = match generate_block(
206        function_name,
207        &input.block,
208        input.sig.asyncness.is_some(),
209        &arguments,
210        &veecle_telemetry_crate,
211    ) {
212        Ok(body) => body,
213        Err(error) => return error.to_compile_error().into(),
214    };
215
216    let ItemFn {
217        attrs, vis, sig, ..
218    } = input;
219
220    let Signature {
221        output: return_type,
222        inputs: params,
223        unsafety,
224        constness,
225        abi,
226        ident,
227        asyncness,
228        generics:
229            Generics {
230                params: gen_params,
231                where_clause,
232                ..
233            },
234        ..
235    } = sig;
236
237    quote::quote!(
238        #(#attrs) *
239        #vis #constness #unsafety #asyncness #abi fn #ident<#gen_params>(#params) #return_type
240        #where_clause
241        {
242            #function_body
243        }
244    )
245    .into()
246}
247
248fn generate_name(
249    function_name: &Ident,
250    arguments: &Arguments,
251    async_closure: bool,
252    veecle_telemetry_crate: &syn::Path,
253) -> syn::Result<proc_macro2::TokenStream> {
254    let span = function_name.span();
255    if let Some(name) = &arguments.name {
256        if name.value().is_empty() {
257            return Err(Error::new(span, "`name` can not be empty"));
258        }
259
260        if arguments.short_name {
261            return Err(Error::new(
262                Span::call_site(),
263                "`name` and `short_name` can not be used together",
264            ));
265        }
266
267        Ok(quote_spanned!(span=>
268            #name
269        ))
270    } else if arguments.short_name {
271        let function_name = function_name.to_string();
272        Ok(quote_spanned!(span=>
273            #function_name
274        ))
275    } else {
276        Ok(quote_spanned!(span=>
277            #veecle_telemetry_crate::macro_helpers::strip_closure_suffix(core::any::type_name_of_val(&|| {}), #async_closure)
278        ))
279    }
280}
281
282fn generate_properties(
283    arguments: &Arguments,
284    veecle_telemetry_crate: &syn::Path,
285) -> proc_macro2::TokenStream {
286    if arguments.properties.is_empty() {
287        return quote::quote!(&[]);
288    }
289
290    let span = arguments.span;
291    let properties = arguments
292        .properties
293        .iter()
294        .map(|Property { key, value, span }| {
295            quote_spanned!(*span=>
296                #veecle_telemetry_crate::value::KeyValue::new(#key, #value)
297            )
298        });
299    let properties = Punctuated::<_, Token![,]>::from_iter(properties);
300    quote_spanned!(span=>
301        &[ #properties ]
302    )
303}
304
305/// Instrument a block
306fn generate_block(
307    func_name: &Ident,
308    block: &Block,
309    async_context: bool,
310    arguments: &Arguments,
311    veecle_telemetry_crate: &syn::Path,
312) -> syn::Result<proc_macro2::TokenStream> {
313    let name = generate_name(func_name, arguments, async_context, veecle_telemetry_crate)?;
314    let properties = generate_properties(arguments, veecle_telemetry_crate);
315
316    // Generate the instrumented function body.
317    // If the function is an `async fn`, this will wrap it in an async block.
318    // Otherwise, this will enter the span and then perform the rest of the body.
319    if async_context {
320        Ok(quote!(
321            #veecle_telemetry_crate::future::FutureExt::with_span(
322                async move { #block },
323                #veecle_telemetry_crate::Span::new(#name, #properties),
324            ).await
325        ))
326    } else {
327        Ok(quote!(
328            let __guard__= #veecle_telemetry_crate::Span::new(#name, #properties).entered();
329            #block
330        ))
331    }
332}
333
334/// Returns a path to the `veecle_telemetry` crate for use when macro users don't set it explicitly.
335fn veecle_telemetry_path() -> syn::Result<syn::Path> {
336    proc_macro_crate::crate_name("veecle-telemetry")
337        .map(|found| match found {
338            proc_macro_crate::FoundCrate::Itself => {
339                // The only place we use `veecle-telemetry` within "itself" is doc-tests, where it needs to be an external
340                // path anyway.
341                syn::parse_quote!(::veecle_telemetry)
342            }
343            proc_macro_crate::FoundCrate::Name(name) => {
344                let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
345                syn::parse_quote!(::#ident)
346            }
347        })
348        .or_else(|_| {
349            proc_macro_crate::crate_name("veecle-os").map(|found| match found {
350                proc_macro_crate::FoundCrate::Itself => {
351                    todo!("unused currently, not sure what behavior will be wanted")
352                }
353                proc_macro_crate::FoundCrate::Name(name) => {
354                    let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
355                    syn::parse_quote!(::#ident::telemetry)
356                }
357            })
358        })
359        .map_err(|_| {
360            syn::Error::new(
361                proc_macro2::Span::call_site(),
362                "could not find either veecle-telemetry or veecle-os crates",
363            )
364        })
365}
veecle_telemetry_macros/lib.rs

veecle_telemetry_macros/
lib.rs
