greentic_types_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{Attribute, ItemFn, LitStr, ReturnType, Type, meta, parse::Parser, spanned::Spanned};
4
5/// Automatically installs Greentic telemetry at runtime entry-points.
6///
7/// ```ignore
8/// #[greentic_types::telemetry::main(service_name = "runner")]
9/// async fn main() -> anyhow::Result<()> {
10///     // tracing/logging is ready here.
11///     Ok(())
12/// }
13/// ```
14#[proc_macro_attribute]
15pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
16    expand_main(args, item).unwrap_or_else(|err| err.to_compile_error().into())
17}
18
19fn expand_main(args: TokenStream, item: TokenStream) -> syn::Result<TokenStream> {
20    let service_name = parse_service_name(args)?;
21
22    let mut item_fn: ItemFn = syn::parse(item)?;
23    ensure_async(&item_fn)?;
24    ensure_no_args(&item_fn)?;
25
26    strip_self_attr(&mut item_fn.attrs);
27
28    let vis = item_fn.vis.clone();
29    let user_ident = item_fn.sig.ident.clone();
30    let inner_ident = format_ident!("__greentic_types_main");
31    item_fn.sig.ident = inner_ident.clone();
32    item_fn.vis = syn::Visibility::Inherited;
33
34    let generics = item_fn.sig.generics.clone();
35    let where_clause = generics.where_clause.clone();
36    let inputs = item_fn.sig.inputs.clone();
37    let output = item_fn.sig.output.clone();
38    let returns_result = is_result_return(&output);
39    let install_stmt = if returns_result {
40        quote! {
41            ::greentic_types::telemetry::install_telemetry(#service_name)
42                .map_err(::core::convert::Into::into)?;
43        }
44    } else {
45        quote! {
46            ::greentic_types::telemetry::install_telemetry(#service_name)
47                .expect("greentic telemetry initialization failed");
48        }
49    };
50
51    let expanded = quote! {
52        #[::greentic_types::telemetry::__tokio_main]
53        #vis async fn #user_ident #generics (#inputs) #output #where_clause {
54            #install_stmt
55            #inner_ident().await
56        }
57
58        #item_fn
59    };
60
61    Ok(expanded.into())
62}
63
64fn parse_service_name(args: TokenStream) -> syn::Result<LitStr> {
65    let mut service_name = None;
66    let parser = meta::parser(|meta| {
67        if meta.path.is_ident("service_name") {
68            let lit: LitStr = meta.value()?.parse()?;
69            if service_name.is_some() {
70                return Err(meta.error("service_name specified more than once"));
71            }
72            service_name = Some(lit);
73            Ok(())
74        } else {
75            Err(meta.error("expected `service_name = \"...\"`"))
76        }
77    });
78
79    parser.parse2(proc_macro2::TokenStream::from(args))?;
80
81    service_name.ok_or_else(|| {
82        syn::Error::new(
83            proc_macro2::Span::call_site(),
84            "missing `service_name = \"...\"` argument",
85        )
86    })
87}
88
89fn ensure_async(item_fn: &ItemFn) -> syn::Result<()> {
90    if item_fn.sig.asyncness.is_none() {
91        return Err(syn::Error::new(
92            item_fn.sig.span(),
93            "`#[greentic_types::telemetry::main]` requires an `async fn`",
94        ));
95    }
96    Ok(())
97}
98
99fn ensure_no_args(item_fn: &ItemFn) -> syn::Result<()> {
100    if !item_fn.sig.inputs.is_empty() {
101        return Err(syn::Error::new(
102            item_fn.sig.inputs.span(),
103            "`main` must not take arguments",
104        ));
105    }
106    Ok(())
107}
108
109fn strip_self_attr(attrs: &mut Vec<Attribute>) {
110    attrs.retain(|attr| !is_self_attr(attr));
111}
112
113fn is_self_attr(attr: &Attribute) -> bool {
114    let path = attr.path();
115    let segments: Vec<_> = path
116        .segments
117        .iter()
118        .map(|seg| seg.ident.to_string())
119        .collect();
120    segments == ["greentic_types", "telemetry", "main"]
121}
122
123fn is_result_return(output: &ReturnType) -> bool {
124    match output {
125        ReturnType::Type(_, ty) => is_result_type(ty),
126        ReturnType::Default => false,
127    }
128}
129
130fn is_result_type(ty: &Type) -> bool {
131    match ty {
132        Type::Path(type_path) => type_path
133            .path
134            .segments
135            .last()
136            .map(|segment| segment.ident == "Result")
137            .unwrap_or(false),
138        _ => false,
139    }
140}
greentic_types_macros/lib.rs

greentic_types_macros/
lib.rs
