Allow enum variants to be used as label values

autometrics-macros/Cargo.toml

@@ -14,6 +14,7 @@ categories = { workspace = true }
 proc-macro = true
 
 [dependencies]
+Inflector = "0.11.4"
 percent-encoding = "2.2"
 proc-macro2 = "1"
 quote = "1"

autometrics-macros/src/lib.rs

@@ -1,9 +1,10 @@
 use crate::parse::{AutometricsArgs, Item};
 use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
-use proc_macro2::TokenStream;
+use proc_macro2::{Ident, TokenStream};
 use quote::quote;
 use std::env;
-use syn::{parse_macro_input, ImplItem, ItemFn, ItemImpl, Result};
+use inflector::Inflector;
+use syn::{parse_macro_input, DeriveInput, ImplItem, ItemFn, ItemImpl, Result, Data, DataEnum, Attribute, Meta, NestedMeta, Lit};
 
 mod parse;
 
@@ -129,6 +130,18 @@ pub fn autometrics(
     output.into()
 }
 
+#[proc_macro_derive(GetLabel, attributes(autometrics))]
+pub fn derive_get_label(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    let result = derive_get_label_impl(input);
+    let output = match result {
+        Ok(output) => output,
+        Err(err) => err.into_compile_error(),
+    };
+
+    output.into()
+}
+
 /// Add autometrics instrumentation to a single function
 fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStream> {
     let sig = item.sig;
@@ -176,10 +189,9 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
         };
         quote! {
             {
-                use autometrics::__private::{CALLER, CounterLabels, GetStaticStrFromIntoStaticStr, GetStaticStr};
+                use autometrics::__private::{CALLER, CounterLabels, GetLabel};
                 let result_label = #result_label;
-                // If the return type implements Into<&'static str>, attach that as a label
-                let value_type = (&result).__autometrics_static_str();
+                let value_type = (&result).get_label().map(|(_, v)| v);
                 CounterLabels::new(
                     #function_name,
                      module_path!(),
@@ -377,3 +389,121 @@ histogram_quantile(0.95, {latency})"
 fn concurrent_calls_query(gauge_name: &str, label_key: &str, label_value: &str) -> String {
     format!("sum by (function, module) {gauge_name}{{{label_key}=\"{label_value}\"}}")
 }
+
+fn derive_get_label_impl(input: DeriveInput) -> Result<TokenStream> {
+    let variants = match input.data {
+        Data::Enum(DataEnum { variants, .. }) => variants,
+        _ => {
+            return Err(syn::Error::new_spanned(input, "#[derive(LabelValues}] is only supported for enums"));
+        },
+    };
+
+    let label_key = {
+        let attrs: Vec<_> = input.attrs.iter().filter(|attr| attr.path.is_ident("autometrics")).collect();
+
+        let key_from_attr = match attrs.len() {
+            0 => None,
+            1 => get_label_attr(attrs[0], "label_key")?,
+            _ => {
+                let mut error =
+                    syn::Error::new_spanned(attrs[1], "redundant `autometrics(label_value)` attribute");
+                error.combine(syn::Error::new_spanned(attrs[0], "note: first one here"));
+                return Err(error);
+            }
+        };
+
+        let key_from_attr = key_from_attr.map(|value| value.to_string());
+
+        // Check casing of the user-provided value
+        if let Some(key) = &key_from_attr {
+            if key.as_str() != key.to_snake_case() {
+                return Err(syn::Error::new_spanned(attrs[0], "label_key should be snake_cased"));
+            }
+        }
+
+        let ident = input.ident.clone();
+        key_from_attr.unwrap_or_else(|| ident.clone().to_string().to_snake_case())
+    };
+
+    let value_match_arms = variants
+        .into_iter()
+        .map(|variant| {
+            let attrs: Vec<_> = variant.attrs.iter().filter(|attr| attr.path.is_ident("autometrics")).collect();
+
+            let value_from_attr = match attrs.len() {
+                0 => None,
+                1 => get_label_attr(attrs[0], "label_value")?,
+                _ => {
+                    let mut error =
+                        syn::Error::new_spanned(attrs[1], "redundant `autometrics(label_value)` attribute");
+                    error.combine(syn::Error::new_spanned(attrs[0], "note: first one here"));
+                    return Err(error);
+                }
+            };
+
+            let value_from_attr = value_from_attr.map(|value| value.to_string());
+
+            // Check casing of the user-provided value
+            if let Some(value) = &value_from_attr {
+                if value.as_str() != value.to_snake_case() {
+                    return Err(syn::Error::new_spanned(attrs[0], "label_value should be snake_cased"));
+                }
+            }
+
+            let ident = variant.ident;
+            let value = value_from_attr.unwrap_or_else(|| ident.clone().to_string().to_snake_case());
+            Ok(quote! {
+                Self::#ident => #value,
+            })
+        })
+        .collect::<Result<TokenStream>>()?;
+
+    let ident = input.ident;
+    Ok(quote! {
+        #[automatically_derived]
+        impl GetLabel for #ident {
+            fn get_label(&self) -> Option<(&'static str, &'static str)> {
+                Some((#label_key, match self {
+                    #value_match_arms
+                }))
+            }
+        }
+    })
+}
+
+fn get_label_attr(attr: &Attribute, attr_name: &str) -> Result<Option<Ident>> {
+    let meta = attr.parse_meta()?;
+    let meta_list = match meta {
+        Meta::List(list) => list,
+        _ => return Err(syn::Error::new_spanned(meta, "expected a list-style attribute")),
+    };
+
+    let nested = match meta_list.nested.len() {
+        // `#[autometrics()]` without any arguments is a no-op
+        0 => return Ok(None),
+        1 => &meta_list.nested[0],
+        _ => {
+            return Err(syn::Error::new_spanned(
+                meta_list.nested,
+                "currently only a single autometrics attribute is supported",
+            ));
+        }
+    };
+
+    let label_value = match nested {
+        NestedMeta::Meta(Meta::NameValue(nv)) => nv,
+        _ => return Err(syn::Error::new_spanned(nested, format!("expected `{attr_name} = \"<value>\"`"))),
+    };
+
+    if !label_value.path.is_ident(attr_name) {
+        return Err(syn::Error::new_spanned(
+            &label_value.path,
+            format!("unsupported autometrics attribute, expected `{attr_name}`"),
+        ));
+    }
+
+    match &label_value.lit {
+        Lit::Str(s) => syn::parse_str(&s.value()).map_err(|e| syn::Error::new_spanned(s, e)),
+        lit => Err(syn::Error::new_spanned(lit, "expected string literal")),
+    }
+}

autometrics/src/labels.rs

@@ -132,10 +132,7 @@ pub trait GetLabelsFromResult {
 
 impl<T, E> GetLabelsFromResult for Result<T, E> {
     fn __autometrics_get_labels(&self) -> Option<ResultAndReturnTypeLabels> {
-        match self {
-            Ok(ok) => Some((OK_KEY, ok.__autometrics_static_str())),
-            Err(err) => Some((ERROR_KEY, err.__autometrics_static_str())),
-        }
+        self.get_label().map(|(k, v)| (k, Some(v)))
     }
 }
 
@@ -202,22 +199,65 @@ macro_rules! impl_trait_for_types {
 
 impl_trait_for_types!(GetLabels);
 
-pub trait GetStaticStrFromIntoStaticStr<'a> {
-    fn __autometrics_static_str(&'a self) -> Option<&'static str>;
+pub trait GetLabel {
+    fn get_label(&self) -> Option<(&'static str, &'static str)> {
+        None
+    }
 }
+impl_trait_for_types!(GetLabel);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use autometrics_macros::GetLabel;
 
-impl<'a, T: 'a> GetStaticStrFromIntoStaticStr<'a> for T
-where
-    &'static str: From<&'a T>,
-{
-    fn __autometrics_static_str(&'a self) -> Option<&'static str> {
-        Some(self.into())
+    #[test]
+    fn custom_trait_implementation() {
+        struct CustomResult;
+
+        impl GetLabel for CustomResult {
+            fn get_label(&self) -> Option<(&'static str, &'static str)> {
+                Some(("ok", "my-result"))
+            }
+        }
+
+        assert_eq!(Some(("ok", "my-result")), CustomResult {}.get_label());
     }
-}
 
-pub trait GetStaticStr {
-    fn __autometrics_static_str(&self) -> Option<&'static str> {
-        None
+    #[test]
+    fn manual_enum() {
+        enum MyFoo {
+            A,
+            B,
+        }
+
+        impl GetLabel for MyFoo {
+            fn get_label(&self) -> Option<(&'static str, &'static str)> {
+                Some(("hello", match self {
+                    MyFoo::A => "a",
+                    MyFoo::B => "b",
+                }))
+            }
+        }
+
+        assert_eq!(Some(("hello", "a")), MyFoo::A.get_label());
+        assert_eq!(Some(("hello", "b")), MyFoo::B.get_label());
+    }
+
+    #[test]
+    fn derived_enum() {
+        #[derive(GetLabel)]
+        #[autometrics(label_key = "my_foo")]
+        enum MyFoo {
+            #[autometrics(label_value = "hello")]
+            Alpha,
+            #[autometrics()]
+            BetaValue,
+            Charlie,
+        }
+
+        assert_eq!(Some(("my_foo", "hello")), MyFoo::Alpha.get_label());
+        assert_eq!(Some(("my_foo", "beta_value")), MyFoo::BetaValue.get_label());
+        assert_eq!(Some(("my_foo", "charlie")), MyFoo::Charlie.get_label());
     }
 }
-impl_trait_for_types!(GetStaticStr);

autometrics/src/lib.rs

@@ -14,6 +14,8 @@ mod task_local;
 mod tracker;
 
 pub use autometrics_macros::autometrics;
+pub use labels::GetLabel;
+pub use objectives::{Objective, ObjectivePercentage, TargetLatency};
 
 // Optional exports
 #[cfg(feature = "prometheus-exporter")]