attributes: support adding arbitrary fields to the instrument macro

tracing-attributes/src/lib.rs

@@ -62,15 +62,15 @@
 #![allow(unused)]
 extern crate proc_macro;
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::iter;
 
 use proc_macro::TokenStream;
 use quote::{quote, quote_spanned, ToTokens};
 use syn::{
     spanned::Spanned, AttributeArgs, FieldPat, FnArg, Ident, ItemFn, Lit, LitInt, Meta, MetaList,
     MetaNameValue, NestedMeta, Pat, PatIdent, PatReference, PatStruct, PatTuple, PatTupleStruct,
-    PatType, Signature,
+    PatType, Path, Signature,
 };
 
 /// Instruments a function to create and enter a `tracing` [span] every time
@@ -193,6 +193,12 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream {
         })
         .filter(|ident| !skips.contains(ident))
         .collect();
+
+    let fields = match fields(&args, &param_names) {
+        Ok(fields) => fields,
+        Err(err) => return quote!(#err).into(),
+    };
+
     let param_names_clone = param_names.clone();
 
     // Generate the instrumented function body.
@@ -225,6 +231,19 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream {
     let target = target(&args);
     let span_name = name(&args, ident_str);
 
+    let mut quoted_fields: Vec<_> = param_names
+        .into_iter()
+        .map(|i| quote!(#i = tracing::field::debug(&#i)))
+        .collect();
+    quoted_fields.extend(fields.into_iter().map(|(key, value)| {
+        let value = match value {
+            Some(value) => quote!(#value),
+            None => quote!(tracing::field::Empty),
+        };
+
+        quote!(#key = #value)
+    }));
+
     quote!(
         #(#attrs) *
         #vis #constness #unsafety #asyncness #abi fn #ident<#gen_params>(#params) #return_type
@@ -234,7 +253,7 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream {
                 target: #target,
                 #level,
                 #span_name,
-                #(#param_names = tracing::field::debug(&#param_names_clone)),*
+                #(#quoted_fields),*
             );
             #body
         }
@@ -350,22 +369,22 @@ fn level(args: &[NestedMeta]) -> impl ToTokens {
 }
 
 fn target(args: &[NestedMeta]) -> impl ToTokens {
-    let mut levels = args.iter().filter_map(|arg| match arg {
+    let mut targets = args.iter().filter_map(|arg| match arg {
         NestedMeta::Meta(Meta::NameValue(MetaNameValue {
             ref path, ref lit, ..
         })) if path.is_ident("target") => Some(lit.clone()),
         _ => None,
     });
-    let level = levels.next();
+    let target = targets.next();
 
-    // If we found more than one arg named "level", that's a syntax error...
-    if let Some(lit) = levels.next() {
+    // If we found more than one arg named "target", that's a syntax error...
+    if let Some(lit) = targets.next() {
         return quote_spanned! {lit.span()=>
             compile_error!("expected only a single `target` argument!")
         };
     }
 
-    match level {
+    match target {
         Some(Lit::Str(ref lit)) => quote!(#lit),
         Some(lit) => quote_spanned! {lit.span()=>
             compile_error!(
@@ -376,6 +395,93 @@ fn target(args: &[NestedMeta]) -> impl ToTokens {
     }
 }
 
+fn fields(
+    args: &[NestedMeta],
+    param_names: &[Ident],
+) -> Result<(Vec<(Ident, Option<Lit>)>), impl ToTokens> {
+    let mut fields = args.iter().filter_map(|arg| match arg {
+        NestedMeta::Meta(Meta::List(MetaList {
+            ref path,
+            ref nested,
+            ..
+        })) if path.is_ident("fields") => Some(nested.clone()),
+        _ => None,
+    });
+    let field_holder = fields.next();
+
+    // If we found more than one arg named "fields", that's a syntax error...
+    if let Some(lit) = fields.next() {
+        return Err(quote_spanned! {lit.span()=>
+            compile_error!("expected only a single `fields` argument!")
+        });
+    }
+
+    match field_holder {
+        Some(fields) => {
+            let mut parsed = Vec::default();
+            let mut visited_keys: HashSet<String> = Default::default();
+            let param_set: HashSet<String> = param_names.iter().map(|i| i.to_string()).collect();
+            for field in fields.into_iter() {
+                let (key, value) = match field {
+                    NestedMeta::Meta(meta) => match meta {
+                        Meta::NameValue(kv) => (kv.path, Some(kv.lit)),
+                        Meta::Path(path) => (path, None),
+                        _ => {
+                            return Err(quote_spanned! {meta.span()=>
+                                compile_error!("each field must be a key with an optional value. Keys must be valid Rust identifiers (nested keys with dots are not supported).")
+                            })
+                        }
+                    },
+                    _ => {
+                        return Err(quote_spanned! {field.span()=>
+                            compile_error!("`fields` argument should be a list of key-value fields")
+                        })
+                    }
+                };
+
+                let key = match key.get_ident() {
+                    Some(key) => key,
+                    None => {
+                        return Err(quote_spanned! {key.span()=>
+                            compile_error!("field keys must be valid Rust identifiers (nested keys with dots are not supported).")
+                        })
+                    }
+                };
+
+                let key_str = key.to_string();
+                if param_set.contains(&key_str) {
+                    return Err(quote_spanned! {key.span()=>
+                        compile_error!("field overlaps with (non-skipped) parameter name")
+                    });
+                }
+
+                if visited_keys.contains(&key_str) {
+                    return Err(quote_spanned! {key.span()=>
+                        compile_error!("each field key must appear at most once")
+                    });
+                } else {
+                    visited_keys.insert(key_str);
+                }
+
+                if let Some(literal) = &value {
+                    match literal {
+                        Lit::Bool(_) | Lit::Str(_) | Lit::Int(_) => {}
+                        _ => {
+                            return Err(quote_spanned! {literal.span()=>
+                                compile_error!("values can be only strings, integers or booleans")
+                            })
+                        }
+                    }
+                }
+
+                parsed.push((key.clone(), value));
+            }
+            Ok(parsed)
+        }
+        None => Ok(Default::default()),
+    }
+}
+
 fn name(args: &[NestedMeta], default_name: String) -> impl ToTokens {
     let mut names = args.iter().filter_map(|arg| match arg {
         NestedMeta::Meta(Meta::NameValue(MetaNameValue {

tracing-attributes/tests/fields.rs

@@ -0,0 +1,63 @@
+mod support;
+use support::*;
+
+use crate::support::field::mock;
+use crate::support::span::NewSpan;
+use tracing::subscriber::with_default;
+use tracing_attributes::instrument;
+
+#[instrument(fields(foo = "bar", dsa = true, num = 1))]
+fn fn_no_param() {}
+
+#[instrument(fields(foo = "bar"))]
+fn fn_param(param: u32) {}
+
+#[instrument(fields(foo = "bar", empty))]
+fn fn_empty_field() {}
+
+#[test]
+fn fields() {
+    let span = span::mock().with_field(
+        mock("foo")
+            .with_value(&"bar")
+            .and(mock("dsa").with_value(&true))
+            .and(mock("num").with_value(&1))
+            .only(),
+    );
+    run_test(span, || {
+        fn_no_param();
+    });
+}
+
+#[test]
+fn parameters_with_fields() {
+    let span = span::mock().with_field(
+        mock("foo")
+            .with_value(&"bar")
+            .and(mock("param").with_value(&format_args!("1")))
+            .only(),
+    );
+    run_test(span, || {
+        fn_param(1);
+    });
+}
+
+#[test]
+fn empty_field() {
+    let span = span::mock().with_field(mock("foo").with_value(&"bar").only());
+    run_test(span, || {
+        fn_empty_field();
+    });
+}
+
+fn run_test<F: FnOnce() -> T, T>(span: NewSpan, fun: F) {
+    let (subscriber, handle) = subscriber::mock()
+        .new_span(span)
+        .enter(span::mock())
+        .exit(span::mock())
+        .done()
+        .run_with_handle();
+
+    with_default(subscriber, fun);
+    handle.assert_finished();
+}