/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 | 82 | fn parse(input: ParseStream) -> syn::Result<Self> { |
37 | 82 | if input.is_empty() { |
38 | 0 | return Ok(Self::None); |
39 | 82 | } |
40 | 82 | let group_str: LitStr = input.parse()?0 ; |
41 | 82 | let group = format_ident!("{}", group_str.value()); |
42 | 82 | Ok(Self::StaticGroupName(group)) |
43 | 82 | } |
44 | | } |
45 | | |
46 | | impl ToTokens for GroupType { |
47 | 380 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
48 | 380 | match self { |
49 | | Self::None => { |
50 | 298 | quote! { "" } |
51 | | } |
52 | 82 | Self::StaticGroupName(group) => quote! { stringify!(#group) }, |
53 | | } |
54 | 380 | .to_tokens(tokens); |
55 | 380 | } |
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 | 380 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
85 | 380 | 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 | 380 | Self::Default => quote! { ::nativelink_metric::MetricKind::Default }, |
90 | | } |
91 | 380 | .to_tokens(tokens); |
92 | 380 | } |
93 | | } |
94 | | |
95 | | /// Holds general information about a specific field that is to be published. |
96 | | struct MetricFieldMetaData<'a> { |
97 | | field_name: &'a Ident, |
98 | | metric_kind: MetricKind, |
99 | | help: Option<LitStr>, |
100 | | group: GroupType, |
101 | | handler: Option<syn::ExprPath>, |
102 | | } |
103 | | |
104 | | impl<'a> MetricFieldMetaData<'a> { |
105 | 380 | fn try_from(field_name: &'a Ident, attr: &Attribute) -> syn::Result<Self> { |
106 | 380 | let mut result = MetricFieldMetaData { |
107 | 380 | field_name, |
108 | 380 | metric_kind: MetricKind::Default, |
109 | 380 | help: None, |
110 | 380 | group: GroupType::None, |
111 | 380 | handler: None, |
112 | 380 | }; |
113 | | // If the attribute is just a path, it has no args, so use defaults. |
114 | 380 | if let syn::Meta::Path(_) = attr.meta { |
115 | 22 | return Ok(result); |
116 | 358 | } |
117 | 360 | attr358 .parse_args_with358 (syn::meta::parser358 (|meta| { |
118 | 360 | if meta.path.is_ident("help") { |
119 | 278 | result.help = meta.value()?0 .parse()?0 ; |
120 | 82 | } else if meta.path.is_ident("kind") { |
121 | 0 | result.metric_kind = meta.value()?.parse()?; |
122 | 82 | } else if meta.path.is_ident("group") { |
123 | 82 | result.group = meta.value()?0 .parse()?0 ; |
124 | 0 | } else if meta.path.is_ident("handler") { |
125 | 0 | result.handler = Some(meta.value()?.parse()?); |
126 | 0 | } |
127 | 360 | Ok(()) |
128 | 360 | }))?0 ; |
129 | 358 | Ok(result) |
130 | 380 | } |
131 | | } |
132 | | |
133 | | /// Holds the template information about the struct. This is needed |
134 | | /// to create the `MetricsComponent` impl. |
135 | | struct Generics<'a> { |
136 | | implementation: ImplGenerics<'a>, |
137 | | ty: TypeGenerics<'a>, |
138 | | where_clause: Option<&'a WhereClause>, |
139 | | } |
140 | | |
141 | | /// Holds metadata about the struct that is having `MetricsComponent` |
142 | | /// implemented. |
143 | | struct MetricStruct<'a> { |
144 | | name: &'a Ident, |
145 | | metric_fields: Vec<MetricFieldMetaData<'a>>, |
146 | | generics: Generics<'a>, |
147 | | } |
148 | | |
149 | | impl ToTokens for MetricStruct<'_> { |
150 | 119 | fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { |
151 | 119 | let name = &self.name; |
152 | 119 | let impl_generics = &self.generics.implementation; |
153 | 119 | let ty_generics = &self.generics.ty; |
154 | 119 | let where_clause = &self.generics.where_clause; |
155 | | |
156 | 380 | let metric_fields119 = self.metric_fields.iter()119 .map119 (|field| { |
157 | 380 | let field_name = &field.field_name; |
158 | 380 | let group = &field.group; |
159 | | |
160 | 380 | let help = field |
161 | 380 | .help |
162 | 380 | .as_ref() |
163 | 380 | .map_or_else(|| quote!102 { "" }102 , |help| quote!278 { #help }278 ); |
164 | | |
165 | 380 | let value = field.handler.as_ref().map_or_else( |
166 | 380 | || quote! { &self.#field_name }, |
167 | 0 | |handler| quote! { &#handler(&self.#field_name) }, |
168 | | ); |
169 | | |
170 | 380 | let metric_kind = &field.metric_kind; |
171 | 380 | quote! { |
172 | | ::nativelink_metric::publish!( |
173 | | stringify!(#field_name), |
174 | | #value, |
175 | | #metric_kind, |
176 | | #help, |
177 | | #group |
178 | | ); |
179 | | } |
180 | 380 | }); |
181 | 119 | quote! { |
182 | | impl #impl_generics ::nativelink_metric::MetricsComponent for #name #ty_generics #where_clause { |
183 | | fn publish(&self, kind: ::nativelink_metric::MetricKind, field_metadata: ::nativelink_metric::MetricFieldData) -> Result<::nativelink_metric::MetricPublishKnownKindData, ::nativelink_metric::Error> { |
184 | | #( #metric_fields )* |
185 | | Ok(::nativelink_metric::MetricPublishKnownKindData::Component) |
186 | | } |
187 | | } |
188 | 119 | }.to_tokens(tokens); |
189 | 119 | } |
190 | | } |
191 | | |
192 | | #[proc_macro_derive(MetricsComponent, attributes(metric))] |
193 | 119 | pub fn metrics_component_derive(input: TokenStream) -> TokenStream { |
194 | 119 | let input = parse_macro_input!(input as DeriveInput); |
195 | 119 | let syn::Data::Struct(data) = &input.data else { |
196 | 0 | panic!("MetricsComponent can only be derived for structs") |
197 | | }; |
198 | | |
199 | 119 | let mut metric_fields = vec![]; |
200 | 119 | match &data.fields { |
201 | 119 | syn::Fields::Named(fields) => { |
202 | 656 | fields.named119 .iter119 ().for_each119 (|field| { |
203 | 656 | field.attrs.iter().for_each(|attr| {651 |
204 | 651 | if attr.path().is_ident("metric") { |
205 | 380 | metric_fields.push( |
206 | 380 | MetricFieldMetaData::try_from(field.ident.as_ref().unwrap(), attr) |
207 | 380 | .unwrap(), |
208 | 380 | ); |
209 | 380 | }271 |
210 | 651 | }); |
211 | 656 | }); |
212 | | } |
213 | | syn::Fields::Unnamed(_) => { |
214 | 0 | panic!("Unnamed fields are not supported"); |
215 | | } |
216 | | syn::Fields::Unit => { |
217 | 0 | panic!("Unit structs are not supported"); |
218 | | } |
219 | | } |
220 | | |
221 | 119 | let (implementation, ty, where_clause) = input.generics.split_for_impl(); |
222 | 119 | let metrics_struct = MetricStruct { |
223 | 119 | name: &input.ident, |
224 | 119 | metric_fields, |
225 | 119 | generics: Generics { |
226 | 119 | implementation, |
227 | 119 | ty, |
228 | 119 | where_clause, |
229 | 119 | }, |
230 | 119 | }; |
231 | | // This line is intentionally left here to make debugging |
232 | | // easier. If you want to see the output of the macro, just |
233 | | // uncomment this line and run the tests. |
234 | | // panic!("{}", quote! { #metrics_struct }); |
235 | 119 | TokenStream::from(quote! { #metrics_struct }) |
236 | 119 | } |