Add #[with_removal] attribute to #[metrics] macro

`#[with_removal]` generates two additional functions that allow to
remove individual label sets from a metric family as well as clear the
family completely. This is useful when metrics are used for dynamic
data, such as peer RTT in a P2P system.

The feature is hidden behind `cfg(foundations_unstable)`, which allows
us to iterate on the API in the future. Both proc-macro unit tests as
well as an integration test are included in the change.

Author: TheJokr

foundations-macros/src/metrics/mod.rs

@@ -58,6 +58,7 @@ struct FnAttrs {
     doc: String,
     ctor: Option<ExprStruct>,
     optional: bool,
+    with_removal: bool,
 }
 
 struct FnArg {
@@ -300,6 +301,7 @@ fn metric_init(foundations: &Path, fn_: &ItemFn) -> proc_macro2::TokenStream {
                 doc,
                 optional,
                 ctor,
+                ..
             },
         ident: field_name,
         args,
@@ -357,7 +359,13 @@ fn metric_init(foundations: &Path, fn_: &ItemFn) -> proc_macro2::TokenStream {
 
 fn metric_fn(foundations: &Path, metrics_struct: &Ident, fn_: &ItemFn) -> proc_macro2::TokenStream {
     let ItemFn {
-        attrs: FnAttrs { cfg, doc, .. },
+        attrs:
+            FnAttrs {
+                cfg,
+                doc,
+                with_removal,
+                ..
+            },
         fn_token,
         vis: fn_vis,
         ident: metric_name,
@@ -366,35 +374,77 @@ fn metric_fn(foundations: &Path, metrics_struct: &Ident, fn_: &ItemFn) -> proc_m
         ty: metric_type,
     } = fn_;
 
-    let fn_args = args.iter().map(|arg| arg.to_arg());
+    let fn_args: Vec<_> = args.iter().map(|arg| arg.to_arg()).collect();
 
-    let fn_body = if args.is_empty() {
-        quote! {
+    let (convert_args, access_metric) = if args.is_empty() {
+        let accessor = quote! {
             ::std::clone::Clone::clone(&#metrics_struct.#metric_name)
-        }
+        };
+        (quote! {}, accessor)
     } else {
         let label_inits = args.iter().map(|arg| arg.to_initializer());
+        let convert = quote! {
+            let __args = #metric_name {
+                #(#label_inits,)*
+            };
+        };
 
-
-        quote! {
+        let accessor = quote! {
             ::std::clone::Clone::clone(
                 &#foundations::reexports_for_macros::prometools::serde::Family::get_or_create(
                     &#metrics_struct.#metric_name,
-                    &#metric_name {
-                        #(#label_inits,)*
-                    },
+                    &__args,
                 )
             )
+        };
+        (convert, accessor)
+    };
+
+    let removal_fns = if cfg!(foundations_unstable) && *with_removal {
+        let remove_ident = format_ident!("{metric_name}_remove");
+        let remove_doc = LitStr::new(
+            &format!("Removes one label set from the `{metric_name}` family."),
+            Span::call_site(),
+        );
+
+        let clear_ident = format_ident!("{metric_name}_clear");
+        let clear_doc = LitStr::new(
+            &format!("Removes all label sets from the `{metric_name}` family."),
+            Span::call_site(),
+        );
+
+        quote! {
+            #[doc = #remove_doc]
+            #(#cfg)*
+            #fn_vis #fn_token #remove_ident(#(#fn_args,)*) #arrow_token bool {
+                #convert_args
+                #foundations::reexports_for_macros::prometools::serde::Family::remove(
+                    &#metrics_struct.#metric_name,
+                    &__args,
+                )
+            }
+
+            #[doc = #clear_doc]
+            #(#cfg)*
+            #fn_vis #fn_token #clear_ident() {
+                #foundations::reexports_for_macros::prometools::serde::Family::clear(
+                    &#metrics_struct.#metric_name,
+                )
+            }
         }
+    } else {
+        quote! {}
     };
 
     quote! {
         #[doc = #doc]
         #(#cfg)*
         #[must_use]
         #fn_vis #fn_token #metric_name(#(#fn_args,)*) #arrow_token #metric_type {
-            #fn_body
+            #convert_args
+            #access_metric
         }
+        #removal_fns
     }
 }
 
@@ -713,15 +763,16 @@ mod tests {
                     message: &'static str,
                     error: impl Into<String>,
                 ) -> Counter {
+                    let __args = connections_errors_total {
+                        endpoint: ::std::clone::Clone::clone(endpoint),
+                        kind,
+                        message,
+                        error: ::std::convert::Into::into(error),
+                    };
                     ::std::clone::Clone::clone(
                         &::foundations::reexports_for_macros::prometools::serde::Family::get_or_create(
                             &__oxy_Metrics.connections_errors_total,
-                            &connections_errors_total {
-                                endpoint: ::std::clone::Clone::clone(endpoint),
-                                kind,
-                                message,
-                                error: ::std::convert::Into::into(error),
-                            },
+                            &__args,
                         )
                     )
                 }
