Allow specifying result labels for enum variants

.gitignore

@@ -1,3 +1,6 @@
 target
 Cargo.lock
 data
+
+# Trybuild compilation error outputs in development phase
+wip

CHANGELOG.md

@@ -10,19 +10,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
--
+
+- `ResultLabels` derive macro allows to specify on an enum whether variants should
+  always be "ok", or "error" for the success rate metrics of functions using them. (#61)
 
 ### Changed
--
+
+- `GetLabels` trait (publicly exported but meant for internal use) changed the signature
+  of its function to accomodate the new `ResultLabels` macro. This change is not significant
+  if you never imported `autometrics::__private` manually (#61)
 
 ### Deprecated
 -
 
 ### Removed
--
+
+- `GetLabelsForResult` trait (publicly exported but meant for internal use) was removed
+  to accomodate the new `ResultLabels` macro. This change is not significant
+  if you never imported `autometrics::__private` manually (#61)
 
 ### Fixed
--
+
+- `#[autometrics]` now works on functions that use type inference in their return statement
+  (#74, #61)
 
 ### Security
 -

autometrics-macros/src/lib.rs

@@ -3,9 +3,10 @@ use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
 use proc_macro2::TokenStream;
 use quote::quote;
 use std::env;
-use syn::{parse_macro_input, ImplItem, ItemFn, ItemImpl, Result};
+use syn::{parse_macro_input, ImplItem, ItemFn, ItemImpl, Result, ReturnType, Type};
 
 mod parse;
+mod result_labels;
 
 const COUNTER_NAME_PROMETHEUS: &str = "function_calls_count";
 const HISTOGRAM_BUCKET_NAME_PROMETHEUS: &str = "function_calls_duration_bucket";
@@ -36,6 +37,14 @@ pub fn autometrics(
     output.into()
 }
 
+#[proc_macro_derive(ResultLabels, attributes(label))]
+pub fn result_labels(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let input = parse_macro_input!(input as syn::DeriveInput);
+    result_labels::expand(input)
+        .unwrap_or_else(syn::Error::into_compile_error)
+        .into()
+}
+
 /// Add autometrics instrumentation to a single function
 fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStream> {
     let sig = item.sig;
@@ -51,6 +60,36 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
     // Build the documentation we'll add to the function's RustDocs
     let metrics_docs = create_metrics_docs(&prometheus_url, &function_name, args.track_concurrency);
 
+    // Type annotation to allow type inference to work on return expressions (such as `.collect()`), as
+    // well as prevent compiler type-inference from selecting the wrong branch in the `spez` macro later.
+    //
+    // Type inference can make the compiler select one of the early cases of `autometrics::result_labels!`
+    // even if the types `T` or `E` do not implement the `GetLabels` trait. That leads to a compilation error
+    // looking like this:
+    // ```
+    // error[E0277]: the trait bound `ApiError: GetLabels` is not satisfied
+    //  --> examples/full-api/src/routes.rs:48:1
+    //   |
+    //48 | #[autometrics(objective = API_SLO)]
+    //   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `GetLabels` is not implemented for `ApiError`
+    //   |
+    //   = help: the trait `create_user::{closure#0}::Match2` is implemented for `&&&&create_user::{closure#0}::Match<&Result<T, E>>`
+    //note: required for `&&&&create_user::{closure#0}::Match<&Result<Json<User>, ApiError>>` to implement `create_user::{closure#0}::Match2`
+    //  --> examples/full-api/src/routes.rs:48:1
+    //   |
+    //48 | #[autometrics(objective = API_SLO)]
+    //   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    //   = note: this error originates in the macro `$crate::__private::spez` which comes from the expansion of the attribute macro `autometrics` (in Nightly builds, run with -Z macro-backtrace for more info)
+    // ```
+    //
+    // specifying the return type makes the compiler select the (correct) fallback case of `ApiError` not being a
+    // `GetLabels` implementor.
+    let return_type = match sig.output {
+        ReturnType::Default => quote! { : () },
+        ReturnType::Type(_, ref t) if matches!(t.as_ref(), &Type::ImplTrait(_)) => quote! {},
+        ReturnType::Type(_, ref t) => quote! { : #t },
+    };
+
     // Wrap the body of the original function, using a slightly different approach based on whether the function is async
     let call_function = if sig.asyncness.is_some() {
         quote! {
@@ -97,12 +136,10 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
             }
         }
     } else {
-        // This will use the traits defined in the `labels` module to determine if
-        // the return value was a `Result` and, if so, assign the appropriate labels
         quote! {
             {
-                use autometrics::__private::{CALLER, CounterLabels, GetLabels, GetLabelsFromResult};
-                let result_labels = (&result).__autometrics_get_labels();
+                use autometrics::__private::{CALLER, CounterLabels, GetLabels};
+                let result_labels = autometrics::get_result_labels_for_value!(&result);
                 CounterLabels::new(
                     #function_name,
                     module_path!(),
@@ -137,7 +174,7 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
                 AutometricsTracker::start(#gauge_labels)
             };
 
-            let result = #call_function;
+            let result #return_type = #call_function;
 
             {
                 use autometrics::__private::{HistogramLabels, TrackMetrics};

autometrics-macros/src/result_labels.rs

@@ -0,0 +1,163 @@
+//! The definition of the ResultLabels derive macro, see
+//! autometrics::ResultLabels for more information.
+
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{
+    punctuated::Punctuated, token::Comma, Attribute, Data, DataEnum, DeriveInput, Error, Ident,
+    Lit, LitStr, Result, Variant,
+};
+
+// These labels must match autometrics::ERROR_KEY and autometrics::OK_KEY,
+// to avoid a dependency loop just for 2 constants we recreate these here.
+const OK_KEY: &str = "ok";
+const ERROR_KEY: &str = "error";
+const RESULT_KEY: &str = "result";
+const ATTR_LABEL: &str = "label";
+const ACCEPTED_LABELS: [&str; 2] = [ERROR_KEY, OK_KEY];
+
+/// Entry point of the ResultLabels macro
+pub(crate) fn expand(input: DeriveInput) -> Result<TokenStream> {
+    let variants = match &input.data {
+        Data::Enum(DataEnum { variants, .. }) => variants,
+        _ => {
+            return Err(Error::new_spanned(
+                input,
+                "ResultLabels only works with 'Enum's.",
+            ))
+        }
+    };
+    let enum_name = &input.ident;
+    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+    let conditional_clauses_for_labels = conditional_label_clauses(variants, enum_name)?;
+
+    Ok(quote! {
+        #[automatically_derived]
+        impl #impl_generics ::autometrics::__private::GetLabels for #enum_name #ty_generics #where_clause {
+            fn __autometrics_get_labels(&self) -> Option<&'static str> {
+                #conditional_clauses_for_labels
+            }
+        }
+    })
+}
+
+/// Build the list of match clauses for the generated code.
+fn conditional_label_clauses(
+    variants: &Punctuated<Variant, Comma>,
+    enum_name: &Ident,
+) -> Result<TokenStream> {
+    let clauses: Vec<TokenStream> = variants
+        .iter()
+        .map(|variant| {
+            let variant_name = &variant.ident;
+            let variant_matcher: TokenStream = match variant.fields {
+                syn::Fields::Named(_) => quote! { #variant_name {..} },
+                syn::Fields::Unnamed(_) => quote! { #variant_name (_) },
+                syn::Fields::Unit => quote! { #variant_name },
+            };
+            if let Some(key) = extract_label_attribute(&variant.attrs)? {
+                Ok(quote! [
+                    else if ::std::matches!(self, & #enum_name :: #variant_matcher) {
+                       Some(#key)
+                    }
+                ])
+            } else {
+                // Let the code flow through the last value
+                Ok(quote! {})
+            }
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    Ok(quote! [
+        if false {
+            None
+        }
+        #(#clauses)*
+        else {
+            None
+        }
+    ])
+}
+
+/// Extract the wanted label from the annotation in the variant, if present.
+/// The function looks for `#[label(result = "ok")]` kind of labels.
+///
+/// ## Error cases
+///
+/// The function will error out with the smallest possible span when:
+///
+/// - The attribute on a variant is not a "list" type (so `#[label]` is not allowed),
+/// - The key in the key value pair is not "result", as it's the only supported keyword
+///   for now (so `#[label(non_existing_label = "ok")]` is not allowed),
+/// - The value for the "result" label is not in the autometrics supported set (so
+///   `#[label(result = "random label that will break queries")]` is not allowed)
+fn extract_label_attribute(attrs: &[Attribute]) -> Result<Option<LitStr>> {
+    attrs
+            .iter()
+            .find_map(|att| match att.parse_meta() {
+                Ok(meta) => match &meta {
+                    syn::Meta::List(list) => {
+                        // Ignore attribute if it's not `label(...)`
+                        if list.path.segments.len() != 1 || list.path.segments[0].ident != ATTR_LABEL {
+                            return None;
+                        }
+
+                        // Only lists are allowed
+                        let pair = match list.nested.first() {
+                            Some(syn::NestedMeta::Meta(syn::Meta::NameValue(pair))) => pair,
+                            _ => return Some(Err(Error::new_spanned(
+                            meta,
+                            format!("Only `{ATTR_LABEL}({RESULT_KEY} = \"RES\")` (RES can be {OK_KEY:?} or {ERROR_KEY:?}) is supported"),
+                            ))),
+                        };
+
+                        // Inside list, only 'result = ...' are allowed
+                        if pair.path.segments.len() != 1 || pair.path.segments[0].ident != RESULT_KEY {
+                            return Some(Err(Error::new_spanned(
+                                pair.path.clone(),
+                            format!("Only `{RESULT_KEY} = \"RES\"` (RES can be {OK_KEY:?} or {ERROR_KEY:?}) is supported"),
+                            )));
+                        }
+
+                        // Inside 'result = val', 'val' must be a string literal
+                        let lit_str = match pair.lit {
+                            Lit::Str(ref lit_str) => lit_str,
+                            _ => {
+                            return Some(Err(Error::new_spanned(
+                                &pair.lit,
+                            format!("Only {OK_KEY:?} or {ERROR_KEY:?}, as string literals, are accepted as result values"),
+                            )));
+                        }
+                        };
+
+                        // Inside 'result = val', 'val' must be one of the allowed string literals
+                        if !ACCEPTED_LABELS.contains(&lit_str.value().as_str()) {
+                            return Some(Err(Error::new_spanned(
+                                    lit_str,
+                            format!("Only {OK_KEY:?} or {ERROR_KEY:?} are accepted as result values"),
+                            )));
+                        }
+
+                        Some(Ok(lit_str.clone()))
+                    },
+                    syn::Meta::NameValue(nv) if nv.path.segments.len() == 1 && nv.path.segments[0].ident == ATTR_LABEL => {
+                        Some(Err(Error::new_spanned(
+                            nv,
+                            format!("Only `{ATTR_LABEL}({RESULT_KEY} = \"RES\")` (RES can be {OK_KEY:?} or {ERROR_KEY:?}) is supported"),
+                        )))
+                    },
+                    syn::Meta::Path(p) if p.segments.len() == 1 && p.segments[0].ident == ATTR_LABEL => {
+                        Some(Err(Error::new_spanned(
+                            p,
+                            format!("Only `{ATTR_LABEL}({RESULT_KEY} = \"RES\")` (RES can be {OK_KEY:?} or {ERROR_KEY:?}) is supported"),
+                        )))
+                    },
+                    _ => None,
+                },
+                Err(e) => Some(Err(Error::new_spanned(
+                    att,
+                    format!("could not parse the meta attribute: {e}"),
+                ))),
+            })
+            .transpose()
+}

autometrics/Cargo.toml

@@ -29,6 +29,7 @@ custom-objective-latency = []
 
 [dependencies]
 autometrics-macros = { workspace = true }
+spez = { version = "0.1.2" }
 
 # Used for opentelemetry feature
 opentelemetry_api = { version = "0.18", default-features = false, features = ["metrics"], optional = true }
@@ -51,6 +52,7 @@ axum = { version = "0.6", features = ["tokio"] }
 regex = "1.7"
 http = "0.2"
 tokio = { version = "1", features = ["full"] }
+trybuild = "1.0"
 vergen = { version = "8.1", features = ["git", "gitcl"] }
 
 [package.metadata.docs.rs]