/build/source/nativelink-metric/nativelink-metric-macro-derive/src/lib.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); |
4 | | // you may not use this file except in compliance with the License. |
5 | | // You may obtain a copy of the License at |
6 | | // |
7 | | // See LICENSE file for details |
8 | | // |
9 | | // Unless required by applicable law or agreed to in writing, software |
10 | | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | | // See the License for the specific language governing permissions and |
13 | | // limitations under the License. |
14 | | |
15 | | use core::panic; |
16 | | |
17 | | use proc_macro::TokenStream; |
18 | | use quote::{ToTokens, format_ident, quote}; |
19 | | use syn::parse::{Parse, ParseStream}; |
20 | | use syn::{ |
21 | | Attribute, DeriveInput, Ident, ImplGenerics, LitStr, TypeGenerics, WhereClause, |
22 | | parse_macro_input, |
23 | | }; |
24 | | |
25 | | /// Holds the type of group for the metric. For example, if a metric |
26 | | /// has no group it'll be `None`, if it has a static group name it'll |
27 | | /// be `StaticGroupName(name_of_metric)`. |
28 | | #[derive(Default, Debug)] |
29 | | enum GroupType { |
30 | | #[default] |
31 | | None, |
32 | | StaticGroupName(Ident), |
33 | | } |
34 | | |
35 | | impl Parse for GroupType { |
36 | 80 | fn parse(input: ParseStream) -> syn::Result<Self> { |
37 | 80 | if input.is_empty() { |
38 | 0 | return Ok(Self::None); |
39 | 80 | } |
40 | 80 | let group_str: LitStr = input.parse()?0 ; |
41 | 80 | let group = format_ident!("{}", group_str.value()); |
42 | 80 | Ok(Self::StaticGroupName(group)) |
43 | 80 | } |
44 | | } |
45 | | |
46 | | impl ToTokens for GroupType { |
47 | 358 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
48 | 358 | match self { |
49 | | Self::None => { |
50 | 278 | quote! { "" } |
51 | | } |
52 | 80 | Self::StaticGroupName(group) => quote! { stringify!(#group) }, |
53 | | } |
54 | 358 | .to_tokens(tokens); |
55 | 358 | } |
56 | | } |
57 | | |
58 | | /// Holds the type of the metric. If the metric was not specified |
59 | | /// it'll be `Default`, which will try to resolve the type from the |
60 | | /// [`MetricsComponent::publish()`] method that got executed based on |
61 | | /// the type of the field. |
62 | | #[derive(Debug)] |
63 | | enum MetricKind { |
64 | | Default, |
65 | | Counter, |
66 | | String, |
67 | | Component, |
68 | | } |
69 | | |
70 | | impl Parse for MetricKind { |
71 | 0 | fn parse(input: ParseStream) -> syn::Result<Self> { |
72 | 0 | let kind_str: LitStr = input.parse()?; |
73 | 0 | match kind_str.value().as_str() { |
74 | 0 | "counter" => Ok(Self::Counter), |
75 | 0 | "string" => Ok(Self::String), |
76 | 0 | "component" => Ok(Self::Component), |
77 | 0 | "default" => Ok(Self::Default), |
78 | 0 | _ => Err(syn::Error::new(kind_str.span(), "Invalid metric type")), |
79 | | } |
80 | 0 | } |
81 | | } |
82 | | |
83 | | impl ToTokens for MetricKind { |
84 | 358 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
85 | 358 | match self { |
86 | 0 | Self::Counter => quote! { ::nativelink_metric::MetricKind::Counter }, |
87 | 0 | Self::String => quote! { ::nativelink_metric::MetricKind::String }, |
88 | 0 | Self::Component => quote! { ::nativelink_metric::MetricKind::Component }, |
89 | 358 | Self::Default => quote! { ::nativelink_metric::MetricKind::Default }, |
90 | | } |
91 | 358 | .to_tokens(tokens); |
92 | 358 | } |
93 | | } |
94 | | |
95 | | /// Holds general information about a specific field that is to be published. |
96 | | #[derive(Debug)] |
97 | | struct MetricFieldMetaData<'a> { |
98 | | field_name: &'a Ident, |
99 | | metric_kind: MetricKind, |
100 | | help: Option<LitStr>, |
101 | | group: GroupType, |
102 | | handler: Option<syn::ExprPath>, |
103 | | } |
104 | | |
105 | | impl<'a> MetricFieldMetaData<'a> { |
106 | 358 | fn try_from(field_name: &'a Ident, attr: &Attribute) -> syn::Result<Self> { |
107 | 358 | let mut result = MetricFieldMetaData { |
108 | 358 | field_name, |
109 | 358 | metric_kind: MetricKind::Default, |
110 | 358 | help: None, |
111 | 358 | group: GroupType::None, |
112 | 358 | handler: None, |
113 | 358 | }; |
114 | | // If the attribute is just a path, it has no args, so use defaults. |
115 | 358 | if let syn::Meta::Path(_) = attr.meta { |
116 | 22 | return Ok(result); |
117 | 336 | } |
118 | 338 | attr336 .parse_args_with336 (syn::meta::parser336 (|meta| { |
119 | 338 | if meta.path.is_ident("help") { |
120 | 258 | result.help = meta.value()?0 .parse()?0 ; |
121 | 80 | } else if meta.path.is_ident("kind") { |
122 | 0 | result.metric_kind = meta.value()?.parse()?; |
123 | 80 | } else if meta.path.is_ident("group") { |
124 | 80 | result.group = meta.value()?0 .parse()?0 ; |
125 | 0 | } else if meta.path.is_ident("handler") { |
126 | 0 | result.handler = Some(meta.value()?.parse()?); |
127 | 0 | } |
128 | 338 | Ok(()) |
129 | 338 | }))?0 ; |
130 | 336 | Ok(result) |
131 | 358 | } |
132 | | } |
133 | | |
134 | | /// Holds the template information about the struct. This is needed |
135 | | /// to create the `MetricsComponent` impl. |
136 | | #[derive(Debug)] |
137 | | struct Generics<'a> { |
138 | | implementation: ImplGenerics<'a>, |
139 | | ty: TypeGenerics<'a>, |
140 | | where_clause: Option<&'a WhereClause>, |
141 | | } |
142 | | |
143 | | /// Holds metadata about the struct that is having `MetricsComponent` |
144 | | /// implemented. |
145 | | #[derive(Debug)] |
146 | | struct MetricStruct<'a> { |
147 | | name: &'a Ident, |
148 | | metric_fields: Vec<MetricFieldMetaData<'a>>, |
149 | | generics: Generics<'a>, |
150 | | } |
151 | | |
152 | | impl ToTokens for MetricStruct<'_> { |
153 | 105 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
154 | 105 | let name = &self.name; |
155 | 105 | let impl_generics = &self.generics.implementation; |
156 | 105 | let ty_generics = &self.generics.ty; |
157 | 105 | let where_clause = &self.generics.where_clause; |
158 | | |
159 | 358 | let metric_fields105 = self.metric_fields.iter()105 .map105 (|field| { |
160 | 358 | let field_name = &field.field_name; |
161 | 358 | let group = &field.group; |
162 | | |
163 | 358 | let help = field |
164 | 358 | .help |
165 | 358 | .as_ref() |
166 | 358 | .map_or_else(|| quote!100 { "" }100 , |help| quote!258 { #help }258 ); |
167 | | |
168 | 358 | let value = field.handler.as_ref().map_or_else( |
169 | 358 | || quote! { &self.#field_name }, |
170 | 0 | |handler| quote! { &#handler(&self.#field_name) }, |
171 | | ); |
172 | | |
173 | 358 | let metric_kind = &field.metric_kind; |
174 | 358 | quote! { |
175 | | ::nativelink_metric::publish!( |
176 | | stringify!(#field_name), |
177 | | #value, |
178 | | #metric_kind, |
179 | | #help, |
180 | | #group |
181 | | ); |
182 | | } |
183 | 358 | }); |
184 | 105 | quote! { |
185 | | impl #impl_generics ::nativelink_metric::MetricsComponent for #name #ty_generics #where_clause { |
186 | | fn publish(&self, kind: ::nativelink_metric::MetricKind, field_metadata: ::nativelink_metric::MetricFieldData) -> Result<::nativelink_metric::MetricPublishKnownKindData, ::nativelink_metric::Error> { |
187 | | #( #metric_fields )* |
188 | | Ok(::nativelink_metric::MetricPublishKnownKindData::Component) |
189 | | } |
190 | | } |
191 | 105 | }.to_tokens(tokens); |
192 | 105 | } |
193 | | } |
194 | | |
195 | | #[proc_macro_derive(MetricsComponent, attributes(metric))] |
196 | 105 | pub fn metrics_component_derive(input: TokenStream) -> TokenStream { |
197 | 105 | let input = parse_macro_input!(input as DeriveInput); |
198 | 105 | let syn::Data::Struct(data) = &input.data else { |
199 | 0 | panic!("MetricsComponent can only be derived for structs") |
200 | | }; |
201 | | |
202 | 105 | let mut metric_fields = vec![]; |
203 | 105 | match &data.fields { |
204 | 105 | syn::Fields::Named(fields) => { |
205 | 569 | fields.named105 .iter105 ().for_each105 (|field| { |
206 | 590 | field.attrs.iter()569 .for_each569 (|attr| { |
207 | 590 | if attr.path().is_ident("metric") { |
208 | 358 | metric_fields.push( |
209 | 358 | MetricFieldMetaData::try_from(field.ident.as_ref().unwrap(), attr) |
210 | 358 | .unwrap(), |
211 | 358 | ); |
212 | 358 | }232 |
213 | 590 | }); |
214 | 569 | }); |
215 | | } |
216 | | syn::Fields::Unnamed(_) => { |
217 | 0 | panic!("Unnamed fields are not supported"); |
218 | | } |
219 | | syn::Fields::Unit => { |
220 | 0 | panic!("Unit structs are not supported"); |
221 | | } |
222 | | } |
223 | | |
224 | 105 | let (implementation, ty, where_clause) = input.generics.split_for_impl(); |
225 | 105 | let metrics_struct = MetricStruct { |
226 | 105 | name: &input.ident, |
227 | 105 | metric_fields, |
228 | 105 | generics: Generics { |
229 | 105 | implementation, |
230 | 105 | ty, |
231 | 105 | where_clause, |
232 | 105 | }, |
233 | 105 | }; |
234 | | // This line is intentionally left here to make debugging |
235 | | // easier. If you want to see the output of the macro, just |
236 | | // uncomment this line and run the tests. |
237 | | // panic!("{}", quote! { #metrics_struct }); |
238 | 105 | TokenStream::from(quote! { #metrics_struct }) |
239 | 105 | } |