opentelemetry_auto_span/
lib.rs

1mod dig;
2mod handle_sqlx;
3mod line;
4mod utils;
5
6use darling::ast::NestedMeta;
7use darling::{Error, FromMeta};
8use proc_macro2::{Ident, Span, TokenStream};
9use quote::{quote, quote_spanned};
10use syn::{
11    parse_macro_input, spanned::Spanned, visit_mut::VisitMut, Expr, ExprAwait, ExprClosure,
12    ExprTry, ItemFn, Signature,
13};
14
15use crate::{dig::find_source_path, line::LineAccess};
16
17#[derive(Default, FromMeta)]
18#[darling(default)]
19struct Opt {
20    pub debug: bool,
21}
22
23#[proc_macro_attribute]
24pub fn auto_span(
25    attr: proc_macro::TokenStream,
26    item: proc_macro::TokenStream,
27) -> proc_macro::TokenStream {
28    let attr_args = match NestedMeta::parse_meta_list(attr.into()) {
29        Ok(v) => v,
30        Err(e) => {
31            return proc_macro::TokenStream::from(Error::from(e).write_errors());
32        }
33    };
34    let opt = match Opt::from_list(&attr_args) {
35        Ok(v) => v,
36        Err(e) => {
37            return proc_macro::TokenStream::from(e.write_errors());
38        }
39    };
40
41    let mut input = parse_macro_input!(item as ItemFn);
42
43    let mut dir = std::path::PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
44    dir.push("src");
45    let line_access = find_source_path(dir, &input).map(LineAccess::new);
46    let mut visitor = AutoSpanVisitor::new(line_access);
47    visitor.visit_item_fn_mut(&mut input);
48
49    insert_function_span(&mut input);
50    let token = quote! {#input};
51
52    if opt.debug {
53        let mut target = std::path::PathBuf::from(
54            std::env::var("CARGO_TARGET_DIR").unwrap_or_else(|_| "/tmp".to_owned()),
55        );
56        target.push("auto-span");
57        std::fs::create_dir_all(&target).unwrap();
58        target.push(format!("{}.rs", input.sig.ident));
59        std::fs::write(&target, format!("{}", token)).unwrap();
60    }
61
62    token.into()
63}
64
65fn insert_function_span(i: &mut ItemFn) {
66    let def_tracer = quote! {
67        let __otel_auto_tracer = ::opentelemetry::global::tracer("");
68    };
69    let span_ident = Ident::new("span", Span::call_site());
70    let start_tracer = otel_start_tracer_token(&format!("fn:{}", i.sig.ident));
71    let ctx = otel_ctx_token(&span_ident);
72    let stmts = &i.block.stmts;
73    let tokens = if i.sig.asyncness.is_some() {
74        quote! {
75            #def_tracer
76            ::opentelemetry::trace::FutureExt::with_context(
77                async {#(#stmts)*},
78                {
79                    let #span_ident = #start_tracer;
80                    #ctx
81                }
82            ).await
83        }
84    } else {
85        quote! {
86            #def_tracer
87            let #span_ident = #start_tracer;
88            let __otel_auto_ctx = #ctx;
89            let __otel_auto_guard = __otel_auto_ctx.clone().attach();
90            #(#stmts)*
91        }
92    };
93    let body: Expr = syn::parse2(quote! {{#tokens}}).unwrap();
94    match body {
95        Expr::Block(block) => {
96            i.block.stmts = block.block.stmts;
97        }
98        _ => unreachable!(),
99    }
100}
101
102fn otel_start_tracer_token(name: &str) -> TokenStream {
103    quote! {
104        ::opentelemetry::trace::Tracer::start(&__otel_auto_tracer, #name)
105    }
106}
107
108fn otel_ctx_token(span_ident: &Ident) -> TokenStream {
109    quote! {
110        <::opentelemetry::Context as ::opentelemetry::trace::TraceContextExt>::current_with_span(#span_ident)
111    }
112}
113
114struct AutoSpanVisitor {
115    line_access: Option<LineAccess>,
116    context: Vec<ReturnTypeContext>,
117}
118
119#[derive(Copy, Clone)]
120enum ReturnTypeContext {
121    Unknown,
122    Result,
123    Option,
124}
125
126impl AutoSpanVisitor {
127    fn new(line_access: Option<LineAccess>) -> AutoSpanVisitor {
128        AutoSpanVisitor {
129            line_access,
130            context: Vec::new(),
131        }
132    }
133
134    fn push_fn_context(&mut self, sig: &Signature) {
135        let rt = match &sig.output {
136            syn::ReturnType::Default => ReturnTypeContext::Unknown,
137            syn::ReturnType::Type(_, ty) => match ty.as_ref() {
138                syn::Type::Path(path) => {
139                    let name = path.path.segments.last().unwrap().ident.to_string();
140                    if name.contains("Result") {
141                        ReturnTypeContext::Result
142                    } else if name.contains("Option") {
143                        ReturnTypeContext::Option
144                    } else {
145                        ReturnTypeContext::Unknown
146                    }
147                }
148                _ => ReturnTypeContext::Unknown,
149            },
150        };
151        self.context.push(rt);
152    }
153
154    pub fn push_closure_context(&mut self) {
155        self.context.push(ReturnTypeContext::Unknown);
156    }
157
158    pub fn pop_context(&mut self) {
159        self.context.pop();
160    }
161
162    pub fn current_context(&self) -> ReturnTypeContext {
163        *self.context.last().unwrap()
164    }
165
166    fn handle_sqlx(&self, expr_await: &mut ExprAwait) -> bool {
167        let mut visitor = handle_sqlx::SqlxVisitor::new();
168        visitor.visit_expr_await_mut(expr_await);
169        visitor.is_mutate()
170    }
171
172    fn get_line_info(&self, span: Span) -> Option<(i64, String)> {
173        self.line_access.as_ref().and_then(|la| la.span(span))
174    }
175
176    fn span_ident(&self) -> Ident {
177        Ident::new("__otel_auto_span", Span::call_site())
178    }
179}
180
181fn add_line_info(tokens: &mut TokenStream, span_ident: &Ident, line_info: Option<(i64, String)>) {
182    if let Some((lineno, line)) = line_info {
183        tokens.extend(quote! {
184            #span_ident.set_attribute(::opentelemetry::KeyValue::new("code.lineno", #lineno));
185            #span_ident.set_attribute(::opentelemetry::KeyValue::new("code.line", #line));
186        });
187    }
188}
189
190impl VisitMut for AutoSpanVisitor {
191    fn visit_expr_mut(&mut self, i: &mut Expr) {
192        let span = i.span();
193
194        let span_ident = self.span_ident();
195        let new_span = |name, line_info, expr| {
196            let start_tracer = otel_start_tracer_token(name);
197            let current_with_span = otel_ctx_token(&span_ident);
198            let mut tokens = quote! {
199                #[allow(unused_import)]
200                use ::opentelemetry::trace::{Span as _};
201                #[allow(unused_mut)]
202                let mut #span_ident = #start_tracer;
203            };
204            add_line_info(&mut tokens, &span_ident, line_info);
205            let tokens = quote_spanned! {
206                span => {
207                    ::opentelemetry::trace::FutureExt::with_context(
208                        async { #expr },
209                        {
210                            #tokens
211                            #current_with_span
212                        }
213                    ).await
214                }
215            };
216            syn::parse2(tokens).unwrap()
217        };
218
219        match i {
220            Expr::Await(expr) => {
221                if self.handle_sqlx(expr) {
222                    *i = new_span("db", self.get_line_info(span), expr);
223                } else {
224                    syn::visit_mut::visit_expr_await_mut(self, expr);
225                }
226            }
227            _ => syn::visit_mut::visit_expr_mut(self, i),
228        };
229    }
230
231    fn visit_expr_closure_mut(&mut self, i: &mut ExprClosure) {
232        self.push_closure_context();
233        syn::visit_mut::visit_expr_closure_mut(self, i);
234        self.pop_context();
235    }
236
237    fn visit_expr_try_mut(&mut self, i: &mut ExprTry) {
238        syn::visit_mut::visit_expr_try_mut(self, i);
239
240        if let ReturnTypeContext::Result = self.current_context() {
241            let span_ident = self.span_ident();
242            let span = i.expr.span();
243            let inner = i.expr.as_ref();
244            let mut tokens = quote! {
245                #span_ident.set_status(::opentelemetry::trace::Status::error(format!("{}", e)));
246            };
247            add_line_info(&mut tokens, &span_ident, self.get_line_info(span));
248            i.expr = Box::new(
249                syn::parse2(quote_spanned! {
250                    span => #inner.map_err(|e| {
251                        ::opentelemetry::trace::get_active_span(|__otel_auto_span| {
252                            #tokens
253                        });
254                        e
255                    })
256                })
257                .unwrap(),
258            );
259        }
260    }
261
262    fn visit_item_fn_mut(&mut self, i: &mut ItemFn) {
263        self.push_fn_context(&i.sig);
264        if self.context.len() == 1 {
265            // skip inner function, because `span` is not shared
266            syn::visit_mut::visit_item_fn_mut(self, i);
267        }
268        self.pop_context();
269    }
270}
opentelemetry_auto_span/lib.rs

opentelemetry_auto_span/
lib.rs
