attributes: fix handling of inner attributes (#2307)

## Motivation

When the `instrument` attribute is used on a function with inner
attributes, the proc macro generates code above the attributes within
the function block that causes compilation errors. These should be
parsed out separately and handled.

Fixes #2294

## Solution

I updated `MaybeItemFn` and `MaybeItemFnRef` to so they hold both the
outer and inner attributes for the instrumented function and updated the
codegen to inlcude them in the appropriate locations.

I couldn't preserve the existing implementation of
`From<&'_ ItemFn> for MaybeItemFnRef<'_, Box<Block>>`, because it is
now necessary to separate the inner and outer attributes of the
`ItemFn` into two separate `Vec`s. That implementation was replaced
with a `From<ItemFn> for MaybeItemFn`, which uses `Iterator::partition`
to separate out the inner and outer attributes.

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

tracing-attributes/src/expand.rs

@@ -11,7 +11,7 @@ use syn::{
 
 use crate::{
     attr::{Field, Fields, FormatMode, InstrumentArgs},
-    MaybeItemFnRef,
+    MaybeItemFn, MaybeItemFnRef,
 };
 
 /// Given an existing function, generate an instrumented version of that function
@@ -25,7 +25,8 @@ pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
     // isn't representable inside a quote!/quote_spanned! macro
     // (Syn's ToTokens isn't implemented for ItemFn)
     let MaybeItemFnRef {
-        attrs,
+        outer_attrs,
+        inner_attrs,
         vis,
         sig,
         block,
@@ -87,10 +88,11 @@ pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
     );
 
     quote!(
-        #(#attrs) *
+        #(#outer_attrs) *
         #vis #constness #unsafety #asyncness #abi fn #ident<#gen_params>(#params) #output
         #where_clause
         {
+            #(#inner_attrs) *
             #warnings
             #body
         }
@@ -657,7 +659,7 @@ impl<'block> AsyncInfo<'block> {
         self,
         args: InstrumentArgs,
         instrumented_function_name: &str,
-    ) -> proc_macro::TokenStream {
+    ) -> Result<proc_macro::TokenStream, syn::Error> {
         // let's rewrite some statements!
         let mut out_stmts: Vec<TokenStream> = self
             .input
@@ -678,12 +680,15 @@ impl<'block> AsyncInfo<'block> {
             // instrument the future by rewriting the corresponding statement
             out_stmts[iter] = match self.kind {
                 // `Box::pin(immediately_invoked_async_fn())`
-                AsyncKind::Function(fun) => gen_function(
-                    fun.into(),
-                    args,
-                    instrumented_function_name,
-                    self.self_type.as_ref(),
-                ),
+                AsyncKind::Function(fun) => {
+                    let fun = MaybeItemFn::from(fun.clone());
+                    gen_function(
+                        fun.as_ref(),
+                        args,
+                        instrumented_function_name,
+                        self.self_type.as_ref(),
+                    )
+                }
                 // `async move { ... }`, optionally pinned
                 AsyncKind::Async {
                     async_expr,
@@ -714,13 +719,13 @@ impl<'block> AsyncInfo<'block> {
         let vis = &self.input.vis;
         let sig = &self.input.sig;
         let attrs = &self.input.attrs;
-        quote!(
+        Ok(quote!(
             #(#attrs) *
             #vis #sig {
                 #(#out_stmts) *
             }
         )
-        .into()
+        .into())
     }
 }
 

tracing-attributes/src/lib.rs

@@ -86,7 +86,7 @@ extern crate proc_macro;
 use proc_macro2::TokenStream;
 use quote::ToTokens;
 use syn::parse::{Parse, ParseStream};
-use syn::{Attribute, Block, ItemFn, Signature, Visibility};
+use syn::{Attribute, ItemFn, Signature, Visibility};
 
 mod attr;
 mod expand;
@@ -587,11 +587,13 @@ fn instrument_precise(
     // check for async_trait-like patterns in the block, and instrument
     // the future instead of the wrapper
     if let Some(async_like) = expand::AsyncInfo::from_fn(&input) {
-        return Ok(async_like.gen_async(args, instrumented_function_name.as_str()));
+        return async_like.gen_async(args, instrumented_function_name.as_str());
     }
 
+    let input = MaybeItemFn::from(input);
+
     Ok(expand::gen_function(
-        (&input).into(),
+        input.as_ref(),
         args,
         instrumented_function_name.as_str(),
         None,
@@ -603,7 +605,8 @@ fn instrument_precise(
 /// which's block is just a `TokenStream` (it may contain invalid code).
 #[derive(Debug, Clone)]
 struct MaybeItemFn {
-    attrs: Vec<Attribute>,
+    outer_attrs: Vec<Attribute>,
+    inner_attrs: Vec<Attribute>,
     vis: Visibility,
     sig: Signature,
     block: TokenStream,
@@ -612,7 +615,8 @@ struct MaybeItemFn {
 impl MaybeItemFn {
     fn as_ref(&self) -> MaybeItemFnRef<'_, TokenStream> {
         MaybeItemFnRef {
-            attrs: &self.attrs,
+            outer_attrs: &self.outer_attrs,
+            inner_attrs: &self.inner_attrs,
             vis: &self.vis,
             sig: &self.sig,
             block: &self.block,
@@ -624,36 +628,50 @@ impl MaybeItemFn {
 /// (just like `ItemFn`, but skips parsing the body).
 impl Parse for MaybeItemFn {
     fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
-        let attrs = input.call(syn::Attribute::parse_outer)?;
+        let outer_attrs = input.call(Attribute::parse_outer)?;
         let vis: Visibility = input.parse()?;
         let sig: Signature = input.parse()?;
+        let inner_attrs = input.call(Attribute::parse_inner)?;
         let block: TokenStream = input.parse()?;
         Ok(Self {
-            attrs,
+            outer_attrs,
+            inner_attrs,
             vis,
             sig,
             block,
         })
     }
 }
 
+impl From<ItemFn> for MaybeItemFn {
+    fn from(
+        ItemFn {
+            attrs,
+            vis,
+            sig,
+            block,
+        }: ItemFn,
+    ) -> Self {
+        let (outer_attrs, inner_attrs) = attrs
+            .into_iter()
+            .partition(|attr| attr.style == syn::AttrStyle::Outer);
+        Self {
+            outer_attrs,
+            inner_attrs,
+            vis,
+            sig,
+            block: block.to_token_stream(),
+        }
+    }
+}
+
 /// A generic reference type for `MaybeItemFn`,
 /// that takes a generic block type `B` that implements `ToTokens` (eg. `TokenStream`, `Block`).
 #[derive(Debug, Clone)]
 struct MaybeItemFnRef<'a, B: ToTokens> {
-    attrs: &'a Vec<Attribute>,
+    outer_attrs: &'a Vec<Attribute>,
+    inner_attrs: &'a Vec<Attribute>,
     vis: &'a Visibility,
     sig: &'a Signature,
     block: &'a B,
 }
-
-impl<'a> From<&'a ItemFn> for MaybeItemFnRef<'a, Box<Block>> {
-    fn from(val: &'a ItemFn) -> Self {
-        MaybeItemFnRef {
-            attrs: &val.attrs,
-            vis: &val.vis,
-            sig: &val.sig,
-            block: &val.block,
-        }
-    }
-}

tracing-attributes/tests/async_fn.rs

@@ -32,6 +32,15 @@ async fn test_ret_impl_trait_err(n: i32) -> Result<impl Iterator<Item = i32>, &'
 #[instrument]
 async fn test_async_fn_empty() {}
 
+// Reproduces a compile error when an instrumented function body contains inner
+// attributes (https://github.com/tokio-rs/tracing/issues/2294).
+#[deny(unused_variables)]
+#[instrument]
+async fn repro_async_2294() {
+    #![allow(unused_variables)]
+    let i = 42;
+}
+
 // Reproduces https://github.com/tokio-rs/tracing/issues/1613
 #[instrument]
 // LOAD-BEARING `#[rustfmt::skip]`! This is necessary to reproduce the bug;

tracing-attributes/tests/instrument.rs

@@ -3,6 +3,15 @@ use tracing::Level;
 use tracing_attributes::instrument;
 use tracing_mock::*;
 
+// Reproduces a compile error when an instrumented function body contains inner
+// attributes (https://github.com/tokio-rs/tracing/issues/2294).
+#[deny(unused_variables)]
+#[instrument]
+fn repro_2294() {
+    #![allow(unused_variables)]
+    let i = 42;
+}
+
 #[test]
 fn override_everything() {
     #[instrument(target = "my_target", level = "debug")]