fix #[async_trait] attribute not being correctly detected when its re-exported from another crate

CHANGELOG.md

@@ -35,6 +35,7 @@ ensure that you match the version that Autometrics uses.**
 - Fixed incorrect duration being recorded when using `#[async_trait]` together with `#[autometrics]` (#161)   
   **Please note that the `#[autometrics]` macro needs to be defined BEFORE `#[async_trait]`.**
 - Fixed value of the `result` label being empty when the function is annotated with `#[async_trait]` (#161)
+- Fixed `#[async_trait]` not being correctly detected if its re-exported in another crate (#164)
 
 ### Autometrics 1.0 compliance
 

autometrics-macros/Cargo.toml

@@ -18,4 +18,5 @@ proc-macro = true
 percent-encoding = "2.2"
 proc-macro2 = "1"
 quote = "1"
+regex = "1.10.2"
 syn =  { version = "2", features = ["full"] }

autometrics-macros/src/lib.rs

@@ -2,6 +2,7 @@ use crate::parse::{AutometricsArgs, Item};
 use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
 use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
+use regex::Regex;
 use std::env;
 use std::str::FromStr;
 use syn::{
@@ -29,7 +30,7 @@ pub fn autometrics(
 
     let result = match item {
         Item::Function(item) => instrument_function(&args, item, &None),
-        Item::Impl(item) => instrument_impl_block(&args, item, async_trait),
+        Item::Impl(item) => instrument_impl_block(&args, item, &async_trait),
     };
 
     let output = match result {
@@ -40,22 +41,21 @@ pub fn autometrics(
     output.into()
 }
 
-fn check_async_trait(input: proc_macro::TokenStream) -> (bool, proc_macro::TokenStream) {
-    let str = input.to_string();
-
-    if str.contains("#[async_trait]") || str.contains("#[async_trait::async_trait]") {
-        // .unwrap is safe because we only remove tokens from the existing stream, we dont add new ones
-        (
-            true,
-            proc_macro::TokenStream::from_str(
-                &str.replace("#[async_trait]", "")
-                    .replace("#[async_trait::async_trait]", ""),
-            )
-            .unwrap(),
-        )
-    } else {
-        (false, input)
-    }
+/// returns a tuple of two containing:
+/// - `async_trait` attributes that have to be re-added after our instrumentation magic has been added
+/// - `input` but without the `async_trait` attributes
+fn check_async_trait(input: proc_macro::TokenStream) -> (String, proc_macro::TokenStream) {
+    let regex = Regex::new(r#"#\[[^\]]*async_trait\]"#).expect("The regex is hardcoded and thus guaranteed to be successfully parseable");
+
+    let original = input.to_string();
+
+    let attributes: Vec<_> = regex.find_iter(&original).map(|m| m.as_str()).collect();
+    let replaced = regex.replace_all(&original, "");
+
+    // .unwrap is safe because we only remove tokens from the existing stream, we dont add new ones
+    let ts = proc_macro::TokenStream::from_str(replaced.as_ref()).unwrap();
+
+    (attributes.join("\n"), ts)
 }
 
 #[proc_macro_derive(ResultLabels, attributes(label))]
@@ -335,16 +335,10 @@ fn instrument_function(
 fn instrument_impl_block(
     args: &AutometricsArgs,
     mut item: ItemImpl,
-    async_trait: bool,
+    attributes_to_re_add: &str,
 ) -> Result<TokenStream> {
     let struct_name = Some(item.self_ty.to_token_stream().to_string());
 
-    let async_trait = if async_trait {
-        quote! { #[async_trait::async_trait] }
-    } else {
-        quote! {}
-    };
-
     // Replace all of the method items in place
     item.items = item
         .items
@@ -379,8 +373,10 @@ fn instrument_impl_block(
         })
         .collect();
 
+    let ts = TokenStream::from_str(attributes_to_re_add)?;
+
     Ok(quote! {
-        #async_trait
+        #ts
         #item
     })
 }