Derive macro to register struct of metrics

Cargo.toml

@@ -15,13 +15,14 @@ default = []
 protobuf = ["dep:prost", "dep:prost-types", "dep:prost-build"]
 
 [workspace]
-members = ["derive-encode"]
+members = ["derive-encode", "derive-register"]
 
 [dependencies]
 dtoa = "1.0"
 itoa = "1.0"
 parking_lot = "0.12"
 prometheus-client-derive-encode = { version = "0.4.1", path = "derive-encode" }
+prometheus-client-derive-register = { version = "0.1.0", path = "derive-register" }
 prost = { version = "0.12.0", optional = true }
 prost-types = { version = "0.12.0", optional = true }
 

derive-register/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "prometheus-client-derive-register"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+darling = "0.20.10"
+proc-macro2 = "1"
+quote = "1"
+syn = "2"
+
+[dev-dependencies]
+prometheus-client = { path = "../" }
+
+[lib]
+proc-macro = true
+

derive-register/src/lib.rs

@@ -0,0 +1,114 @@
+use darling::{ast, util::Flag, FromDeriveInput, FromField};
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{DeriveInput, Expr, Generics, Ident, Lit, Meta, Type};
+
+#[derive(Debug, FromDeriveInput)]
+#[darling(attributes(register), supports(struct_named))]
+struct Register {
+    ident: Ident,
+    generics: Generics,
+    data: ast::Data<(), RegisterField>,
+}
+
+#[derive(Debug, FromField)]
+#[darling(attributes(register), forward_attrs(doc))]
+struct RegisterField {
+    ident: Option<Ident>,
+    ty: Type,
+    attrs: Vec<syn::Attribute>,
+    skip: Flag,
+    unit: Option<String>,
+    name: Option<String>,
+    help: Option<String>,
+}
+
+#[proc_macro_derive(Register, attributes(register))]
+pub fn derive_register(input: TokenStream) -> TokenStream {
+    let ast: DeriveInput = syn::parse(input).unwrap();
+    let info = Register::from_derive_input(&ast).unwrap();
+
+    let name = info.ident;
+    let (impl_generics, ty_generics, where_clause) = info.generics.split_for_impl();
+
+    let field_register = info
+        .data
+        .take_struct()
+        .unwrap()
+        .into_iter()
+        .filter(|x| !x.skip.is_present())
+        .map(|field| {
+            let mut help = String::new();
+            for attr in field.attrs {
+                let path = attr.path();
+                if path.is_ident("doc") && help.is_empty() {
+                    if let Some(doc) = extract_doc_comment(&attr.meta) {
+                        help = doc.trim().to_string();
+                    }
+                }
+            }
+
+            if let Some(custom_help) = field.help {
+                help = custom_help;
+            }
+
+            let ident = field.ident.unwrap();
+            let ty = field.ty;
+            let name = if let Some(name) = field.name {
+                name
+            } else {
+                ident.to_string()
+            };
+
+            let unit = if let Some(unit) = field.unit {
+                quote!(Some(::prometheus_client::registry::Unit::Other(#unit.to_string())))
+            } else {
+                quote!(None)
+            };
+
+            quote! {
+                <#ty as ::prometheus_client::registry::RegisterField>::register_field(
+                    &self.#ident,
+                    #name,
+                    #help,
+                    #unit,
+                    registry,
+                )
+            }
+        });
+
+    quote! {
+        impl #impl_generics ::prometheus_client::registry::Register for #name #ty_generics #where_clause {
+            fn register(&self, registry: &mut ::prometheus_client::registry::Registry) {
+                #(#field_register);*
+            }
+        }
+
+        impl #impl_generics ::prometheus_client::registry::RegisterField for #name #ty_generics #where_clause {
+            fn register_field<N: ::std::convert::Into<::std::string::String>, H: ::std::convert::Into<::std::string::String>>(
+                &self,
+                name: N,
+                help: H,
+                unit: Option<::prometheus_client::registry::Unit>,
+                registry: &mut ::prometheus_client::registry::Registry)
+            {
+                let name = name.into();
+                let mut registry = registry.sub_registry_with_prefix(name);
+                <Self as ::prometheus_client::registry::Register>::register(&self, &mut registry);
+            }
+        }
+    }.into()
+}
+
+fn extract_doc_comment(meta: &Meta) -> Option<String> {
+    let Meta::NameValue(nv) = meta else {
+        return None;
+    };
+    let Expr::Lit(lit) = &nv.value else {
+        return None;
+    };
+    let Lit::Str(lit_str) = &lit.lit else {
+        return None;
+    };
+    Some(lit_str.value())
+}

