zipkin_macros/
lib.rs

1//  Copyright 2020 Palantir Technologies, Inc.
2//
3//  Licensed under the Apache License, Version 2.0 (the "License");
4//  you may not use this file except in compliance with the License.
5//  You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14//! Macros for use with `zipkin`.
15//!
16//! You should not depend on this crate directly.
17extern crate proc_macro;
18
19use proc_macro::TokenStream;
20use proc_macro2::Span;
21use quote::{quote, ToTokens};
22use syn::parse::{Parse, ParseStream};
23use syn::punctuated::Punctuated;
24use syn::{parse_macro_input, Error, Expr, ImplItemFn, Lit, LitStr, Meta, Stmt, Token};
25
26/// Wraps the execution of a function or method in a span.
27///
28/// Both normal and `async` methods and functions are supported. The name of the span is specified as an argument
29/// to the macro attribute.
30///
31/// Requires the `macros` Cargo feature.
32///
33/// # Examples
34///
35/// ```ignore
36/// #[zipkin::spanned(name = "shave yaks")]
37/// fn shave_some_yaks(yaks: &mut [Yak]) {
38///     // ...
39/// }
40///
41/// #[zipkin::spanned(name = "asynchronously shave yaks")]
42/// async fn shave_some_other_yaks(yaks: &mut [Yak]) {
43///     // ...
44/// }
45///
46/// struct Yak;
47///
48/// impl Yak {
49///     #[zipkin::spanned(name = "shave a yak")]
50///     fn shave(&mut self) {
51///         // ...
52///     }
53///
54///     #[zipkin::spanned(name = "asynchronously shave a yak")]
55///     async fn shave_nonblocking(&mut self) {
56///          // ...
57///     }
58/// }
59/// ```
60#[proc_macro_attribute]
61pub fn spanned(args: TokenStream, item: TokenStream) -> TokenStream {
62    let options = parse_macro_input!(args as Options);
63    let func = parse_macro_input!(item as ImplItemFn);
64
65    spanned_impl(options, func).unwrap_or_else(|e| e.to_compile_error().into())
66}
67
68fn spanned_impl(options: Options, mut func: ImplItemFn) -> Result<TokenStream, Error> {
69    let name = &options.name;
70
71    if func.sig.asyncness.is_some() {
72        let stmts = &func.block.stmts;
73        func.block.stmts = vec![
74            syn::parse2(quote! {
75                let __macro_impl_span = zipkin::next_span()
76                    .with_name(#name)
77                    .detach();
78            })
79            .unwrap(),
80            Stmt::Expr(
81                syn::parse2(quote! {
82                    __macro_impl_span.bind(async move { #(#stmts)* }).await
83                })
84                .unwrap(),
85                None,
86            ),
87        ];
88    } else {
89        let stmt = quote! {
90            let __macro_impl_span = zipkin::next_span().with_name(#name);
91        };
92        func.block.stmts.insert(0, syn::parse2(stmt).unwrap());
93    };
94
95    Ok(func.into_token_stream().into())
96}
97
98struct Options {
99    name: LitStr,
100}
101
102impl Parse for Options {
103    fn parse(input: ParseStream) -> syn::Result<Self> {
104        let args = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;
105
106        let mut name = None;
107
108        for arg in args {
109            let meta = match arg {
110                Meta::NameValue(meta) => meta,
111                _ => return Err(Error::new_spanned(&arg, "invalid attribute syntax")),
112            };
113
114            if meta.path.is_ident("name") {
115                match meta.value {
116                    Expr::Lit(lit) => match lit.lit {
117                        Lit::Str(lit) => name = Some(lit),
118                        lit => return Err(Error::new_spanned(&lit, "expected a string literal")),
119                    },
120                    _ => return Err(Error::new_spanned(meta, "expected `name = \"...\"`")),
121                }
122            } else {
123                return Err(Error::new_spanned(meta.path, "unknown option"));
124            }
125        }
126
127        Ok(Options {
128            name: name.ok_or_else(|| Error::new(Span::call_site(), "missing `name` option"))?,
129        })
130    }
131}
zipkin_macros/lib.rs

zipkin_macros/
lib.rs