@@ -827,15 +878,111 @@ mod tests {
                 pub fn requests_per_connection(
                     endpoint: String,
                 ) -> Histogram {
+                    let __args = requests_per_connection { endpoint, };
                     ::std::clone::Clone::clone(
                         &::foundations::reexports_for_macros::prometools::serde::Family::get_or_create(
                             &__oxy_Metrics.requests_per_connection,
-                            &requests_per_connection {
-                                endpoint,
+                            &__args,
+                        )
+                    )
+                }
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
+    #[cfg(foundations_unstable)]
+    #[test]
+    fn expand_with_removal() {
+        let attr = parse_attr! {
+            #[metrics]
+        };
+
+        let src = parse_quote! {
+            pub(crate) mod oxy {
+                /// Total number of requests
+                #[with_removal]
+                pub(crate) fn requests_total(status: u16) -> Counter;
+            }
+        };
+
+        let actual = expand_from_parsed(attr, src).to_string();
+
+        let expected = code_str! {
+            pub(crate) mod oxy {
+                use super::*;
+
+                #[allow(non_camel_case_types)]
+                struct __oxy_Metrics {
+                    requests_total:
+                        ::foundations::reexports_for_macros::prometools::serde::Family<
+                            requests_total,
+                            Counter,
+                        >,
+                }
+
+                #[allow(non_camel_case_types)]
+                #[derive(
+                    ::std::clone::Clone,
+                    ::std::cmp::Eq,
+                    ::std::hash::Hash,
+                    ::std::cmp::PartialEq,
+                    ::foundations::reexports_for_macros::serde::Serialize,
+                )]
+                #[serde(crate = ":: foundations :: reexports_for_macros :: serde")]
+                struct requests_total {
+                    status: u16,
+                }
+
+                #[allow(non_upper_case_globals)]
+                static __oxy_Metrics: ::foundations::reexports_for_macros::once_cell::sync::Lazy<__oxy_Metrics> =
+                    ::foundations::reexports_for_macros::once_cell::sync::Lazy::new(|| {
+                        let registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), false, true);
+
+                        __oxy_Metrics {
+                            requests_total: {
+                                let metric = ::std::default::Default::default();
+
+                                ::foundations::reexports_for_macros::prometheus_client::registry::Registry::register(
+                                    registry,
+                                    ::std::stringify!(requests_total),
+                                    str::trim(" Total number of requests"),
+                                    ::std::boxed::Box::new(::std::clone::Clone::clone(&metric))
+                                );
+
+                                metric
                             },
+                        }
+                    });
+
+                #[doc = " Total number of requests"]
+                #[must_use]
+                pub(crate) fn requests_total(status: u16,) -> Counter {
+                    let __args = requests_total { status, };
+                    ::std::clone::Clone::clone(
+                        &::foundations::reexports_for_macros::prometools::serde::Family::get_or_create(
+                            &__oxy_Metrics.requests_total,
+                            &__args,
                         )
                     )
                 }
+
+                #[doc = "Removes one label set from the `requests_total` family."]
+                pub(crate) fn requests_total_remove(status: u16,) -> bool {
+                    let __args = requests_total { status, };
+                    ::foundations::reexports_for_macros::prometools::serde::Family::remove(
+                        &__oxy_Metrics.requests_total,
+                        &__args,
+                    )
+                }
+
+                #[doc = "Removes all label sets from the `requests_total` family."]
+                pub(crate) fn requests_total_clear() {
+                    ::foundations::reexports_for_macros::prometools::serde::Family::clear(
+                        &__oxy_Metrics.requests_total,
+                    )
+                }
             }
         };
 

foundations-macros/src/metrics/parsing.rs

@@ -1,26 +1,27 @@
 use super::{ArgAttrs, ArgMode, FnArg, FnAttrs, ItemFn, MacroArgs, Mod};
 use crate::common::{error, parse_attr_value, parse_meta_list, Result};
 use darling::FromMeta;
-use quote::ToTokens as _;
 use syn::parse::{Parse, ParseStream};
 use syn::punctuated::Punctuated;
 use syn::{
-    braced, parenthesized, AngleBracketedGenericArguments, Attribute, GenericArgument, LitBool,
-    LitStr, PathArguments, Token, TraitBound, TraitBoundModifier, Type, TypeImplTrait,
-    TypeParamBound,
+    braced, parenthesized, AngleBracketedGenericArguments, Attribute, GenericArgument,
+    PathArguments, Token, TraitBound, TraitBoundModifier, Type, TypeImplTrait, TypeParamBound,
 };
 
 const IMPL_TRAIT_ERROR: &str = "Only `impl Into<T>` is allowed";
 
 const FN_ATTR_ERROR: &str =
