bity_ic_canister_tracing_macros/
lib.rs

1//! Module for procedural macros that add tracing capabilities to canister functions.
2//!
3//! This module provides macros that automatically add tracing instrumentation to functions,
4//! making it easier to debug and monitor canister behavior. It wraps functions with tracing
5//! capabilities while preserving their original functionality.
6//!
7//! # Example
8//! ```
9//! use bity_ic_canister_tracing_macros::trace;
10//!
11//! #[trace]
12//! async fn my_function(arg1: u64, arg2: String) -> Result<(), String> {
13//!     // Function implementation
14//!     Ok(())
15//! }
16//! ```
17
18use proc_macro::TokenStream;
19use proc_macro2::Ident;
20use quote::{format_ident, quote};
21use syn::{parse_macro_input, FnArg, ItemFn, Pat, PatIdent, PatType, Signature};
22
23/// A procedural macro attribute that adds tracing capabilities to a function.
24///
25/// This macro wraps the target function with tracing instrumentation, automatically
26/// logging function entry, arguments, and return values at the trace level.
27///
28/// # Usage
29/// Add the `#[trace]` attribute above any function you want to trace:
30/// ```rust
31/// #[trace]
32/// fn my_function(arg: u64) -> u64 {
33///     arg * 2
34/// }
35/// ```
36///
37/// # Features
38/// * Automatically traces function entry and exit
39/// * Logs all function arguments
40/// * Logs the return value
41/// * Works with both synchronous and asynchronous functions
42/// * Preserves the original function signature
43#[proc_macro_attribute]
44pub fn trace(_: TokenStream, item: TokenStream) -> TokenStream {
45    let mut inner = parse_macro_input!(item as ItemFn);
46
47    // We will wrap the original fn in a new fn whose signature matches the original fn
48    #[allow(clippy::redundant_clone)] // clippy doesn't realise that this is used in the macro
49    let wrapper_sig = inner.sig.clone();
50
51    // Change the name of the inner fn so that it doesn't clash with the wrapper fn
52    let inner_method_name = format_ident!("{}_inner_", inner.sig.ident);
53    inner.sig.ident = inner_method_name.clone();
54
55    let is_async = inner.sig.asyncness.is_some();
56    let arg_names = get_arg_names(&inner.sig);
57
58    let function_call = if is_async {
59        quote! { #inner_method_name ( #(#arg_names),* ) .await }
60    } else {
61        quote! { #inner_method_name ( #(#arg_names),* ) }
62    };
63
64    let expanded = quote! {
65        #[allow(unused_mut)]
66        #[tracing::instrument(level = "trace")]
67        #wrapper_sig {
68            let result = #function_call;
69            tracing::trace!(?result);
70            result
71        }
72        #inner
73    };
74
75    TokenStream::from(expanded)
76}
77
78/// Extracts argument names from a function signature.
79///
80/// This helper function processes a function signature to extract the names of all arguments,
81/// including the receiver (self) if present.
82///
83/// # Arguments
84/// * `signature` - The function signature to process
85///
86/// # Returns
87/// A vector of identifiers representing the argument names
88///
89/// # Panics
90/// Panics if it encounters an argument pattern that it cannot process
91fn get_arg_names(signature: &Signature) -> Vec<Ident> {
92    signature
93        .inputs
94        .iter()
95        .map(|arg| match arg {
96            FnArg::Receiver(r) => r.self_token.into(),
97            FnArg::Typed(PatType { pat, .. }) => {
98                if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() {
99                    ident.clone()
100                } else {
101                    panic!("Unable to determine arg name");
102                }
103            }
104        })
105        .collect()
106}
bity_ic_canister_tracing_macros/lib.rs

bity_ic_canister_tracing_macros/
lib.rs
