micromegas_proc_macros/
lib.rs

1//! Top-level procedural macros for micromegas
2//!
3//! This crate provides high-level procedural macros that integrate multiple
4//! micromegas components for a seamless developer experience.
5
6use proc_macro2::TokenStream;
7use quote::quote;
8use syn::{Expr, ExprLit, ItemFn, Lit, Meta};
9
10/// micromegas_main: Creates a tokio runtime with proper micromegas tracing callbacks and telemetry setup
11///
12/// This is a drop-in replacement for `#[tokio::main]` that automatically configures:
13/// - Tokio runtime with proper micromegas tracing thread lifecycle callbacks
14/// - Telemetry guard with sensible defaults (ctrl-c handling, debug level)
15/// - Automatic authentication configuration from environment variables
16///
17/// # Authentication
18///
19/// The macro automatically configures telemetry authentication based on environment variables:
20///
21/// - **API Key:** Set `MICROMEGAS_INGESTION_API_KEY=your-key`
22/// - **OIDC Client Credentials:** Set `MICROMEGAS_OIDC_TOKEN_ENDPOINT`, `MICROMEGAS_OIDC_CLIENT_ID`, `MICROMEGAS_OIDC_CLIENT_SECRET`
23/// - **No auth:** If no env vars are set, telemetry is sent unauthenticated (requires `--disable-auth` on ingestion server)
24///
25/// # Parameters
26///
27/// - `ctrlc_handling`: bool (default: `true`) — enable Ctrl-C graceful shutdown
28/// - `install_log_capture`: bool (default: `false`) — capture `log` crate output
29/// - `interop_max_level`: string (e.g., `"info"`) — interop max level override
30/// - `local_sink_enabled`: bool (default: `true`) — enable local stderr sink
31/// - `local_sink_max_level`: string (default: `"debug"`) — max level for local sink
32/// - `max_level_override`: string (e.g., `"warn"`) — global max level override
33/// - `system_metrics`: bool (default: `true`) — collect system metrics
34/// - `telemetry_url`: string — override the telemetry ingestion URL
35/// - `api_key`: string — embed a literal API key (takes precedence over env-var auth)
36///
37/// # Examples
38///
39/// ```ignore
40/// use micromegas::tracing::prelude::*;
41///
42/// #[micromegas_main]
43/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
44///     info!("Server starting - telemetry already configured!");
45///     Ok(())
46/// }
47///
48/// #[micromegas_main(interop_max_level = "info")]
49/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
50///     info!("Server starting with info interop level!");
51///     Ok(())
52/// }
53///
54/// #[micromegas_main(max_level_override = "warn", interop_max_level = "info")]
55/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
56///     info!("Server starting with both level overrides!");
57///     Ok(())
58/// }
59///
60/// #[micromegas_main(telemetry_url = "http://localhost:9000", api_key = "my-secret-key")]
61/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
62///     info!("Server with explicit URL and embedded API key!");
63///     Ok(())
64/// }
65/// ```
66#[proc_macro_attribute]
67pub fn micromegas_main(
68    args: proc_macro::TokenStream,
69    input: proc_macro::TokenStream,
70) -> proc_macro::TokenStream {
71    expand_micromegas_main(args.into(), input.into())
72        .unwrap_or_else(|err| err.to_compile_error())
73        .into()
74}
75
76fn expand_micromegas_main(
77    args: TokenStream,
78    input: TokenStream,
79) -> Result<TokenStream, syn::Error> {
80    use syn::parse::Parser;
81
82    let args: Vec<Meta> = syn::punctuated::Punctuated::<Meta, syn::Token![,]>::parse_terminated
83        .parse2(args)?
84        .into_iter()
85        .collect();
86
87    let function: ItemFn = syn::parse2(input)?;
88
89    if function.sig.asyncness.is_none() {
90        return Err(syn::Error::new_spanned(
91            function.sig.fn_token,
92            "micromegas_main can only be applied to async functions",
93        ));
94    }
95
96    if function.sig.ident != "main" {
97        return Err(syn::Error::new_spanned(
98            &function.sig.ident,
99            "micromegas_main can only be applied to the main function",
100        ));
101    }
102
103    let mut interop_max_level: Option<syn::LitStr> = None;
104    let mut max_level_override: Option<syn::LitStr> = None;
105    let mut ctrlc_handling: bool = true;
106    let mut local_sink_enabled: bool = true;
107    let mut local_sink_max_level: Option<syn::LitStr> = None;
108    let mut install_log_capture: bool = false;
109    let mut system_metrics: bool = true;
110    let mut telemetry_url: Option<String> = None;
111    let mut api_key: Option<String> = None;
112
113    for arg in args {
114        match arg {
115            Meta::NameValue(nv) if nv.path.is_ident("interop_max_level") => {
116                if let Expr::Lit(ExprLit {
117                    lit: Lit::Str(lit_str),
118                    ..
119                }) = &nv.value
120                {
121                    interop_max_level = Some(lit_str.clone());
122                } else {
123                    return Err(syn::Error::new_spanned(
124                        &nv.value,
125                        "interop_max_level must be a string literal",
126                    ));
127                }
128            }
129            Meta::NameValue(nv) if nv.path.is_ident("max_level_override") => {
130                if let Expr::Lit(ExprLit {
131                    lit: Lit::Str(lit_str),
132                    ..
133                }) = &nv.value
134                {
135                    max_level_override = Some(lit_str.clone());
136                } else {
137                    return Err(syn::Error::new_spanned(
138                        &nv.value,
139                        "max_level_override must be a string literal",
140                    ));
141                }
142            }
143            Meta::NameValue(nv) if nv.path.is_ident("ctrlc_handling") => {
144                if let Expr::Lit(ExprLit {
145                    lit: Lit::Bool(lit_bool),
146                    ..
147                }) = &nv.value
148                {
149                    ctrlc_handling = lit_bool.value();
150                } else {
151                    return Err(syn::Error::new_spanned(
152                        &nv.value,
153                        "ctrlc_handling must be a bool literal",
154                    ));
155                }
156            }
157            Meta::NameValue(nv) if nv.path.is_ident("local_sink_enabled") => {
158                if let Expr::Lit(ExprLit {
159                    lit: Lit::Bool(lit_bool),
160                    ..
161                }) = &nv.value
162                {
163                    local_sink_enabled = lit_bool.value();
164                } else {
165                    return Err(syn::Error::new_spanned(
166                        &nv.value,
167                        "local_sink_enabled must be a bool literal",
168                    ));
169                }
170            }
171            Meta::NameValue(nv) if nv.path.is_ident("local_sink_max_level") => {
172                if let Expr::Lit(ExprLit {
173                    lit: Lit::Str(lit_str),
174                    ..
175                }) = &nv.value
176                {
177                    local_sink_max_level = Some(lit_str.clone());
178                } else {
179                    return Err(syn::Error::new_spanned(
180                        &nv.value,
181                        "local_sink_max_level must be a string literal",
182                    ));
183                }
184            }
185            Meta::NameValue(nv) if nv.path.is_ident("install_log_capture") => {
186                if let Expr::Lit(ExprLit {
187                    lit: Lit::Bool(lit_bool),
188                    ..
189                }) = &nv.value
190                {
191                    install_log_capture = lit_bool.value();
192                } else {
193                    return Err(syn::Error::new_spanned(
194                        &nv.value,
195                        "install_log_capture must be a bool literal",
196                    ));
197                }
198            }
199            Meta::NameValue(nv) if nv.path.is_ident("system_metrics") => {
200                if let Expr::Lit(ExprLit {
201                    lit: Lit::Bool(lit_bool),
202                    ..
203                }) = &nv.value
204                {
205                    system_metrics = lit_bool.value();
206                } else {
207                    return Err(syn::Error::new_spanned(
208                        &nv.value,
209                        "system_metrics must be a bool literal",
210                    ));
211                }
212            }
213            Meta::NameValue(nv) if nv.path.is_ident("telemetry_url") => {
214                if let Expr::Lit(ExprLit {
215                    lit: Lit::Str(lit_str),
216                    ..
217                }) = &nv.value
218                {
219                    telemetry_url = Some(lit_str.value());
220                } else {
221                    return Err(syn::Error::new_spanned(
222                        &nv.value,
223                        "telemetry_url must be a string literal",
224                    ));
225                }
226            }
227            Meta::NameValue(nv) if nv.path.is_ident("api_key") => {
228                if let Expr::Lit(ExprLit {
229                    lit: Lit::Str(lit_str),
230                    ..
231                }) = &nv.value
232                {
233                    api_key = Some(lit_str.value());
234                } else {
235                    return Err(syn::Error::new_spanned(
236                        &nv.value,
237                        "api_key must be a string literal",
238                    ));
239                }
240            }
241            other => {
242                return Err(syn::Error::new_spanned(
243                    &other,
244                    "Unsupported attribute argument. Supported: api_key, ctrlc_handling, install_log_capture, interop_max_level, local_sink_enabled, local_sink_max_level, max_level_override, system_metrics, telemetry_url",
245                ));
246            }
247        }
248    }
249
250    let original_block = &function.block;
251    let return_type = &function.sig.output;
252
253    let level_to_filter = |lit: &syn::LitStr| -> Result<TokenStream, syn::Error> {
254        Ok(match lit.value().to_lowercase().as_str() {
255            "trace" => quote! { micromegas::tracing::levels::LevelFilter::Trace },
256            "debug" => quote! { micromegas::tracing::levels::LevelFilter::Debug },
257            "info" => quote! { micromegas::tracing::levels::LevelFilter::Info },
258            "warn" => quote! { micromegas::tracing::levels::LevelFilter::Warn },
259            "error" => quote! { micromegas::tracing::levels::LevelFilter::Error },
260            "off" => quote! { micromegas::tracing::levels::LevelFilter::Off },
261            _ => {
262                return Err(syn::Error::new_spanned(
263                    lit,
264                    "Invalid level value. Must be one of: trace, debug, info, warn, error, off",
265                ));
266            }
267        })
268    };
269
270    let mut builder_calls = vec![quote! {
271        .with_process_property("version".to_string(), env!("CARGO_PKG_VERSION").to_string())
272    }];
273
274    if ctrlc_handling {
275        builder_calls.push(quote! { .with_ctrlc_handling() });
276    }
277
278    if !local_sink_enabled {
279        builder_calls.push(quote! { .with_local_sink_enabled(false) });
280    }
281
282    {
283        let level_filter = match &local_sink_max_level {
284            Some(lit) => level_to_filter(lit)?,
285            None => quote! { micromegas::tracing::levels::LevelFilter::Debug },
286        };
287        builder_calls.push(quote! { .with_local_sink_max_level(#level_filter) });
288    }
289
290    if install_log_capture {
291        builder_calls.push(quote! { .with_install_log_capture(true) });
292    }
293
294    if !system_metrics {
295        builder_calls.push(quote! { .with_system_metrics_enabled(false) });
296    }
297
298    if let Some(url) = telemetry_url {
299        builder_calls.push(quote! { .with_telemetry_sink_url(#url.to_string()) });
300    }
301
302    if let Some(key) = api_key {
303        builder_calls.push(quote! {
304            .with_request_decorator(std::boxed::Box::new(move || std::sync::Arc::new(
305                micromegas::telemetry_sink::api_key_decorator::ApiKeyRequestDecorator::new(#key.to_string())
306            )))
307        });
308    } else {
309        builder_calls.push(quote! { .with_auth_from_env() });
310    }
311
312    if let Some(lit) = &max_level_override {
313        let level_filter = level_to_filter(lit)?;
314        builder_calls.push(quote! { .with_max_level_override(#level_filter) });
315    }
316
317    if let Some(lit) = &interop_max_level {
318        let level_filter = level_to_filter(lit)?;
319        builder_calls.push(quote! { .with_interop_max_level_override(#level_filter) });
320    }
321
322    let telemetry_guard_builder = quote! {
323        micromegas::telemetry_sink::TelemetryGuardBuilder::default()
324            #(#builder_calls)*
325            .build()
326    };
327
328    Ok(quote! {
329        fn main() #return_type {
330            let cpu_tracing_enabled = std::env::var("MICROMEGAS_ENABLE_CPU_TRACING")
331                .map(|v| v == "true")
332                .unwrap_or(false);
333
334            let _telemetry_guard = #telemetry_guard_builder;
335
336            let runtime = {
337                use micromegas::tracing::runtime::TracingRuntimeExt;
338                let mut builder = tokio::runtime::Builder::new_multi_thread();
339                builder.enable_all();
340                builder.thread_name(env!("CARGO_PKG_NAME"));
341                if cpu_tracing_enabled {
342                    builder.with_tracing_callbacks();
343                }
344                builder.build().expect("Failed to build tokio runtime")
345            };
346
347            runtime.block_on(async move {
348                #original_block
349            })
350        }
351    })
352}
353
354#[cfg(test)]
355mod tests {
356    use super::expand_micromegas_main;
357    use quote::quote;
358
359    fn expand(args: proc_macro2::TokenStream) -> String {
360        let input = quote! { async fn main() {} };
361        expand_micromegas_main(args, input)
362            .expect("expansion should succeed")
363            .to_string()
364    }
365
366    #[test]
367    fn default_produces_standard_calls() {
368        let out = expand(quote! {});
369        assert!(out.contains("with_auth_from_env"));
370        assert!(out.contains("with_ctrlc_handling"));
371        assert!(out.contains("with_local_sink_max_level"));
372    }
373
374    #[test]
375    fn api_key_replaces_env_auth() {
376        let out = expand(quote! { api_key = "secret" });
377        assert!(out.contains("ApiKeyRequestDecorator"));
378        assert!(!out.contains("with_auth_from_env"));
379    }
380
381    #[test]
382    fn ctrlc_handling_false_omits_call() {
383        let out = expand(quote! { ctrlc_handling = false });
384        assert!(!out.contains("with_ctrlc_handling"));
385    }
386
387    #[test]
388    fn telemetry_url_emits_call() {
389        let out = expand(quote! { telemetry_url = "http://localhost:9000" });
390        assert!(out.contains("with_telemetry_sink_url"));
391    }
392
393    #[test]
394    fn local_sink_disabled_emits_call() {
395        let out = expand(quote! { local_sink_enabled = false });
396        assert!(out.contains("with_local_sink_enabled"));
397    }
398
399    #[test]
400    fn system_metrics_false_emits_call() {
401        let out = expand(quote! { system_metrics = false });
402        assert!(out.contains("with_system_metrics_enabled"));
403    }
404
405    #[test]
406    fn install_log_capture_true_emits_call() {
407        let out = expand(quote! { install_log_capture = true });
408        assert!(out.contains("with_install_log_capture"));
409    }
410
411    #[test]
412    fn local_sink_max_level_custom_emits_correct_filter() {
413        let out = expand(quote! { local_sink_max_level = "info" });
414        assert!(out.contains("LevelFilter :: Info"));
415    }
416
417    fn expand_err(args: proc_macro2::TokenStream) -> syn::Error {
418        let input = quote! { async fn main() {} };
419        expand_micromegas_main(args, input).expect_err("expansion should fail")
420    }
421
422    #[test]
423    fn bad_ctrlc_type_is_error() {
424        let err = expand_err(quote! { ctrlc_handling = "not_a_bool" });
425        assert_eq!(err.to_string(), "ctrlc_handling must be a bool literal");
426    }
427
428    #[test]
429    fn unknown_arg_is_error() {
430        let err = expand_err(quote! { unknown_arg = true });
431        assert!(err.to_string().contains("Unsupported attribute argument"));
432    }
433
434    #[test]
435    fn invalid_level_is_error() {
436        let err = expand_err(quote! { max_level_override = "verbose" });
437        assert!(err.to_string().contains("Invalid level value"));
438    }
439
440    #[test]
441    fn non_async_fn_is_error() {
442        let err = expand_micromegas_main(quote! {}, quote! { fn main() {} })
443            .expect_err("non-async main should fail");
444        assert!(err.to_string().contains("async functions"));
445    }
446
447    #[test]
448    fn malformed_args_is_error() {
449        // Garbage tokens that are not valid attribute meta items.
450        let err = expand_micromegas_main(quote! { = = = }, quote! { async fn main() {} })
451            .expect_err("malformed args should fail");
452        // A parse error carries a span-anchored message rather than a panic.
453        assert!(!err.to_string().is_empty());
454    }
455}
micromegas_proc_macros/lib.rs

micromegas_proc_macros/
lib.rs
