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,13 @@
 use crate::parse::{AutometricsArgs, Item};
+use inflector::Inflector;
 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 syn::{
+    parse_macro_input, Attribute, Data, DataEnum, DeriveInput, Error, ImplItem, ItemFn, ItemImpl,
+    Lit, Meta, NestedMeta, Result,
+};
 
 mod parse;
 
@@ -129,6 +133,18 @@ pub fn autometrics(
     output.into()
 }
 
+#[proc_macro_derive(AutometricsLabel, attributes(autometrics_label))]
+pub fn derive_autometrics_label(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    let result = derive_autometrics_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 +192,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!(),
@@ -194,8 +209,8 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
         // 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, GetLabel};
+                let result_labels = (&result).get_label().map(|(k, v)| (k, Some(v)));
                 CounterLabels::new(
                     #function_name,
                     module_path!(),
@@ -221,14 +236,14 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
 
         #vis #sig {
             let __autometrics_tracker = {
-                use autometrics::__private::{AutometricsTracker, TrackMetrics};
+                use ::autometrics::__private::{AutometricsTracker, TrackMetrics};
                 AutometricsTracker::start(#gauge_labels)
             };
 
             let result = #call_function;
 
             {
-                use autometrics::__private::{HistogramLabels, TrackMetrics};
+                use ::autometrics::__private::{HistogramLabels, TrackMetrics};
                 let counter_labels = #counter_labels;
                 let histogram_labels = HistogramLabels::new(
                     #function_name,
@@ -377,3 +392,149 @@ 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_autometrics_label_impl(input: DeriveInput) -> Result<TokenStream> {
+    let variants = match input.data {
+        Data::Enum(DataEnum { variants, .. }) => variants,
+        _ => {
+            return Err(Error::new_spanned(
+                input,
+                "#[derive(AutometricsLabel}] is only supported for enums",
+            ));
+        }
+    };
+
+    // Use the key provided or the snake case version of the enum name
+    let label_key = {
+        let attrs: Vec<_> = input
+            .attrs
+            .iter()
+            .filter(|attr| attr.path.is_ident("autometrics_label"))
+            .collect();
+
+        let key_from_attr = match attrs.len() {
+            0 => None,
+            1 => get_label_attr(attrs[0], "key")?,
+            _ => {
+                let mut error = syn::Error::new_spanned(
+                    attrs[1],
+                    "redundant `autometrics_label(key)` 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(Error::new_spanned(attrs[0], "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_label"))
+                .collect();
+
+            let value_from_attr = match attrs.len() {
+                0 => None,
+                1 => get_label_attr(attrs[0], "value")?,
+                _ => {
+                    let mut error = Error::new_spanned(
+                        attrs[1],
+                        "redundant `autometrics_label(value)` attribute",
+                    );
+                    error.combine(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(Error::new_spanned(attrs[0], "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! {
+        use ::autometrics::__private::{GetLabel, COUNTER_LABEL_KEYS, linkme};
+
+        #[linkme::distributed_slice(COUNTER_LABEL_KEYS)]
+        #[linkme(crate = ::autometrics::__private::linkme)]
+        pub static COUNTER_LABEL_KEY: &'static str = #label_key;
+
+        #[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(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(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(Error::new_spanned(
+                nested,
+                format!("expected `{attr_name} = \"<value>\"`"),
+            ))
+        }
+    };
+
+    if !label_value.path.is_ident(attr_name) {
+        return Err(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| Error::new_spanned(s, e)),
+        lit => Err(Error::new_spanned(lit, "expected string literal")),
+    }
+}

autometrics/Cargo.toml

@@ -44,6 +44,9 @@ prometheus = { version = "0.13", default-features = false, optional = true }
 # Used for prometheus feature
 const_format = { version = "0.2", features = ["rust_1_51"], optional = true }
 
+# Used to enumerate all generated counter types
+linkme = "0.3.8"
+
 [package.metadata.docs.rs]
 all-features = true
 rustdoc-args = ["--cfg", "docsrs"]

autometrics/src/constants.rs

@@ -13,8 +13,6 @@ pub const FUNCTION_KEY: &'static str = "function";
 pub const MODULE_KEY: &'static str = "module";
 pub const CALLER_KEY: &'static str = "caller";
 pub const RESULT_KEY: &'static str = "result";
-pub const OK_KEY: &'static str = "ok";
-pub const ERROR_KEY: &'static str = "error";
 pub const OBJECTIVE_NAME: &'static str = "objective.name";
 pub const OBJECTIVE_PERCENTILE: &'static str = "objective.percentile";
 pub const OBJECTIVE_LATENCY_THRESHOLD: &'static str = "objective.latency_threshold";

autometrics/src/labels.rs

@@ -1,9 +1,13 @@
 use crate::{constants::*, objectives::*};
+use linkme::distributed_slice;
 use std::ops::Deref;
 
 pub(crate) type Label = (&'static str, &'static str);
 type ResultAndReturnTypeLabels = (&'static str, Option<&'static str>);
 
+#[distributed_slice]
+pub static COUNTER_LABEL_KEYS: [&'static str] = [..];
+
 /// These are the labels used for the `function.calls.count` metric.
 pub struct CounterLabels {
     pub(crate) function: &'static str,
@@ -112,33 +116,6 @@ impl GaugeLabels {
     }
 }
 
-// The following is a convoluted way to figure out if the return type resolves to a Result
-// or not. We cannot simply parse the code using syn to figure out if it's a Result
-// because syn doesn't do type resolution and thus would count any renamed version
-// of Result as a different type. Instead, we define two traits with intentionally
-// conflicting method names and use a trick based on the order in which Rust resolves
-// method names to return a different value based on whether the return value is
-// a Result or anything else.
-// This approach is based on dtolnay's answer to this question:
-// https://users.rust-lang.org/t/how-to-check-types-within-macro/33803/5
-// and this answer explains why it works:
-// https://users.rust-lang.org/t/how-to-check-types-within-macro/33803/8
-
-pub trait GetLabelsFromResult {
-    fn __autometrics_get_labels(&self) -> Option<ResultAndReturnTypeLabels> {
-        None
-    }
-}
-
-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())),
-        }
-    }
-}
-
 pub enum LabelArray {
     Three([Label; 3]),
     Four([Label; 4]),
@@ -157,12 +134,6 @@ impl Deref for LabelArray {
     }
 }
 
-pub trait GetLabels {
-    fn __autometrics_get_labels(&self) -> Option<ResultAndReturnTypeLabels> {
-        None
-    }
-}
-
 /// Implement the given trait for &T and all primitive types.
 macro_rules! impl_trait_for_types {
     ($trait:ident) => {
@@ -200,24 +171,19 @@ macro_rules! impl_trait_for_types {
     };
 }
 
-impl_trait_for_types!(GetLabels);
-
-pub trait GetStaticStrFromIntoStaticStr<'a> {
-    fn __autometrics_static_str(&'a self) -> Option<&'static str>;
-}
-
-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())
+pub trait GetLabel {
+    fn get_label(&self) -> Option<(&'static str, &'static str)> {
+        None
     }
 }
 
-pub trait GetStaticStr {
-    fn __autometrics_static_str(&self) -> Option<&'static str> {
-        None
+impl<T: GetLabel, E: GetLabel> GetLabel for Result<T, E> {
+    fn get_label(&self) -> Option<(&'static str, &'static str)> {
+        match self {
+            Ok(v) => (*v).get_label(),
+            Err(v) => (*v).get_label(),
+        }
     }
 }
-impl_trait_for_types!(GetStaticStr);
+
+impl_trait_for_types!(GetLabel);

autometrics/src/lib.rs

@@ -13,7 +13,12 @@ mod prometheus_exporter;
 mod task_local;
 mod tracker;
 
-pub use autometrics_macros::autometrics;
+pub extern crate linkme;
+
+pub use labels::GetLabel;
+
+pub extern crate autometrics_macros;
+pub use autometrics_macros::{autometrics, AutometricsLabel};
 
 // Optional exports
 #[cfg(feature = "prometheus-exporter")]
@@ -47,4 +52,9 @@ pub mod __private {
 
         LocalKey { inner: CALLER_KEY }
     };
+
+    /// Re-export the linkme crate so that it can be used in the code generated by the autometrics macro
+    pub mod linkme {
+        pub use ::linkme::*;
+    }
 }