Skip to content

Commit

Permalink
Allow label key to also be provided.
Browse files Browse the repository at this point in the history
  • Loading branch information
sagacity committed Mar 15, 2023
1 parent 59ee750 commit 1fd408b
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 50 deletions.
1 change: 1 addition & 0 deletions autometrics-macros/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ categories = { workspace = true }
proc-macro = true

[dependencies]
Inflector = "0.11.4"
percent-encoding = "2.2"
proc-macro2 = "1"
quote = "1"
Expand Down
76 changes: 56 additions & 20 deletions autometrics-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use proc_macro2::{Ident, TokenStream};
use quote::quote;
use std::env;
use inflector::Inflector;
use syn::{parse_macro_input, DeriveInput, ImplItem, ItemFn, ItemImpl, Result, Data, DataEnum, Attribute, Meta, NestedMeta, Lit};

mod parse;
Expand Down Expand Up @@ -129,10 +130,10 @@ pub fn autometrics(
output.into()
}

#[proc_macro_derive(LabelValues, attributes(autometrics))]
pub fn derive_label_values(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
#[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_label_values_impl(input);
let result = derive_get_label_impl(input);
let output = match result {
Ok(output) => output,
Err(err) => err.into_compile_error(),
Expand Down Expand Up @@ -188,9 +189,9 @@ fn instrument_function(args: &AutometricsArgs, item: ItemFn) -> Result<TokenStre
};
quote! {
{
use autometrics::__private::{CALLER, CounterLabels, GetLabelValue};
use autometrics::__private::{CALLER, CounterLabels, GetLabel};
let result_label = #result_label;
let value_type = (&result).get_label_value();
let value_type = (&result).get_label().map(|(_, v)| v);
CounterLabels::new(
#function_name,
module_path!(),
Expand Down Expand Up @@ -389,22 +390,49 @@ fn concurrent_calls_query(gauge_name: &str, label_key: &str, label_value: &str)
format!("sum by (function, module) {gauge_name}{{{label_key}=\"{label_value}\"}}")
}

fn derive_label_values_impl(input: DeriveInput) -> Result<TokenStream> {
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 match_arms = variants
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_value_attr(attrs[0])?,
1 => get_label_attr(attrs[0], "label_value")?,
_ => {
let mut error =
syn::Error::new_spanned(attrs[1], "redundant `autometrics(label_value)` attribute");
Expand All @@ -413,29 +441,37 @@ fn derive_label_values_impl(input: DeriveInput) -> Result<TokenStream> {
}
};

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());
let value = value.to_string();
let value = value_from_attr.unwrap_or_else(|| ident.clone().to_string().to_snake_case());
Ok(quote! {
Self::#ident => Some(#value),
Self::#ident => #value,
})
})
.collect::<Result<TokenStream>>()?;

let ident = input.ident;
Ok(quote! {
#[automatically_derived]
impl GetLabelValue for #ident {
fn get_label_value(&self) -> Option<&'static str> {
match self {
#match_arms
}
impl GetLabel for #ident {
fn get_label(&self) -> Option<(&'static str, &'static str)> {
Some((#label_key, match self {
#value_match_arms
}))
}
}
})
}

fn get_label_value_attr(attr: &Attribute) -> Result<Option<Ident>> {
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,
Expand All @@ -456,13 +492,13 @@ fn get_label_value_attr(attr: &Attribute) -> Result<Option<Ident>> {

let label_value = match nested {
NestedMeta::Meta(Meta::NameValue(nv)) => nv,
_ => return Err(syn::Error::new_spanned(nested, "expected `label_value = \"<value>\"`")),
_ => return Err(syn::Error::new_spanned(nested, format!("expected `{attr_name} = \"<value>\"`"))),
};

if !label_value.path.is_ident("label_value") {
if !label_value.path.is_ident(attr_name) {
return Err(syn::Error::new_spanned(
&label_value.path,
"unsupported autometrics attribute, expected `label_value`",
format!("unsupported autometrics attribute, expected `{attr_name}`"),
));
}

Expand Down
56 changes: 27 additions & 29 deletions autometrics/src/labels.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.get_label_value())),
Err(err) => Some((ERROR_KEY, err.get_label_value())),
}
self.get_label().map(|(k, v)| (k, Some(v)))
}
}

Expand Down Expand Up @@ -202,64 +199,65 @@ macro_rules! impl_trait_for_types {

impl_trait_for_types!(GetLabels);

pub trait GetLabelValue {
fn get_label_value(&self) -> Option<&'static str> {
pub trait GetLabel {
fn get_label(&self) -> Option<(&'static str, &'static str)> {
None
}
}
impl_trait_for_types!(GetLabelValue);
impl_trait_for_types!(GetLabel);

#[cfg(test)]
mod tests {
use super::*;
use autometrics_macros::LabelValues;
use autometrics_macros::GetLabel;

#[test]
fn custom_trait_implementation() {
struct CustomResult;

impl GetLabelValue for CustomResult {
fn get_label_value(&self) -> Option<&'static str> {
Some("my-result")
impl GetLabel for CustomResult {
fn get_label(&self) -> Option<(&'static str, &'static str)> {
Some(("ok", "my-result"))
}
}

assert_eq!(Some("my-result"), CustomResult {}.get_label_value());
assert_eq!(Some(("ok", "my-result")), CustomResult {}.get_label());
}

#[test]
fn manual_enum() {
enum Foo {
enum MyFoo {
A,
B,
}

impl GetLabelValue for Foo {
fn get_label_value(&self) -> Option<&'static str> {
match self {
Foo::A => Some("a"),
Foo::B => Some("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("a"), Foo::A.get_label_value());
assert_eq!(Some("b"), Foo::B.get_label_value());
assert_eq!(Some(("hello", "a")), MyFoo::A.get_label());
assert_eq!(Some(("hello", "b")), MyFoo::B.get_label());
}

#[test]
fn derived_enum() {
#[derive(LabelValues)]
enum Foo {
#[derive(GetLabel)]
#[autometrics(label_key = "my_foo")]
enum MyFoo {
#[autometrics(label_value = "hello")]
A,
Alpha,
#[autometrics()]
B,
C,
BetaValue,
Charlie,
}

assert_eq!(Some("hello"), Foo::A.get_label_value());
assert_eq!(Some("B"), Foo::B.get_label_value());
assert_eq!(Some("C"), Foo::C.get_label_value());
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());
}
}
2 changes: 1 addition & 1 deletion autometrics/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ mod task_local;
mod tracker;

pub use autometrics_macros::autometrics;
pub use labels::GetLabelValue;
pub use labels::GetLabel;
pub use objectives::{Objective, ObjectivePercentage, TargetLatency};

// Optional exports
Expand Down

0 comments on commit 1fd408b

Please sign in to comment.