Skip to content

Commit

Permalink
add full support for async-trait (#161)
Browse files Browse the repository at this point in the history
* add full support for async-trait

* move async trait support from compile fail to pass

* update changelog

* update changelog again
  • Loading branch information
mellowagain authored Nov 30, 2023
1 parent 1af33f2 commit 69e64f8
Show file tree
Hide file tree
Showing 4 changed files with 34 additions and 52 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@ ensure that you match the version that Autometrics uses.**
has been improved (#149)
- Fixed missing feature flag for `opentelemetry-otlp` when autometrics feature
`otel-push-exporter` is enabled
- 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)

### Autometrics 1.0 compliance

Expand Down
48 changes: 31 additions & 17 deletions autometrics-macros/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
use crate::parse::{AutometricsArgs, Item};
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use proc_macro2::{Span, TokenStream};
use proc_macro2::TokenStream;
use quote::{quote, ToTokens};
use std::env;
use std::str::FromStr;
use syn::{
parse_macro_input, GenericArgument, ImplItem, ItemFn, ItemImpl, PathArguments, Result,
ReturnType, Type,
Expand All @@ -21,19 +22,14 @@ pub fn autometrics(
args: proc_macro::TokenStream,
item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
let args = parse_macro_input!(args as parse::AutometricsArgs);

// if the input has yet to be modified by the `async-trait` crate, we will fail parsing in
// in the next line, so lets error early now
if let Err(err) = check_async_trait(&item) {
return err.into_compile_error().into();
}
let args = parse_macro_input!(args as AutometricsArgs);

let (async_trait, item) = check_async_trait(item);
let item = parse_macro_input!(item as Item);

let result = match item {
Item::Function(item) => instrument_function(&args, item, &None),
Item::Impl(item) => instrument_impl_block(&args, item),
Item::Impl(item) => instrument_impl_block(&args, item, async_trait),
};

let output = match result {
Expand All @@ -44,16 +40,21 @@ pub fn autometrics(
output.into()
}

fn check_async_trait(input: &proc_macro::TokenStream) -> Result<()> {
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]") {
Err(syn::Error::new(
Span::call_site(),
"#[async_trait] must be defined BEFORE #[autometrics]",
))
// .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 {
Ok(())
(false, input)
}
}

Expand Down Expand Up @@ -331,9 +332,19 @@ fn instrument_function(
}

/// Add autometrics instrumentation to an entire impl block
fn instrument_impl_block(args: &AutometricsArgs, mut item: ItemImpl) -> Result<TokenStream> {
fn instrument_impl_block(
args: &AutometricsArgs,
mut item: ItemImpl,
async_trait: bool,
) -> 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
Expand Down Expand Up @@ -368,7 +379,10 @@ fn instrument_impl_block(args: &AutometricsArgs, mut item: ItemImpl) -> Result<T
})
.collect();

Ok(quote! { #item })
Ok(quote! {
#async_trait
#item
})
}

/// Create Prometheus queries for the generated metric and
Expand Down
35 changes: 0 additions & 35 deletions autometrics/tests/result_labels/fail/async_trait_support.stderr

This file was deleted.

0 comments on commit 69e64f8

Please sign in to comment.