support instrument attribute macros

.github/workflows/publish.yml

@@ -17,6 +17,8 @@ jobs:
       - uses: actions/checkout@v3
 
       - name: Publish to crates.io
-        run: cargo publish
+        run: |
+          cargo publish -p await-tree-attributes
+          cargo publish -p await-tree
         env:
           CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}

CHANGELOG.md

@@ -5,6 +5,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+- Added `#[instrument("<fmt>", args)]` attribute macro for automatic instrumentation of async functions ([#16](https://github.com/risingwavelabs/await-tree/pull/32))
+
 ## [0.3.0] - 2025-04-09
 
 ### Added

Cargo.toml

@@ -1,3 +1,6 @@
+[workspace]
+members = [".", "await-tree-attributes"]
+
 [package]
 name = "await-tree"
 version = "0.3.1"
@@ -12,8 +15,10 @@ license = "Apache-2.0"
 [features]
 serde = ["dep:serde", "flexstr/serde"]
 tokio = ["dep:tokio"]
+attributes = ["dep:await-tree-attributes"]
 
 [dependencies]
+await-tree-attributes = { path = "await-tree-attributes", version = "0.1", optional = true }
 coarsetime = "0.1"
 derive_builder = "0.20"
 easy-ext = "1"
@@ -59,3 +64,7 @@ required-features = ["tokio"]
 [[example]]
 name = "global"
 required-features = ["tokio"]
+
+[[example]]
+name = "instrument"
+required-features = ["tokio", "attributes"]

README.md

@@ -79,6 +79,22 @@ println!("{tree}");
   });
   ```
 
+- `attributes`: Enables the `#[instrument]` attribute macro for automatic instrumentation of async functions. This provides a convenient way to add await-tree spans to functions without manual instrumentation.
+
+  ```rust
+  // Enable the attributes feature in Cargo.toml
+  // await-tree = { version = "<version>", features = ["attributes"] }
+
+  // Then you can use the instrument attribute
+  use await_tree::instrument;
+
+  #[instrument("my_function({})", arg)]
+  async fn my_function(arg: i32) {
+      // Your async code here
+      work().await;
+  }
+  ```
+
 ## Compared to `async-backtrace`
 
 [`tokio-rs/async-backtrace`](https://github.com/tokio-rs/async-backtrace) is a similar crate that also provides the ability to dump the execution tree of async tasks. Here are some differences between `await-tree` and `async-backtrace`:

await-tree-attributes/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "await-tree-attributes"
+version = "0.1.0"
+edition = "2021"
+description = "Procedural attributes for await-tree instrumentation"
+repository = "https://github.com/risingwavelabs/await-tree"
+keywords = ["async", "tokio", "backtrace", "actor", "attributes"]
+categories = ["development-tools::debugging"]
+license = "Apache-2.0"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1.0"
+quote = "1.0"
+syn = { version = "2.0", features = ["full", "extra-traits"] }
+
+[dev-dependencies]
+tokio = { version = "1", features = ["full"] }
+await-tree = { path = ".." }

await-tree-attributes/README.md

@@ -0,0 +1,114 @@
+# await-tree-attributes
+
+Procedural attributes for the [`await-tree`](https://crates.io/crates/await-tree) crate.
+
+## Overview
+
+This crate provides the `#[instrument]` attribute macro that automatically instruments async functions with await-tree spans, similar to how `tracing::instrument` works but specifically designed for await-tree.
+
+## Usage
+
+Add this to your `Cargo.toml`:
+
+```toml
+[dependencies]
+await-tree = { version = "0.3", features = ["attributes"] }
+```
+
+Then use the `#[instrument]` attribute on your async functions:
+
+```rust
+use await_tree::{instrument, InstrumentAwait};
+
+#[instrument("fetch_data({})", id)]
+async fn fetch_data(id: u32) -> String {
+    // Your async code here
+    format!("data_{}", id)
+}
+
+#[instrument(long_running, verbose, "complex_task({}, {})", name, value)]
+async fn complex_task(name: &str, value: i32) -> String {
+    format!("{}: {}", name, value)
+}
+
+#[instrument]
+async fn simple_function() -> String {
+    "hello".to_string()
+}
+```
+
+## Attribute Expansion
+
+The `#[instrument]` macro transforms your async function by:
+
+1. Creating an await-tree span with the provided format arguments
+2. Wrapping the original function body in an async block
+3. Instrumenting the async block with the span
+
+For example:
+
+```rust
+#[instrument("span_name({})", arg1)]
+async fn foo(arg1: i32, arg2: String) {
+    // original function body
+}
+```
+
+Expands to:
+
+```rust
+async fn foo(arg1: i32, arg2: String) {
+    let span = await_tree::span!("span_name({})", arg1);
+    let fut = async move {
+        // original function body
+    };
+    fut.instrument_await(span).await
+}
+```
+
+## Features
+
+- **Format arguments**: Pass format strings and arguments just like `format!()` or `println!()`
+- **No argument parsing**: Format arguments are passed directly to `await_tree::span!()` without modification
+- **Function name fallback**: If no arguments are provided, uses the function name as the span name
+- **Preserves function attributes**: All function attributes and visibility modifiers are preserved
+- **Method chaining**: Support for chaining any method calls on the span
+
+### Method Chaining
+
+You can chain method calls on the span by including identifiers before the format arguments:
+
+```rust
+// Chain span methods
+#[instrument(long_running, "slow_task")]
+async fn slow_task() { /* ... */ }
+
+// Chain multiple methods
+#[instrument(long_running, verbose, "complex_task({})", id)]
+async fn complex_task(id: u32) { /* ... */ }
+
+// Method calls without format args
+#[instrument(long_running, verbose)]
+async fn keywords_only() { /* ... */ }
+
+// Any method name works (will fail at compile time if method doesn't exist)
+#[instrument(custom_attribute, "task")]
+async fn custom_task() { /* ... */ }
+```
+
+The identifiers are processed in order and result in method calls on the span:
+- `long_running` → `.long_running()`
+- `verbose` → `.verbose()`
+- `custom_attribute` → `.custom_attribute()`
+
+If a method doesn't exist on the `Span` type, the code will fail to compile with a clear error message.
+
+## Requirements
+
+- The macro can only be applied to `async` functions
+- You must import `InstrumentAwait` trait to use the generated code
+- The `attributes` feature must be enabled in the `await-tree` dependency
+
+## License
+
+Licensed under the Apache License, Version 2.0.

await-tree-attributes/src/lib.rs

@@ -0,0 +1,171 @@
+// Copyright 2025 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Procedural attributes for await-tree instrumentation.
+
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{parse_macro_input, Ident, ItemFn, Token};
+
+/// Parse the attribute arguments to extract method calls and format args
+struct InstrumentArgs {
+    method_calls: Vec<Ident>,
+    format_args: Option<proc_macro2::TokenStream>,
+}
+
+impl syn::parse::Parse for InstrumentArgs {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        let mut method_calls = Vec::new();
+        let mut format_args = None;
+
+        // Parse identifiers first (these will become method calls)
+        while input.peek(Ident) {
+            // Look ahead to see if this looks like a method call identifier
+            let fork = input.fork();
+            let ident: Ident = fork.parse()?;
+
+            // Check if the next token after the identifier is a comma or end
+            // If it's something else (like a parenthesis or string), treat as format args
+            if fork.peek(Token![,]) || fork.is_empty() {
+                // This is a method call identifier
+                input.parse::<Ident>()?; // consume the identifier
+                method_calls.push(ident);
+                if input.peek(Token![,]) {
+                    input.parse::<Token![,]>()?;
+                }
+            } else {
+                // This looks like the start of format arguments
+                break;
+            }
+        }
+
+        // Parse remaining tokens as format arguments
+        if !input.is_empty() {
+            let remaining: proc_macro2::TokenStream = input.parse()?;
+            format_args = Some(remaining);
+        }
+
+        Ok(InstrumentArgs {
+            method_calls,
+            format_args,
+        })
+    }
+}
+
+/// Instruments an async function with await-tree spans.
+///
+/// This attribute macro transforms an async function to automatically create
+/// an await-tree span and instrument the function's execution.
+///
+/// # Usage
+///
+/// ```rust,ignore
+/// #[await_tree::instrument("span_name({})", arg1)]
+/// async fn foo(arg1: i32, arg2: String) {
+///     // function body
+/// }
+/// ```
+///
+/// With keywords:
+///
+/// ```rust,ignore
+/// #[await_tree::instrument(long_running, verbose, "span_name({})", arg1)]
+/// async fn foo(arg1: i32, arg2: String) {
+///     // function body
+/// }
+/// ```
+///
+/// The above will be expanded to:
+///
+/// ```rust,ignore
+/// async fn foo(arg1: i32, arg2: String) {
+///     let span = await_tree::span!("span_name({})", arg1).long_running().verbose();
+///     let fut = async move {
+///         // original function body
+///     };
+///     fut.instrument_await(span).await
+/// }
+/// ```
+///
+/// # Arguments
+///
+/// The macro accepts format arguments similar to `format!` or `println!`:
+/// - The first argument is the format string
+/// - Subsequent arguments are the values to be formatted
+///
+/// The format arguments are passed directly to the `await_tree::span!` macro
+/// without any parsing or modification.
+#[proc_macro_attribute]
+pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream {
+    let input_fn = parse_macro_input!(input as ItemFn);
+
+    // Validate that this is an async function
+    if input_fn.sig.asyncness.is_none() {
+        return syn::Error::new_spanned(
+            &input_fn.sig.fn_token,
+            "the `instrument` attribute can only be applied to async functions",
+        )
+        .to_compile_error()
+        .into();
+    }
+
+    // Parse the arguments
+    let parsed_args = if args.is_empty() {
+        InstrumentArgs {
+            method_calls: Vec::new(),
+            format_args: None,
+        }
+    } else {
+        match syn::parse::<InstrumentArgs>(args) {
+            Ok(args) => args,
+            Err(e) => return e.to_compile_error().into(),
+        }
+    };
+
+    // Extract the span format arguments
+    let span_args = if let Some(format_args) = parsed_args.format_args {
+        quote! { #format_args }
+    } else {
+        // If no format arguments provided, use the function name as span
+        let fn_name = &input_fn.sig.ident;
+        quote! { stringify!(#fn_name) }
+    };
+
+    // Build span creation with method calls
+    let mut span_creation = quote! { ::await_tree::span!(#span_args) };
+
+    // Chain all method calls
+    for method_name in parsed_args.method_calls {
+        span_creation = quote! { #span_creation.#method_name() };
+    }
+
+    // Extract function components
+    let fn_vis = &input_fn.vis;
+    let fn_sig = &input_fn.sig;
+    let fn_block = &input_fn.block;
+    let fn_attrs = &input_fn.attrs;
+
+    // Generate the instrumented function
+    let result = quote! {
+        #(#fn_attrs)*
+        #fn_vis #fn_sig {
+            use ::await_tree::SpanExt as _;
+            let __at_span: ::await_tree::Span = #span_creation;
+            let __at_fut = async move #fn_block;
+            ::await_tree::InstrumentAwait::instrument_await(__at_fut, __at_span).await
+        }
+    };
+
+    result.into()
+}