mz_ore_proc/
instrument.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
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 in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! instrument macro with improved default safety
17//!
18//! This wraps the `tracing::instrument` macro and:
19//! - adds `skip_all`
20//! - errors on `skip`
21//!
22//! Its purpose is to prevent accidentally including large function arguments in tracing spans. By
23//! enforcing the use of skip_all, users must use the `fields` argument of the `tracing::instrument`
24//! macro to manually select their desired fields.
25
26use proc_macro::TokenStream;
27use proc_macro2::{TokenStream as TokenStream2, TokenTree};
28use quote::quote;
29
30/// Implementation for the `#[instrument]` macro.
31pub fn instrument_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
32    let attr = TokenStream2::from(attr);
33    let item = TokenStream2::from(item);
34
35    // syn appears to not be able to parse the `%` part of things like `#[instrument(fields(shard =
36    // %id))]`, so we use the more naive proc_macro crate and look for strings.
37    let mut iter = attr.into_iter();
38    let mut args: TokenStream2 = quote! { skip_all };
39    while let Some(tok) = iter.next() {
40        match &tok {
41            TokenTree::Ident(ident) => match ident.to_string().as_str() {
42                "skip_all" => panic!("skip_all already included; remove it"),
43                "skip" => panic!("skip prohibited; use fields"),
44                _ => {}
45            },
46            _ => {}
47        }
48        args.extend([tok])
49    }
50    quote! {
51        #[allow(clippy::disallowed_macros)]
52        #[::tracing::instrument(
53            #args
54        )]
55        #item
56    }
57    .into()
58}
mz_ore_proc/instrument.rs

mz_ore_proc/
instrument.rs