derive-register/tests/lib.rs

@@ -0,0 +1,60 @@
+use prometheus_client::{
+    encoding::text::encode,
+    metrics::{counter::Counter, gauge::Gauge},
+    registry::{Register, RegisterDefault, Registry},
+};
+
+#[derive(Register, Default)]
+struct Metrics {
+    /// This is my counter
+    my_counter: Counter,
+    nested: NestedMetrics,
+    #[register(skip)]
+    skipped: Counter,
+    #[register(unit = "bytes")]
+    custom_unit: Counter,
+    #[register(name = "my_custom_name")]
+    custom_name: Counter,
+    /// This will get ignored
+    #[register(help = "my custom help")]
+    custom_help: Counter,
+}
+
+#[derive(Register, Default)]
+struct NestedMetrics {
+    /// This is my gauge
+    my_gauge: Gauge,
+}
+
+#[test]
+fn basic_flow() {
+    let mut registry = Registry::default();
+
+    let metrics = Metrics::register_default(&mut registry);
+
+    metrics.my_counter.inc();
+    metrics.nested.my_gauge.set(23);
+
+    // Encode all metrics in the registry in the text format.
+    let mut buffer = String::new();
+    encode(&mut buffer, &registry).unwrap();
+
+    let expected = "# HELP my_counter This is my counter.\n".to_owned()
+        + "# TYPE my_counter counter\n"
+        + "my_counter_total 1\n"
+        + "# HELP custom_unit_bytes .\n"
+        + "# TYPE custom_unit_bytes counter\n"
+        + "# UNIT custom_unit_bytes bytes\n"
+        + "custom_unit_bytes_total 0\n"
+        + "# HELP my_custom_name .\n"
+        + "# TYPE my_custom_name counter\n"
+        + "my_custom_name_total 0\n"
+        + "# HELP custom_help my custom help.\n"
+        + "# TYPE custom_help counter\n"
+        + "custom_help_total 0\n"
+        + "# HELP nested_my_gauge This is my gauge.\n"
+        + "# TYPE nested_my_gauge gauge\n"
+        + "nested_my_gauge 23\n"
+        + "# EOF\n";
+    assert_eq!(expected, buffer);
+}

src/registry.rs

@@ -2,6 +2,8 @@
 //!
 //! See [`Registry`] for details.
 
+pub use prometheus_client_derive_register::*;
+
 use std::borrow::Cow;
 
 use crate::collector::Collector;
@@ -389,3 +391,51 @@
 
 impl<T> Metric for T where T: crate::encoding::EncodeMetric + Send + Sync + std::fmt::Debug + 'static
 {}
+
+pub trait Register {
+    fn register(&self, registry: &mut Registry);
+}
+
+pub trait RegisterDefault {
+    fn register_default(registry: &mut Registry) -> Self;
+}
+
+impl<T> RegisterDefault for T
+where
+    T: Register + Default,
+{
+    fn register_default(registry: &mut Registry) -> Self {
+        let this = Self::default();
+        this.register(registry);
+        this
+    }
+}
+
+pub trait RegisterField {
+    fn register_field<N: Into<String>, H: Into<String>>(
+        &self,
+        name: N,
+        help: H,
+        unit: Option<Unit>,
+        registry: &mut Registry,
+    );
+}
+
+impl<T> RegisterField for T
+where
+    T: Metric + Clone,
+{
+    fn register_field<N: Into<String>, H: Into<String>>(
+        &self,
+        name: N,
+        help: H,
+        unit: Option<Unit>,
+        registry: &mut Registry,
+    ) {
+        if let Some(unit) = unit {
+            registry.register_with_unit(name, help, unit, self.clone())
+        } else {
+            registry.register(name, help, self.clone())
+        }
+    }
+}