diff --git a/Cargo.toml b/Cargo.toml index 2c1f889..763adea 100644 --- a/Cargo.toml +++ b/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 } diff --git a/derive-register/Cargo.toml b/derive-register/Cargo.toml new file mode 100644 index 0000000..3da6989 --- /dev/null +++ b/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 + diff --git a/derive-register/src/lib.rs b/derive-register/src/lib.rs new file mode 100644 index 0000000..44064c7 --- /dev/null +++ b/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, + ty: Type, + attrs: Vec, + skip: Flag, + unit: Option, + name: Option, + help: Option, +} + +#[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, 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); + ::register(&self, &mut registry); + } + } + }.into() +} + +fn extract_doc_comment(meta: &Meta) -> Option { + 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()) +} diff --git a/derive-register/tests/lib.rs b/derive-register/tests/lib.rs new file mode 100644 index 0000000..e6e828c --- /dev/null +++ b/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, ®istry).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); +} diff --git a/src/registry.rs b/src/registry.rs index c7dec3b..b9a71c1 100644 --- a/src/registry.rs +++ b/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 @@ pub trait Metric: crate::encoding::EncodeMetric + Send + Sync + std::fmt::Debug impl 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 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, H: Into>( + &self, + name: N, + help: H, + unit: Option, + registry: &mut Registry, + ); +} + +impl RegisterField for T +where + T: Metric + Clone, +{ + fn register_field, H: Into>( + &self, + name: N, + help: H, + unit: Option, + registry: &mut Registry, + ) { + if let Some(unit) = unit { + registry.register_with_unit(name, help, unit, self.clone()) + } else { + registry.register(name, help, self.clone()) + } + } +}