Add unprefixed flag to #[metrics] macro

TheJokr · TheJokr · commit f0194db389a7 · 2025-11-06T20:06:40.000+01:00
This change is twofold. First, it refactors the registry selection in
`metrics::internal::Registries` to use boolean flags rather than
separate methods to select a registry to use. As part of this, we add a
combination of flags to return a registry without the global service
name prefix that is normally present.

Second, we add a new `unprefixed` argument to the `metrics` proc-macro.
If present, this argument causes all metrics in the associated module to
be registered in the registry mentioned above. This works for both the
"main" and the "optional" registry.

I've added both a proc-macro expansion test and an integration test
to validate the flag's functionality.
diff --git a/foundations-macros/src/metrics/mod.rs b/foundations-macros/src/metrics/mod.rs
@@ -1,7 +1,8 @@
+use darling::util::Flag;
 use darling::FromMeta;
 use proc_macro::TokenStream;
 use proc_macro2::Span;
-use quote::{quote, ToTokens};
+use quote::{format_ident, quote, ToTokens};
 use syn::punctuated::Punctuated;
 use syn::{
     parse_macro_input, parse_quote, Attribute, ExprStruct, Ident, LitStr, Path, Token, Type,
@@ -13,13 +14,15 @@ mod validation;
 
 #[derive(FromMeta)]
 struct MacroArgs {
+    unprefixed: Flag,
     #[darling(default = "Self::default_crate_path")]
     crate_path: Path,
 }
 
 impl Default for MacroArgs {
     fn default() -> Self {
         Self {
+            unprefixed: Flag::default(),
             crate_path: Self::default_crate_path(),
         }
     }
@@ -86,8 +89,11 @@ pub(crate) fn expand(args: TokenStream, item: TokenStream) -> TokenStream {
 
 fn expand_from_parsed(args: MacroArgs, extern_: Mod) -> proc_macro2::TokenStream {
     let MacroArgs {
+        unprefixed,
         crate_path: foundations,
     } = &args;
+    let with_service_prefix =
+        format_ident!("{}", !unprefixed.is_present(), span = unprefixed.span());
 
     let Mod {
         attrs: mod_attrs,
@@ -107,24 +113,26 @@ fn expand_from_parsed(args: MacroArgs, extern_: Mod) -> proc_macro2::TokenStream
         .iter()
         .filter_map(|fn_| label_set_struct(foundations, fn_));
 
-    let registry_init = |var: &str, kind: &str| {
+    let registry_init = |var: &str, opt: bool| {
         let var = Ident::new(var, Span::call_site());
-        let method = Ident::new(&format!("get_{kind}_subsystem"), Span::call_site());
+        let optional = format_ident!("{opt}");
 
         quote! {
-            let #var = &mut *#foundations::telemetry::metrics::internal::Registries::#method(stringify!(#mod_name));
+            let #var = &mut *#foundations::telemetry::metrics::internal::Registries::get_subsystem(
+                stringify!(#mod_name), #optional, #with_service_prefix
+            );
         }
     };
 
     let init_registry = fns
         .iter()
         .any(|fn_| !fn_.attrs.optional)
-        .then(|| registry_init("registry", "main"));
+        .then(|| registry_init("registry", false));
 
     let init_opt_registry = fns
         .iter()
         .any(|fn_| fn_.attrs.optional)
-        .then(|| registry_init("opt_registry", "opt"));
+        .then(|| registry_init("opt_registry", true));
 
     let metric_inits = fns.iter().map(|fn_| metric_init(foundations, fn_));
 
@@ -453,7 +461,7 @@ mod tests {
                 #[allow(non_upper_case_globals)]
                 static __oxy_Metrics: tarmac::reexports_for_macros::once_cell::sync::Lazy<__oxy_Metrics> =
                     tarmac::reexports_for_macros::once_cell::sync::Lazy::new(|| {
-                        let registry = &mut *tarmac::telemetry::metrics::internal::Registries::get_main_subsystem(stringify!(oxy));
+                        let registry = &mut *tarmac::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), false, true);
 
                         __oxy_Metrics {
                             connections_total: {
@@ -510,9 +518,83 @@ mod tests {
                 #[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 opt_registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_opt_subsystem(stringify!(oxy));
+                        let opt_registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), true, true);
+
+                        __oxy_Metrics {
+                            connections_total: {
+                                let metric = ::std::default::Default::default();
+
+                                ::foundations::reexports_for_macros::prometheus_client::registry::Registry::register(
+                                    opt_registry,
+                                    ::std::stringify!(connections_total),
+                                    str::trim(" Total number of connections"),
+                                    ::std::boxed::Box::new(::std::clone::Clone::clone(&metric))
+                                );
+
+                                metric
+                            },
+                        }
+                    });
+
+                #[doc = " Total number of connections"]
+                #[must_use]
+                pub(crate) fn connections_total() -> Counter {
+                    ::std::clone::Clone::clone(&__oxy_Metrics.connections_total)
+                }
+            }
+        };
+
+        assert_eq!(actual, expected);
+    }
+
+    #[test]
+    fn expand_unprefixed_mixed() {
+        let attr = parse_attr! {
+            #[metrics(unprefixed)]
+        };
+
+        let src = parse_quote! {
+            pub(crate) mod oxy {
+                /// Total number of requests
+                pub(crate) fn requests_total() -> Counter;
+
+                /// Total number of connections
+                #[optional]
+                pub(crate) fn connections_total() -> 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: Counter,
+                    connections_total: Counter,
+                }
+
+                #[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, false);
+                        let opt_registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), true, false);
 
                         __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
+                            },
                             connections_total: {
                                 let metric = ::std::default::Default::default();
 
@@ -528,6 +610,12 @@ mod tests {
                         }
                     });
 
+                #[doc = " Total number of requests"]
+                #[must_use]
+                pub(crate) fn requests_total() -> Counter {
+                    ::std::clone::Clone::clone(&__oxy_Metrics.requests_total)
+                }
+
                 #[doc = " Total number of connections"]
                 #[must_use]
                 pub(crate) fn connections_total() -> Counter {
@@ -596,7 +684,7 @@ mod tests {
                 #[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_main_subsystem(stringify!(oxy));
+                        let registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), false, true);
 
                         __oxy_Metrics {
                             connections_errors_total: {
@@ -691,7 +779,7 @@ mod tests {
                 #[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_main_subsystem(stringify!(oxy));
+                        let registry = &mut *::foundations::telemetry::metrics::internal::Registries::get_subsystem(stringify!(oxy), false, true);
 
                         __oxy_Metrics {
                             connections_latency: {
diff --git a/foundations/src/telemetry/metrics/internal.rs b/foundations/src/telemetry/metrics/internal.rs
@@ -6,36 +6,53 @@ use prometheus_client::encoding::text::{encode, EncodeMetric};
 use prometheus_client::registry::Registry;
 use prometools::serde::InfoGauge;
 use std::any::TypeId;
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::ops::DerefMut;
 
 static REGISTRIES: OnceCell<Registries> = OnceCell::new();
 
+enum MetricsServiceName {
+    Prefix(String),
+    Label(String, String),
+}
+
+impl MetricsServiceName {
+    fn new(name: &str, format: ServiceNameFormat) -> Self {
+        let name = name.to_owned();
+        match format {
+            ServiceNameFormat::MetricPrefix => Self::Prefix(name),
+            ServiceNameFormat::LabelWithName(label) => Self::Label(label, name),
+        }
+    }
+}
+
 #[doc(hidden)]
 pub struct Registries {
     // NOTE: we intentionally use a lock without poisoning here to not
     // panic the threads if they just share telemetry with failed thread.
     main: parking_lot::RwLock<Registry>,
     opt: parking_lot::RwLock<Registry>,
     pub(super) info: parking_lot::RwLock<HashMap<TypeId, Box<dyn ErasedInfoMetric>>>,
-    extra_label: Option<(String, String)>,
+    service_name: MetricsServiceName,
     extra_producers: parking_lot::RwLock<Vec<Box<dyn ExtraProducer>>>,
 }
 
 impl Registries {
     pub(super) fn init(service_info: &ServiceInfo, settings: &MetricsSettings) {
-        let extra_label = match &settings.service_name_format {
-            ServiceNameFormat::MetricPrefix => None,
-            ServiceNameFormat::LabelWithName(name) => {
-                Some((name.clone(), service_info.name_in_metrics.clone()))
-            }
-        };
+        let service_name = MetricsServiceName::new(
+            &service_info.name_in_metrics,
+            settings.service_name_format.clone(),
+        );
 
+        // FIXME(nox): Due to prometheus-client 0.18 not supporting the creation of
+        // registries with specific label values, we use `MetricsServiceName::Label`
+        // directly in `Registries::get_subsystem`.
         REGISTRIES.get_or_init(|| Registries {
-            main: new_registry(&service_info.name_in_metrics, &settings.service_name_format),
-            opt: new_registry(&service_info.name_in_metrics, &settings.service_name_format),
+            main: Default::default(),
+            opt: Default::default(),
             info: Default::default(),
-            extra_label,
+            service_name,
             extra_producers: Default::default(),
         });
     }
@@ -72,24 +89,31 @@ impl Registries {
         encode_registry(buffer, &registry)
     }
 
-    pub fn get_main_subsystem(subsystem: &str) -> impl DerefMut<Target = Registry> + '_ {
+    pub fn get_subsystem(
+        subsystem: &str,
+        optional: bool,
+        with_service_prefix: bool,
+    ) -> impl DerefMut<Target = Registry> + 'static {
         let registries = Self::get();
+        let registry = if optional {
+            &registries.opt
+        } else {
+            &registries.main
+        };
 
-        get_subsystem(
-            Self::get().main.write(),
-            subsystem,
-            registries.extra_label.clone(),
-        )
-    }
-
-    pub fn get_opt_subsystem(subsystem: &str) -> impl DerefMut<Target = Registry> + '_ {
-        let registries = Self::get();
+        let mut prefix = Cow::Borrowed(subsystem);
+        if with_service_prefix {
+            if let MetricsServiceName::Prefix(service) = &registries.service_name {
+                prefix = format!("{service}_{subsystem}").into();
+            }
+        }
 
-        get_subsystem(
-            Self::get().opt.write(),
-            subsystem,
-            registries.extra_label.clone(),
-        )
+        parking_lot::RwLockWriteGuard::map(registry.write(), move |mut reg| {
+            if let MetricsServiceName::Label(name, val) = &registries.service_name {
+                reg = reg.sub_registry_with_label((name.into(), val.into()));
+            }
+            reg.sub_registry_with_prefix(prefix)
+        })
     }
 
     pub fn add_extra_producer(&self, producer: Box<dyn ExtraProducer>) {
@@ -98,45 +122,15 @@ impl Registries {
 
     pub(super) fn get() -> &'static Registries {
         REGISTRIES.get_or_init(|| Registries {
-            main: new_registry("undefined", &ServiceNameFormat::MetricPrefix),
-            opt: new_registry("undefined", &ServiceNameFormat::MetricPrefix),
+            main: Default::default(),
+            opt: Default::default(),
             info: Default::default(),
-            extra_label: None,
+            service_name: MetricsServiceName::Prefix("undefined".to_owned()),
             extra_producers: Default::default(),
         })
     }
 }
 
-fn new_registry(
-    service_name_in_metrics: &str,
-    service_name_format: &ServiceNameFormat,
-) -> parking_lot::RwLock<Registry> {
-    parking_lot::RwLock::new(match service_name_format {
-        ServiceNameFormat::MetricPrefix => match service_name_in_metrics {
-            "" => Registry::default(),
-            _ => Registry::with_prefix(service_name_in_metrics),
-        },
-        // FIXME(nox): Due to prometheus-client 0.18 not supporting the creation of
-        // registries with specific label values, we use this service identifier
-        // format directly in `Registries::get_main` and `Registries::get_optional`.
-        ServiceNameFormat::LabelWithName(_) => Registry::default(),
-    })
-}
-
-fn get_subsystem<'a>(
-    registry: parking_lot::RwLockWriteGuard<'a, Registry>,
-    subsystem: &str,
-    extra_label: Option<(String, String)>,
-) -> impl DerefMut<Target = Registry> + 'a {
-    parking_lot::RwLockWriteGuard::map(registry, move |mut registry| {
-        if let Some((name, value)) = extra_label {
-            registry = registry.sub_registry_with_label((name.into(), value.into()));
-        }
-
-        registry.sub_registry_with_prefix(subsystem)
-    })
-}
-
 /// Build and version information
 #[info_metric(crate_path = "crate")]
 pub(super) struct BuildInfo {
diff --git a/foundations/src/telemetry/metrics/mod.rs b/foundations/src/telemetry/metrics/mod.rs
@@ -54,6 +54,9 @@ pub fn collect(settings: &MetricsSettings) -> Result<String> {
 /// name becomes `<global prefix>_<module name>_<bodyless function name>`and function's
 /// Rust doc comment is reported as metric description to Prometheus.
 ///
+/// The `<global_prefix>` can be disabled by passing the `unprefixed` flag to the macro
+/// invocation, like `#[metrics(unprefixed)]`. The module name is a mandatory prefix.
+///
 /// # Labels
 /// Arguments of the bodyless functions become labels for that metric.
 ///
@@ -85,7 +88,6 @@ pub fn collect(settings: &MetricsSettings) -> Result<String> {
 /// `collect_optional` argument of [`collect`] is set to `true`, or, in case the [telemetry server]
 /// is used, if [`MetricsSettings::report_optional`] is set to `true`.
 ///
-///
 /// Can be used for heavy-weight metrics (e.g. with high cardinality) that don't need to be reported
 /// on a regular basis.
 ///
diff --git a/foundations/tests/metrics.rs b/foundations/tests/metrics.rs
