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

micromegas_proc_macros/
lib.rs
