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 quote::quote;
7use syn::{AttributeArgs, ItemFn, Lit, Meta, NestedMeta, parse_macro_input};
8
9/// micromegas_main: Creates a tokio runtime with proper micromegas tracing callbacks and telemetry setup
10///
11/// This is a drop-in replacement for `#[tokio::main]` that automatically configures:
12/// - Tokio runtime with proper micromegas tracing thread lifecycle callbacks
13/// - Telemetry guard with sensible defaults (ctrl-c handling, debug level)
14/// - Automatic authentication configuration from environment variables
15///
16/// # Authentication
17///
18/// The macro automatically configures telemetry authentication based on environment variables:
19///
20/// - **API Key:** Set `MICROMEGAS_INGESTION_API_KEY=your-key`
21/// - **OIDC Client Credentials:** Set `MICROMEGAS_OIDC_TOKEN_ENDPOINT`, `MICROMEGAS_OIDC_CLIENT_ID`, `MICROMEGAS_OIDC_CLIENT_SECRET`
22/// - **No auth:** If no env vars are set, telemetry is sent unauthenticated (requires `--disable-auth` on ingestion server)
23///
24/// # Parameters
25///
26/// - `interop_max_level`: Optional interop max level override (e.g., "info", "debug", "warn")
27/// - `max_level_override`: Optional max level override (e.g., "info", "debug", "warn")
28///
29/// # Examples
30///
31/// ```ignore
32/// use micromegas::tracing::prelude::*;
33///
34/// #[micromegas_main]
35/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
36///     info!("Server starting - telemetry already configured!");
37///     Ok(())
38/// }
39///
40/// #[micromegas_main(interop_max_level = "info")]
41/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
42///     info!("Server starting with info interop level!");
43///     Ok(())
44/// }
45///
46/// #[micromegas_main(max_level_override = "warn", interop_max_level = "info")]
47/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
48///     info!("Server starting with both level overrides!");
49///     Ok(())
50/// }
51/// ```
52#[proc_macro_attribute]
53pub fn micromegas_main(
54    args: proc_macro::TokenStream,
55    input: proc_macro::TokenStream,
56) -> proc_macro::TokenStream {
57    let args = parse_macro_input!(args as AttributeArgs);
58    let function = parse_macro_input!(input as ItemFn);
59
60    // Ensure the function is async and named main
61    if function.sig.asyncness.is_none() {
62        panic!("micromegas_main can only be applied to async functions");
63    }
64
65    if function.sig.ident != "main" {
66        panic!("micromegas_main can only be applied to the main function");
67    }
68
69    // Parse the level override parameters if provided
70    let mut interop_max_level: Option<String> = None;
71    let mut max_level_override: Option<String> = None;
72
73    for arg in args {
74        match arg {
75            NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("interop_max_level") => {
76                if let Lit::Str(lit_str) = nv.lit {
77                    interop_max_level = Some(lit_str.value());
78                } else {
79                    panic!("interop_max_level must be a string literal");
80                }
81            }
82            NestedMeta::Meta(Meta::NameValue(nv)) if nv.path.is_ident("max_level_override") => {
83                if let Lit::Str(lit_str) = nv.lit {
84                    max_level_override = Some(lit_str.value());
85                } else {
86                    panic!("max_level_override must be a string literal");
87                }
88            }
89            _ => panic!(
90                "Unsupported attribute argument. Supported: interop_max_level, max_level_override"
91            ),
92        }
93    }
94
95    let original_block = &function.block;
96    let return_type = &function.sig.output;
97
98    // Helper function to convert level string to LevelFilter token
99    let level_to_filter = |level: &str| -> proc_macro2::TokenStream {
100        match level.to_lowercase().as_str() {
101            "trace" => quote! { micromegas::tracing::levels::LevelFilter::Trace },
102            "debug" => quote! { micromegas::tracing::levels::LevelFilter::Debug },
103            "info" => quote! { micromegas::tracing::levels::LevelFilter::Info },
104            "warn" => quote! { micromegas::tracing::levels::LevelFilter::Warn },
105            "error" => quote! { micromegas::tracing::levels::LevelFilter::Error },
106            "off" => quote! { micromegas::tracing::levels::LevelFilter::Off },
107            _ => {
108                panic!("Invalid level value. Must be one of: trace, debug, info, warn, error, off")
109            }
110        }
111    };
112
113    // Generate the telemetry guard builder with optional level overrides
114    let mut builder_calls = vec![
115        quote! { .with_ctrlc_handling() },
116        quote! { .with_local_sink_max_level(micromegas::tracing::levels::LevelFilter::Debug) },
117        quote! { .with_process_property("version".to_string(), env!("CARGO_PKG_VERSION").to_string()) },
118        quote! { .with_auth_from_env() },
119    ];
120
121    if let Some(level) = max_level_override {
122        let level_filter = level_to_filter(&level);
123        builder_calls.push(quote! { .with_max_level_override(#level_filter) });
124    }
125
126    if let Some(level) = interop_max_level {
127        let level_filter = level_to_filter(&level);
128        builder_calls.push(quote! { .with_interop_max_level_override(#level_filter) });
129    }
130
131    let telemetry_guard_builder = quote! {
132        micromegas::telemetry_sink::TelemetryGuardBuilder::default()
133            #(#builder_calls)*
134            .build()
135    };
136
137    let expanded = quote! {
138        fn main() #return_type {
139            // Check CPU tracing setting before building runtime
140            let cpu_tracing_enabled = std::env::var("MICROMEGAS_ENABLE_CPU_TRACING")
141                .map(|v| v == "true")
142                .unwrap_or(false); // Default to disabled for minimal overhead
143
144            // Set up telemetry guard BEFORE building tokio runtime
145            // This ensures dispatch is initialized before worker threads start
146            let _telemetry_guard = #telemetry_guard_builder;
147
148            // Build the runtime with conditional tracing callbacks
149            let runtime = {
150                use micromegas::tracing::runtime::TracingRuntimeExt;
151                let mut builder = tokio::runtime::Builder::new_multi_thread();
152                builder.enable_all();
153                builder.thread_name(env!("CARGO_PKG_NAME"));
154                if cpu_tracing_enabled {
155                    builder.with_tracing_callbacks();
156                }
157                builder.build().expect("Failed to build tokio runtime")
158            };
159
160            runtime.block_on(async move {
161                // Execute the original main function body
162                #original_block
163            })
164        }
165    };
166
167    proc_macro::TokenStream::from(expanded)
168}
micromegas_proc_macros/lib.rs

micromegas_proc_macros/
lib.rs