-    "Only `#[cfg]`, `#[doc]`, `#[ctor]` and `#[optional]` are allowed on functions";
+    "Only `#[cfg]`, `#[doc]`, `#[ctor]`, `#[optional]`, and `#[with_removal]` are allowed on functions";
 
 const DUPLICATE_CTOR_ATTR_ERROR: &str = "Duplicate `#[ctor]` attribute";
 const DUPLICATE_OPTIONAL_ATTR_ERROR: &str = "Duplicate `#[optional]` attribute";
+const DUPLICATE_WITH_REMOVAL_ATTR_ERROR: &str = "Duplicate `#[with_removal]` attribute";
 const DUPLICATE_SERDE_ATTR_ERROR: &str = "Duplicate `#[serde]` attribute";
 const DUPLICATE_SERDE_AS_ATTR_ERROR: &str = "Duplicate `#[serde_as]` attribute";
 
 const ARG_ATTR_ERROR: &str = "Only `#[serde]` and `#[serde_as]` are allowed on function arguments";
+const WITH_REMOVAL_NO_ARGS_ERROR: &str =
+    "`#[with_removal]` can only be used on functions with arguments";
 
 impl Parse for MacroArgs {
     fn parse(input: ParseStream) -> Result<Self> {
@@ -58,60 +59,63 @@ impl Parse for Mod {
     }
 }
 
-impl Parse for ItemFn {
-    fn parse(input: ParseStream) -> Result<Self> {
-        fn parse_attrs(attrs: Vec<Attribute>) -> Result<FnAttrs> {
-            let mut cfg = vec![];
-            let mut doc = "".to_owned();
-            let mut ctor = None;
-            let mut optional = None;
-
-            for attr in attrs {
-                let path = attr.path();
-
-                if path.is_ident("cfg") {
-                    cfg.push(attr);
-                } else if path.is_ident("doc") {
-                    doc.push_str(&parse_attr_value::<LitStr>(attr)?.value());
-                } else if path.is_ident("ctor") {
+impl FnAttrs {
+    fn from_attrs(attrs: Vec<Attribute>) -> Result<Self> {
+        let mut cfg = vec![];
+        let mut doc = "".to_owned();
+        let mut ctor = None;
+        let mut optional = None;
+        let mut with_removal = None;
+
+        for attr in attrs {
+            let path = attr.path().get_ident().map(ToString::to_string);
+
+            match path.as_deref() {
+                Some("cfg") => cfg.push(attr),
+                Some("doc") => doc.push_str(&String::from_meta(&attr.meta)?),
+                Some("ctor") => {
                     if ctor.is_some() {
                         return error(&attr, DUPLICATE_CTOR_ATTR_ERROR);
                     }
 
                     ctor = Some(parse_attr_value(attr)?);
-                } else if path.is_ident("optional") {
+                }
+                Some("optional") => {
                     if optional.is_some() {
                         return error(&attr, DUPLICATE_OPTIONAL_ATTR_ERROR);
                     }
 
-                    if attr.to_token_stream().is_empty() {
-                        optional = Some(true);
-                    } else {
-                        optional = Some(
-                            parse_attr_value::<LitBool>(attr)
-                                .map(|l| l.value)
-                                .unwrap_or(true),
-                        );
+                    optional = Some(bool::from_meta(&attr.meta).unwrap_or(true));
+                }
+                Some("with_removal") => {
+                    if with_removal.is_some() {
+                        return error(&attr, DUPLICATE_WITH_REMOVAL_ATTR_ERROR);
                     }
-                } else {
-                    return error(&attr, FN_ATTR_ERROR);
+
+                    with_removal = Some(bool::from_meta(&attr.meta)?);
                 }
+                _ => return error(&attr, FN_ATTR_ERROR),
             }
-
-            Ok(FnAttrs {
-                cfg,
-                doc,
-                ctor,
-                optional: optional.unwrap_or(false),
-            })
         }
 
-        let attrs = parse_attrs(input.call(Attribute::parse_outer)?)?;
+        Ok(Self {
+            cfg,
+            doc,
+            ctor,
+            optional: optional.unwrap_or(false),
+            with_removal: with_removal.unwrap_or(false),
+        })
+    }
+}
+
+impl Parse for ItemFn {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let attrs = FnAttrs::from_attrs(input.call(Attribute::parse_outer)?)?;
         let vis = input.parse()?;
         let fn_token = input.parse()?;
         let ident = input.parse()?;
         let args_content;
-        let _paren_token = parenthesized!(args_content in input);
+        let paren_token = parenthesized!(args_content in input);
         let mut args = Punctuated::new();
 
         while !args_content.is_empty() {
@@ -128,6 +132,10 @@ impl Parse for ItemFn {
         let ty = input.parse()?;
         let _semi_token = input.parse::<Token![;]>()?;
 
+        if attrs.with_removal && args.is_empty() {
+            return error(&paren_token.span, WITH_REMOVAL_NO_ARGS_ERROR);
+        }
+
         Ok(ItemFn {
             attrs,
             vis,