async_inspect_macros/
lib.rs

1//! Procedural macros for async-inspect
2//!
3//! This crate provides attribute macros for automatic instrumentation of async functions.
4
5use proc_macro::TokenStream;
6use quote::quote;
7use syn::{parse_macro_input, visit_mut::VisitMut, Expr, ItemFn};
8
9/// Attribute macro to automatically instrument async functions
10///
11/// # Example
12///
13/// ```rust,ignore
14/// #[async_inspect::trace]
15/// async fn fetch_user(id: u64) -> User {
16///     let profile = fetch_profile(id).await;  // Automatically tracked!
17///     let posts = fetch_posts(id).await;      // Automatically tracked!
18///     User { profile, posts }
19/// }
20/// ```
21///
22/// This macro will:
23/// - Register the function as a tracked task
24/// - Automatically label each `.await` point
25/// - Track execution time
26/// - Report completion or failure
27#[proc_macro_attribute]
28pub fn trace(_attr: TokenStream, item: TokenStream) -> TokenStream {
29    let mut input = parse_macro_input!(item as ItemFn);
30
31    // Ensure it's an async function
32    if input.sig.asyncness.is_none() {
33        return syn::Error::new_spanned(
34            input.sig.fn_token,
35            "#[async_inspect::trace] can only be applied to async functions",
36        )
37        .to_compile_error()
38        .into();
39    }
40
41    let fn_name = &input.sig.ident;
42    let fn_name_str = fn_name.to_string();
43    let vis = &input.vis;
44    let sig = &input.sig;
45
46    // Instrument the function body
47    let mut instrumenter = AwaitInstrumenter {
48        counter: 0,
49        fn_name: fn_name_str.clone(),
50    };
51    instrumenter.visit_block_mut(&mut input.block);
52
53    let instrumented_block = &input.block;
54
55    let output = quote! {
56        #vis #sig {
57            // Register this function as a task
58            let __inspect_task_id = ::async_inspect::inspector::Inspector::global()
59                .register_task(#fn_name_str.to_string());
60
61            ::async_inspect::instrument::set_current_task_id(__inspect_task_id);
62
63            // Execute the original function
64            let __inspect_result = async move #instrumented_block;
65
66            let __result = __inspect_result.await;
67
68            // Mark task as completed
69            ::async_inspect::inspector::Inspector::global().task_completed(__inspect_task_id);
70            ::async_inspect::instrument::clear_current_task_id();
71
72            __result
73        }
74    };
75
76    output.into()
77}
78
79/// Visitor that instruments `.await` expressions
80struct AwaitInstrumenter {
81    counter: usize,
82    fn_name: String,
83}
84
85impl VisitMut for AwaitInstrumenter {
86    fn visit_expr_mut(&mut self, expr: &mut Expr) {
87        // First, recursively visit child expressions BEFORE processing this node
88        // This prevents infinite recursion when we replace await expressions
89        syn::visit_mut::visit_expr_mut(self, expr);
90
91        // Now check if THIS expression is an await that needs instrumentation
92        if let Expr::Await(await_expr) = expr {
93            self.counter += 1;
94            let label = format!("{}::await#{}", self.fn_name, self.counter);
95
96            // Get the source location
97            let location = format!("{}:{}", file!(), line!());
98
99            // Wrap the await with inspection
100            // Clone the base to avoid borrow issues
101            let base = await_expr.base.clone();
102
103            *expr = syn::parse_quote! {
104                {
105                    ::async_inspect::instrument::inspect_await_start(#label, Some(#location.to_string()));
106                    let __result = #base.await;
107                    ::async_inspect::instrument::inspect_await_end(#label);
108                    __result
109                }
110            };
111        }
112    }
113}
114
115/// Attribute macro for inspecting specific code blocks
116///
117/// # Example
118///
119/// ```rust,ignore
120/// #[async_inspect::inspect]
121/// async fn process_data(data: Vec<u8>) -> Result<(), Error> {
122///     // This entire block will be tracked
123///     let parsed = parse(data)?;
124///     let validated = validate(parsed).await?;
125///     Ok(validated)
126/// }
127/// ```
128#[proc_macro_attribute]
129pub fn inspect(_attr: TokenStream, item: TokenStream) -> TokenStream {
130    // For now, just an alias to trace
131    trace(_attr, item)
132}
133
134#[cfg(test)]
135mod tests {
136    #[test]
137    fn test_placeholder() {
138        assert_eq!(2 + 2, 4);
139    }
140}
async_inspect_macros/lib.rs

async_inspect_macros/
lib.rs
