Repository: risingwavelabs/await-tree Branch: main Commit: cea540a0d710 Files: 38 Total size: 111.1 KB Directory structure: gitextract_cf93mxty/ ├── .github/ │ └── workflows/ │ ├── check-test.yaml │ └── publish.yml ├── .gitignore ├── .vscode/ │ └── settings.json ├── CHANGELOG.md ├── CODEOWNERS ├── Cargo.toml ├── LICENSE ├── README.md ├── await-tree-attributes/ │ ├── Cargo.toml │ ├── README.md │ ├── src/ │ │ └── lib.rs │ └── tests/ │ ├── expansion.rs │ └── integration.rs ├── benches/ │ └── basic.rs ├── examples/ │ ├── basic.rs │ ├── detach.rs │ ├── global.rs │ ├── instrument.rs │ ├── long_running.rs │ ├── multiple.rs │ ├── serde.rs │ ├── spawn.rs │ └── verbose.rs ├── rust-toolchain.toml ├── rustfmt.toml └── src/ ├── context.rs ├── future.rs ├── global.rs ├── lib.rs ├── obj_utils.rs ├── registry.rs ├── root.rs ├── span.rs ├── spawn.rs ├── tests/ │ ├── functionality.rs │ └── spawn.rs └── tests.rs ================================================ FILE CONTENTS ================================================ ================================================ FILE: .github/workflows/check-test.yaml ================================================ name: Check and Test on: push: branches: ["main"] pull_request: branches: ["main"] env: CARGO_TERM_COLOR: always jobs: cargo-check-test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - name: Format run: cargo fmt --check # Check with all features - name: Build (all features) run: cargo build --all-targets --all-features - name: Clippy (all features) run: cargo clippy --all-targets --all-features # Check without any features - name: Build (no features) run: cargo build --all-targets - name: Clippy (no features) run: cargo clippy --all-targets # Run tests with all features - name: Run tests run: cargo test --all-features # Run examples - name: Run examples run: | for example in $(ls examples/ | sed 's/\.rs$//'); do echo "Running example $example with all features" cargo run --example $example --all-features done semver: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - uses: cargo-bins/cargo-binstall@main - name: Install cargo-semver-checks run: cargo binstall -y cargo-semver-checks - name: Check run: cargo semver-checks check-release -p await-tree ================================================ FILE: .github/workflows/publish.yml ================================================ name: Publish to crates.io on: push: tags: - "v*" env: CARGO_TERM_COLOR: always jobs: # Publish to crates.io after checks pass publish: name: Publish to crates.io runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 - name: Publish to crates.io run: | cargo publish -p await-tree-attributes cargo publish -p await-tree env: CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} ================================================ FILE: .gitignore ================================================ /target /Cargo.lock ================================================ FILE: .vscode/settings.json ================================================ { "rust-analyzer.cargo.features": "all" } ================================================ FILE: CHANGELOG.md ================================================ # Changelog 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("", 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 - Added `span!` macro as a replacement for `format!` with better performance ([#21](https://github.com/risingwavelabs/await-tree/pull/21)) - Added support for global registry and updated its documentation ([#17](https://github.com/risingwavelabs/await-tree/pull/17), [#18](https://github.com/risingwavelabs/await-tree/pull/18)) - Implemented `serde::Serialize` for tree to provide structured output and made it an optional feature ([#22](https://github.com/risingwavelabs/await-tree/pull/22), [#24](https://github.com/risingwavelabs/await-tree/pull/24)) - Added attributes `long_running` and `verbose` on span, removed `verbose_instrument_await` ([#20](https://github.com/risingwavelabs/await-tree/pull/20)) ### Changed - Only depend on `tokio` if spawn capability is desired ([#25](https://github.com/risingwavelabs/await-tree/pull/25)) - Added workflow for publishing crates ([#23](https://github.com/risingwavelabs/await-tree/pull/23)) ### Fixed - Fixed examples and added CI for running examples ([#19](https://github.com/risingwavelabs/await-tree/pull/19)) [0.3.0]: https://github.com/risingwavelabs/await-tree/compare/v0.2.1...v0.3.0 ================================================ FILE: CODEOWNERS ================================================ * @BugenZhao ================================================ FILE: Cargo.toml ================================================ [workspace] members = [".", "await-tree-attributes"] [package] name = "await-tree" version = "0.3.2-alpha.2" edition = "2021" description = "Generate accurate and informative tree dumps of asynchronous tasks." repository = "https://github.com/risingwavelabs/await-tree" keywords = ["async", "tokio", "backtrace", "actor"] categories = ["development-tools::debugging"] license = "Apache-2.0" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [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.0-alpha.2", optional = true } coarsetime = "0.1" derive_builder = "0.20" easy-ext = "1" flexstr = "0.9" indextree = "4" itertools = "0.12" parking_lot = "0.12" pin-project = "1" serde = { version = "1", features = ["derive"], optional = true } task-local = "0.1" tokio = { version = "1", features = ["rt"], optional = true } tracing = "0.1" weak-table = "0.3.2" [dev-dependencies] criterion = { version = "0.5", features = ["async", "async_tokio"] } futures = { version = "0.3", default-features = false, features = ["alloc"] } serde_json = "1" tokio = { version = "1", features = ["full"] } [[bench]] name = "basic" harness = false [profile.bench] opt-level = 3 debug = false codegen-units = 1 lto = 'fat' incremental = false debug-assertions = false overflow-checks = false rpath = false [[example]] name = "serde" required-features = ["serde"] [[example]] name = "spawn" required-features = ["tokio"] [[example]] name = "global" required-features = ["tokio"] [[example]] name = "instrument" required-features = ["tokio", "attributes"] ================================================ FILE: LICENSE ================================================ Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] 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. ================================================ FILE: README.md ================================================ # await-tree [![Crate](https://img.shields.io/crates/v/await-tree.svg)](https://crates.io/crates/await-tree) [![Docs](https://docs.rs/await-tree/badge.svg)](https://docs.rs/await-tree) The `Future`s in Async Rust can be arbitrarily composited or nested to achieve a variety of control flows. Assuming that the execution of each `Future` is represented as a node, then the asynchronous execution of an async task can be organized into a **logical tree**, which is constantly transformed over the polling, completion, and cancellation of `Future`s. `await-tree` allows developers to dump this execution tree at runtime, with the span of each `Future` annotated by `instrument_await`. A basic example is shown below, and more examples of complex control flows can be found in the [examples](./examples) directory. ```rust async fn bar(i: i32) { // static span baz(i).instrument_await("baz in bar").await } async fn baz(i: i32) { // runtime `String` span is also supported work().instrument_await(span!("working in baz {i}")).await } async fn foo() { // spans of joined futures will be siblings in the tree join( // attribute the span with `long_running` or `verbose` bar(3).instrument_await("bar".long_running()), baz(2).instrument_await("baz".verbose()), ) .await; } // Init the global registry to start tracing the tasks. await_tree::init_global_registry(Default::default()); // Spawn a task with root span "foo" and key "foo". // Note: The `spawn` function requires the `tokio` feature to be enabled. await_tree::spawn("foo", "foo", foo()); // Let the tasks run for a while. sleep(Duration::from_secs(1)).await; // Get the tree of the task with key "foo". let tree = Registry::current().get("foo").unwrap(); // foo [1.006s] // bar [1.006s] // baz in bar [1.006s] // working in baz 3 [1.006s] // baz [1.006s] // working in baz 2 [1.006s] println!("{tree}"); ``` ## Features `await-tree` provides the following optional features: - `serde`: Enables serialization of the tree structure using serde. This allows you to serialize the tree to formats like JSON, as shown in the [serde example](./examples/serde.rs). ```rust // Enable the serde feature in Cargo.toml // await-tree = { version = "", features = ["serde"] } // Then you can serialize the tree let tree = Registry::current().get("foo").unwrap(); let json = serde_json::to_string_pretty(&tree).unwrap(); println!("{json}"); ``` - `tokio`: Enables integration with the Tokio runtime, providing task spawning capabilities through `spawn` and `spawn_anonymous` functions. This feature is required for the examples that demonstrate spawning tasks. ```rust // Enable the tokio feature in Cargo.toml // await-tree = { version = "", features = ["tokio"] } // Then you can spawn tasks with await-tree instrumentation await_tree::spawn("task-key", "root_span", async { // Your async code here work().instrument_await("work_span").await; }); ``` - `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 = "", 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`: **Pros of `await-tree`**: - `await-tree` support customizing the span with runtime `String`, while `async-backtrace` only supports function name and line number. This is useful when we want to annotate the span with some dynamic information, such as the identifier of a shared resource (e.g., a lock), to see how the contention happens among different tasks. - `await-tree` support almost all kinds of async control flows with arbitrary `Future` topology, while `async-backtrace` fails to handle some of them. For example, it's common to use `&mut impl Future` as an arm of `select` to avoid problems led by cancellation unsafety. To further resolve this `Future` after the `select` completes, we may move it to another place and `await` it there. `async-backtrace` fails to track this `Future` again due to the change of its parent. See [`examples/detach.rs`](./examples/detach.rs) for more details. - `await-tree` maintains the tree structure with an [arena-based data structure](https://crates.io/crates/indextree), with zero extra `unsafe` code. For comparison, `async-backtrace` crafts it by hand and there's potential memory unsafety for unhandled topologies mentioned above. It's worth pointing out that `await-tree` has been applied in the production deployment of [RisingWave](https://github.com/risingwavelabs/risingwave), a distributed streaming database, for a long time. - `await-tree` maintains the tree structure separately from the `Future` itself, which enables developers to dump the tree at any time with nearly no contention, no matter the `Future` is under active polling or has been pending. For comparison, `async-backtrace` has to [wait](https://docs.rs/async-backtrace/0.2.5/async_backtrace/fn.taskdump_tree.html) for the polling to complete before dumping the tree, which may cause a long delay. **Pros of `async-backtrace`**: - `async-backtrace` is under the Tokio organization. ## License `await-tree` is distributed under the Apache License (Version 2.0). Please refer to [LICENSE](./LICENSE) for more information. ================================================ FILE: await-tree-attributes/Cargo.toml ================================================ [package] name = "await-tree-attributes" version = "0.1.0-alpha.2" 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 = ".." } ================================================ FILE: await-tree-attributes/README.md ================================================ # 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. ================================================ FILE: await-tree-attributes/src/lib.rs ================================================ // 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 #[derive(Default)] struct InstrumentArgs { method_calls: Vec, format_args: Option, boxed: bool, } impl syn::parse::Parse for InstrumentArgs { fn parse(input: syn::parse::ParseStream) -> syn::Result { let mut method_calls = Vec::new(); let mut format_args = None; let mut boxed = false; // Parse identifiers first (these will become method calls or special keywords) 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 or special keyword input.parse::()?; // consume the identifier // Check for special "boxed" keyword if ident == "boxed" { boxed = true; } else { method_calls.push(ident); } if input.peek(Token![,]) { input.parse::()?; } } 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, boxed, }) } } /// 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 attributes on the span: /// /// ```rust,ignore /// #[await_tree::instrument(long_running, verbose, "span_name({})", arg1)] /// async fn foo(arg1: i32, arg2: String) { /// // function body /// } /// ``` /// /// With the `boxed` keyword to `Box::pin` the function body before calling `instrument_await`, /// which can help reducing the stack usage if you encounter stack overflow: /// /// ```rust,ignore /// #[await_tree::instrument(boxed, "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 /// }; /// let fut = Box::pin(fut); // if `boxed` is specified /// 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::default() } else { match syn::parse::(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 boxed = (parsed_args.boxed).then(|| quote! { let __at_fut = ::std::boxed::Box::pin(__at_fut); }); 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; #boxed ::await_tree::InstrumentAwait::instrument_await(__at_fut, __at_span).await } }; result.into() } ================================================ FILE: await-tree-attributes/tests/expansion.rs ================================================ // 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. //! Test to verify the attribute expansion works correctly. use await_tree_attributes::instrument; // Test that the macro generates the expected code structure #[instrument("test_expansion({})", value)] async fn test_expansion(value: i32) -> i32 { value * 2 } #[tokio::test] async fn test_attribute_expansion() { // This test verifies that the attribute expansion compiles and runs correctly let result = test_expansion(21).await; assert_eq!(result, 42); } // Test with no arguments #[instrument] async fn no_args_function() -> String { "success".to_string() } #[tokio::test] async fn test_no_args_expansion() { let result = no_args_function().await; assert_eq!(result, "success"); } // Test with long_running keyword #[instrument(long_running, "long_running_task({})", id)] async fn long_running_task(id: u32) -> u32 { id * 10 } // Test with verbose keyword #[instrument(verbose, "verbose_task")] async fn verbose_task() -> String { "verbose".to_string() } // Test with both keywords #[instrument(long_running, verbose, "complex_task({}, {})", name, value)] async fn complex_task(name: &str, value: i32) -> String { format!("{}: {}", name, value) } // Test with keywords but no format args #[instrument(long_running, verbose)] async fn keywords_only_task() -> i32 { 42 } #[tokio::test] async fn test_keywords() { let result = long_running_task(5).await; assert_eq!(result, 50); let result = verbose_task().await; assert_eq!(result, "verbose"); let result = complex_task("test", 123).await; assert_eq!(result, "test: 123"); let result = keywords_only_task().await; assert_eq!(result, 42); } // Test with boxed keyword #[instrument(boxed, "boxed_task({})", value)] async fn boxed_task(value: i32) -> i32 { value * 3 } // Test with boxed and other keywords #[instrument(boxed, long_running, "boxed_long_running_task")] async fn boxed_long_running_task() -> String { "boxed and long running".to_owned() } // Test with boxed but no format args #[instrument(boxed)] async fn boxed_no_args_task() -> i32 { 100 } #[tokio::test] async fn test_boxed_keyword() { let result = boxed_task(7).await; assert_eq!(result, 21); let result = boxed_long_running_task().await; assert_eq!(result, "boxed and long running"); let result = boxed_no_args_task().await; assert_eq!(result, 100); } // Note: The attribute now accepts any identifiers as method names. // If the methods don't exist on Span, it will fail at compile time, which is the desired behavior. // For example, this would fail to compile: // #[instrument(custom_method, another_method, "arbitrary_methods")] // async fn arbitrary_methods_task() -> String { "this compiles".to_string() } ================================================ FILE: await-tree-attributes/tests/integration.rs ================================================ // 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. use await_tree_attributes::instrument; // Test basic usage with format string and arguments #[instrument("test_function({})", arg1)] async fn test_function(arg1: i32, arg2: String) -> i32 { tokio::time::sleep(std::time::Duration::from_millis(10)).await; arg1 + arg2.len() as i32 } // Test with no arguments (should use function name) #[instrument] async fn simple_function() -> String { tokio::time::sleep(std::time::Duration::from_millis(5)).await; "hello".to_string() } // Test with complex format string #[instrument("complex_operation({}, {})", name, value)] async fn complex_function(name: &str, value: u64) -> String { tokio::time::sleep(std::time::Duration::from_millis(1)).await; format!("{}: {}", name, value) } #[tokio::test] async fn test_instrument_attribute() { // These tests mainly verify that the attribute compiles correctly // and the functions can be called normally let result = test_function(42, "test".to_string()).await; assert_eq!(result, 46); let result = simple_function().await; assert_eq!(result, "hello"); let result = complex_function("test", 123).await; assert_eq!(result, "test: 123"); } // Test that the macro preserves function visibility and attributes #[instrument("public_fn")] pub async fn public_function() -> i32 { 42 } #[allow(dead_code)] #[instrument("private_fn")] async fn private_function() -> i32 { 24 } #[tokio::test] async fn test_visibility_and_attributes() { let result = public_function().await; assert_eq!(result, 42); let result = private_function().await; assert_eq!(result, 24); } ================================================ FILE: benches/basic.rs ================================================ // Copyright 2023 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. use std::time::Duration; use await_tree::{Config, ConfigBuilder, InstrumentAwait, Registry}; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use tokio::runtime::{Builder, Runtime}; use tokio::task::yield_now; fn runtime() -> Runtime { Builder::new_current_thread().enable_time().build().unwrap() } async fn test() { async fn test_inner() { futures::future::join( async { yield_now().await; black_box(1) } .instrument_await("fut1"), async { yield_now().await; yield_now().await; black_box(2) } .instrument_await("fut2"), ) .instrument_await("join") .await; } for _ in 0..10000 { test_inner().await; } } async fn test_baseline() { async fn test_inner() { futures::future::join( async { yield_now().await; black_box(1) }, async { yield_now().await; yield_now().await; black_box(2) }, ) .await; } for _ in 0..10000 { test_inner().await; } } async fn spawn_many(size: usize) { let registry = Registry::new(Config::default()); let mut handles = vec![]; for i in 0..size { let task = async { tokio::time::sleep(Duration::from_millis(10)).await; }; handles.push(tokio::spawn( registry.register(i, "new_task").instrument(task), )); } futures::future::try_join_all(handles) .await .expect("failed to join background task"); } async fn spawn_many_baseline(size: usize) { let mut handles = vec![]; for _ in 0..size { let task = async { tokio::time::sleep(Duration::from_millis(10)).await; }; handles.push(tokio::spawn(task)); } futures::future::try_join_all(handles) .await .expect("failed to join background task"); } // time: [6.5488 ms 6.5541 ms 6.5597 ms] // change: [+6.5978% +6.7838% +6.9299%] (p = 0.00 < 0.05) // Performance has regressed. fn bench_basic(c: &mut Criterion) { c.bench_function("basic", |b| { b.to_async(runtime()).iter(|| async { let config = ConfigBuilder::default().verbose(false).build().unwrap(); let registry = Registry::new(config); let root = registry.register(233, "root"); root.instrument(test()).await; }) }); } fn bench_basic_baseline(c: &mut Criterion) { c.bench_function("basic_baseline", |b| { b.to_async(runtime()).iter(|| async { let config = ConfigBuilder::default().verbose(false).build().unwrap(); let registry = Registry::new(config); let root = registry.register(233, "root"); black_box(root); test_baseline().await }) }); } criterion_group!(benches, bench_basic, bench_basic_baseline); // with_register_to_root time: [15.993 ms 16.122 ms 16.292 ms] // baseline time: [13.940 ms 13.961 ms 13.982 ms] fn bench_many_baseline(c: &mut Criterion) { c.bench_function("with_register_to_root_baseline", |b| { b.to_async(runtime()) .iter(|| async { black_box(spawn_many_baseline(10000)).await }) }); } fn bench_many_exp(c: &mut Criterion) { c.bench_function("with_register_to_root", |b| { b.to_async(runtime()) .iter(|| async { black_box(spawn_many(10000)).await }) }); } criterion_group!( name = bench_many; config = Criterion::default().sample_size(50); targets = bench_many_exp, bench_many_baseline ); criterion_main!(benches, bench_many); ================================================ FILE: examples/basic.rs ================================================ // Copyright 2023 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. //! This example shows the basic usage of `await-tree`. use std::time::Duration; use await_tree::{span, Config, InstrumentAwait, Registry}; use futures::future::{join, pending}; use tokio::time::sleep; async fn bar(i: i32) { // `&'static str` span baz(i).instrument_await("baz in bar").await } async fn baz(i: i32) { // runtime `String` span is also supported pending() .instrument_await(span!("pending in baz {i}")) .await } async fn foo() { // spans of joined futures will be siblings in the tree join( bar(3).instrument_await("bar"), baz(2).instrument_await("baz"), ) .await; } #[tokio::main] async fn main() { let registry = Registry::new(Config::default()); let root = registry.register((), "foo"); tokio::spawn(root.instrument(foo())); sleep(Duration::from_secs(1)).await; let tree = registry.get(()).unwrap().to_string(); // foo [1.006s] // bar [1.006s] // baz in bar [1.006s] // pending in baz 3 [1.006s] // baz [1.006s] // pending in baz 2 [1.006s] println!("{tree}"); } ================================================ FILE: examples/detach.rs ================================================ // Copyright 2023 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. //! This example shows how a span can be detached from and remounted to the tree. use std::time::Duration; use await_tree::{Config, InstrumentAwait, Registry}; use futures::channel::oneshot::{self, Receiver}; use futures::future::{pending, select}; use futures::FutureExt; use tokio::time::sleep; async fn work(rx: Receiver<()>) { let mut fut = pending().instrument_await("fut"); // poll `fut` under the `select` span let _ = select( sleep(Duration::from_millis(500)) .instrument_await("sleep") .boxed(), &mut fut, ) .instrument_await("select") .await; // `select` span closed so `fut` is detached // the elapsed time of `fut` should be preserved // wait for the signal to continue rx.instrument_await("rx").await.unwrap(); // poll `fut` under the root `work` span, and it'll be remounted fut.await } #[tokio::main] async fn main() { let registry = Registry::new(Config::default()); let root = registry.register((), "work"); let (tx, rx) = oneshot::channel(); tokio::spawn(root.instrument(work(rx))); sleep(Duration::from_millis(100)).await; let tree = registry.get(()).unwrap().to_string(); // work [106.290ms] // select [106.093ms] // sleep [106.093ms] // fut [106.093ms] println!("{tree}"); sleep(Duration::from_secs(1)).await; let tree = registry.get(()).unwrap().to_string(); // work [1.112s] // rx [606.944ms] // [Detached 4] // fut [1.112s] println!("{tree}"); tx.send(()).unwrap(); sleep(Duration::from_secs(1)).await; let tree = registry.get(()).unwrap().to_string(); // work [2.117s] // fut [2.117s] println!("{tree}"); } ================================================ FILE: examples/global.rs ================================================ // Copyright 2023 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. //! This example shows the usage of the global registry. //! //! Note: This example requires the `tokio` feature to be enabled. //! Run with: `cargo run --example global --features tokio` #![cfg(feature = "tokio")] use std::time::Duration; use await_tree::{init_global_registry, Config, InstrumentAwait, Registry}; use futures::future::pending; async fn bar() { pending::<()>().instrument_await("pending").await; } async fn foo() { await_tree::spawn_anonymous("spawn bar", bar()); bar().instrument_await("bar").await; } async fn print() { tokio::time::sleep(Duration::from_secs(1)).await; // Access the registry anywhere and collect all trees. for (key, tree) in Registry::current().collect_all() { // [Actor 42] // foo [1.003s] // bar [1.003s] // pending [1.003s] // // [Anonymous #2] // spawn bar [1.003s] // pending [1.003s] // // [Print] // print [1.003s] println!("[{}]\n{}\n", key, tree); } } #[tokio::main] async fn main() { init_global_registry(Config::default()); // After global registry is initialized, the tasks can be spawned everywhere, being // registered in the global registry. await_tree::spawn("Actor 42", "foo", foo()); // The line above is a shorthand for the following: tokio::spawn( Registry::current() .register("Print", "print") .instrument(print()), ) .await .unwrap(); } ================================================ FILE: examples/instrument.rs ================================================ // 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. //! Example demonstrating the use of the `#[await_tree::instrument]` attribute macro. use await_tree::{init_global_registry, instrument, spawn_derived_root, ConfigBuilder, Registry}; use std::time::Duration; use tokio::time::sleep; #[instrument(long_running, "fetch_data({})", id)] async fn fetch_data(id: u32) -> String { sleep(Duration::from_millis(100)).await; format!("data_{}", id) } #[instrument(verbose, "process_item({}, {})", name, value)] async fn process_item(name: &str, value: i32) -> i32 { sleep(Duration::from_millis(50)).await; value * 2 } #[instrument(long_running, verbose, "complex_operation")] async fn complex_operation() -> Vec { let mut results = Vec::new(); for i in 1..=3 { let data = fetch_data(i).await; results.push(data); } let processed = process_item("test", 42).await; results.push(format!("processed: {}", processed)); results } #[instrument] async fn simple_task() -> String { sleep(Duration::from_millis(100)).await; "simple result".to_string() } #[tokio::main] async fn main() { // Initialize the global registry init_global_registry(ConfigBuilder::default().verbose(true).build().unwrap()); // Spawn tasks with instrumentation spawn_derived_root("complex", complex_operation()); spawn_derived_root("simple", simple_task()); // Let the tasks run for a while sleep(Duration::from_millis(50)).await; // Print the await trees if let Some(tree) = Registry::current().get("complex") { println!("Complex task tree:"); println!("{}", tree); println!(); } if let Some(tree) = Registry::current().get("simple") { println!("Simple task tree:"); println!("{}", tree); } } ================================================ FILE: examples/long_running.rs ================================================ // 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. //! This example shows how to mark a span as "long_running", so that it will //! not be marked as "!!!" if it takes too long to complete. use std::time::Duration; use await_tree::{Config, InstrumentAwait, Registry, SpanExt}; use futures::future::{pending, select}; use futures::FutureExt; use tokio::time::sleep; async fn long_running_child() { pending() .instrument_await("long_running_child".long_running()) .await } async fn child() { pending().instrument_await("child").await } async fn foo() { select(long_running_child().boxed(), child().boxed()).await; } async fn work() -> String { let registry = Registry::new(Config::default()); let root = registry.register((), "foo"); tokio::spawn(root.instrument(foo())); // The default threshold is 10 seconds. sleep(Duration::from_secs(11)).await; registry.get(()).unwrap().to_string() } #[tokio::main] async fn main() { // foo [11.006s] // long_running_child [11.006s] // child [!!! 11.006s] println!("{}", work().await); } ================================================ FILE: examples/multiple.rs ================================================ // Copyright 2023 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. //! This example shows how to use `await-tree` for multiple actors. use std::time::Duration; use await_tree::{span, Config, InstrumentAwait, Registry}; use futures::future::pending; use itertools::Itertools; use tokio::time::sleep; async fn work(i: i32) { foo().instrument_await(span!("actor work {i}")).await } async fn foo() { pending().instrument_await("pending").await } #[tokio::main] async fn main() { let registry = Registry::new(Config::default()); for i in 0_i32..3 { let root = registry.register(i, format!("actor {i}")); tokio::spawn(root.instrument(work(i))); } sleep(Duration::from_secs(1)).await; // actor 0 [1.007s] // actor work 0 [1.007s] // pending [1.007s] // // actor 1 [1.007s] // actor work 1 [1.007s] // pending [1.007s] // // actor 2 [1.007s] // actor work 2 [1.007s] // pending [1.007s] for (_, tree) in registry .collect::() .into_iter() .sorted_by_key(|(i, _)| *i) { println!("{tree}"); } } ================================================ FILE: examples/serde.rs ================================================ // 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. //! This example shows the serialization format of the tree. //! //! The execution flow is the same as `examples/detach.rs`. //! //! Note: This example requires the `serde` feature to be enabled. //! Run with: `cargo run --example serde --features serde` #![cfg(feature = "serde")] use std::time::Duration; use await_tree::{Config, InstrumentAwait, Registry}; use futures::channel::oneshot::{self, Receiver}; use futures::future::{pending, select}; use futures::FutureExt; use tokio::time::sleep; async fn work(rx: Receiver<()>) { let mut fut = pending().instrument_await("fut"); // poll `fut` under the `select` span let _ = select( sleep(Duration::from_millis(500)) .instrument_await("sleep") .boxed(), &mut fut, ) .instrument_await("select") .await; // `select` span closed so `fut` is detached // the elapsed time of `fut` should be preserved // wait for the signal to continue rx.instrument_await("rx").await.unwrap(); // poll `fut` under the root `work` span, and it'll be remounted fut.await } #[tokio::main] async fn main() { let registry = Registry::new(Config::default()); let root = registry.register((), "work"); let (tx, rx) = oneshot::channel(); tokio::spawn(root.instrument(work(rx))); sleep(Duration::from_millis(100)).await; let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); // { // "current": 1, // "tree": { // "id": 1, // "span": { // "name": "work", // "is_verbose": false, // "is_long_running": true // }, // "elapsed_ns": 105404875, // "children": [ // { // "id": 2, // "span": { // "name": "select", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 105287624, // "children": [ // { // "id": 3, // "span": { // "name": "sleep", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 105267874, // "children": [] // }, // { // "id": 4, // "span": { // "name": "fut", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 105264874, // "children": [] // } // ] // } // ] // }, // "detached": [] // } println!("{tree}"); sleep(Duration::from_secs(1)).await; let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); // { // "current": 1, // "tree": { // "id": 1, // "span": { // "name": "work", // "is_verbose": false, // "is_long_running": true // }, // "elapsed_ns": 1108552791, // "children": [ // { // "id": 3, // "span": { // "name": "rx", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 603081749, // "children": [] // } // ] // }, // "detached": [ // { // "id": 4, // "span": { // "name": "fut", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 1108412791, // "children": [] // } // ] // } println!("{tree}"); tx.send(()).unwrap(); sleep(Duration::from_secs(1)).await; let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); // { // "current": 1, // "tree": { // "id": 1, // "span": { // "name": "work", // "is_verbose": false, // "is_long_running": true // }, // "elapsed_ns": 2114497458, // "children": [ // { // "id": 4, // "span": { // "name": "fut", // "is_verbose": false, // "is_long_running": false // }, // "elapsed_ns": 2114366458, // "children": [] // } // ] // }, // "detached": [] // } println!("{tree}"); } ================================================ FILE: examples/spawn.rs ================================================ // Copyright 2023 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. //! This example shows how to spawn tasks with `await_tree::spawn` that are automatically registered //! to the current registry of the scope. //! //! Note: This example requires the `tokio` feature to be enabled. //! Run with: `cargo run --example spawn --features tokio` #![cfg(feature = "tokio")] use std::time::Duration; use await_tree::{Config, InstrumentAwait, Registry}; use futures::future::pending; use tokio::time::sleep; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] struct Actor(usize); async fn actor(i: usize) { // Since we're already inside the scope of a registered/instrumented task, we can directly spawn // new tasks with `await_tree::spawn` to also register them in the same registry. await_tree::spawn_anonymous(format!("background task {i}"), async { pending::<()>().await; }) .instrument_await("waiting for background task") .await .unwrap(); } #[tokio::main] async fn main() { let registry = Registry::new(Config::default()); for i in 0..3 { let root = registry.register(Actor(i), format!("actor {i}")); tokio::spawn(root.instrument(actor(i))); } sleep(Duration::from_secs(1)).await; for (_actor, tree) in registry.collect::() { // actor 0 [1.004s] // waiting for background task [1.004s] println!("{tree}"); } for tree in registry.collect_anonymous() { // background task 0 [1.004s] println!("{tree}"); } } ================================================ FILE: examples/verbose.rs ================================================ // Copyright 2023 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. //! This example shows how to mark a span as "verbose", so that it's conditionally //! enabled based on the config. use std::time::Duration; use await_tree::{ConfigBuilder, InstrumentAwait, Registry, SpanExt}; use futures::future::pending; use tokio::time::sleep; async fn foo() { // verbose span will be disabled if the `verbose` flag in the config is false pending().instrument_await("pending".verbose()).await } async fn work(verbose: bool) -> String { let config = ConfigBuilder::default().verbose(verbose).build().unwrap(); let registry = Registry::new(config); let root = registry.register((), "foo"); tokio::spawn(root.instrument(foo())); sleep(Duration::from_secs(1)).await; registry.get(()).unwrap().to_string() } #[tokio::main] async fn main() { // foo [1.001s] println!("{}", work(false).await); // foo [1.004s] // pending [1.004s] println!("{}", work(true).await); } ================================================ FILE: rust-toolchain.toml ================================================ [toolchain] channel = "1.84" profile = "default" ================================================ FILE: rustfmt.toml ================================================ comment_width = 120 format_code_in_doc_comments = true format_macro_bodies = true format_macro_matchers = true normalize_comments = true normalize_doc_attributes = true imports_granularity = "Module" group_imports = "StdExternalCrate" reorder_impl_items = true reorder_imports = true tab_spaces = 4 wrap_comments = true ================================================ FILE: src/context.rs ================================================ // Copyright 2023 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. use std::fmt::{Debug, Write}; use std::sync::atomic::{AtomicU64, Ordering}; use indextree::{Arena, NodeId}; use itertools::Itertools; use parking_lot::{Mutex, MutexGuard}; use crate::root::current_context; use crate::Span; /// Node in the span tree. #[derive(Debug, Clone)] struct SpanNode { /// The span value. span: Span, /// The time when this span was started, or the future was first polled. start_time: coarsetime::Instant, } impl SpanNode { /// Create a new node with the given value. fn new(span: Span) -> Self { Self { span, start_time: coarsetime::Instant::now(), } } } /// The id of an await-tree context. /// /// We will check the id recorded in the instrumented future against the current task-local context /// before trying to update the tree. /// /// Also used as the key for anonymous trees in the registry. Intentionally made private to prevent /// users from reusing the same id when registering a new tree. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub(crate) struct ContextId(pub(crate) u64); /// An await-tree for a task. #[derive(Debug, Clone)] pub struct Tree { /// The arena for allocating span nodes in this context. arena: Arena, /// The root span node. root: NodeId, /// The current span node. This is the node that is currently being polled. current: NodeId, } #[cfg(feature = "serde")] mod serde_impl { use serde::ser::SerializeStruct as _; use serde::Serialize; use super::*; struct SpanNodeSer<'a> { arena: &'a Arena, node: NodeId, } impl Serialize for SpanNodeSer<'_> { fn serialize(&self, serializer: S) -> Result where S: serde::Serializer, { let inner = self.arena[self.node].get(); let mut s = serializer.serialize_struct("Span", 4)?; // Basic info. let id: usize = self.node.into(); s.serialize_field("id", &id)?; s.serialize_field("span", &inner.span)?; s.serialize_field("elapsed_ns", &inner.start_time.elapsed().as_nanos())?; // Children. let children = (self.node.children(self.arena)) .map(|node| SpanNodeSer { arena: self.arena, node, }) .sorted_by_key(|child| { let inner = self.arena[child.node].get(); inner.start_time }) .collect_vec(); s.serialize_field("children", &children)?; s.end() } } impl Serialize for Tree { fn serialize(&self, serializer: S) -> Result where S: serde::Serializer, { let mut s = serializer.serialize_struct("Tree", 3)?; let current_id: usize = self.current.into(); s.serialize_field("current", ¤t_id)?; // The main tree. s.serialize_field( "tree", &SpanNodeSer { arena: &self.arena, node: self.root, }, )?; // The detached subtrees. let detached = self .detached_roots() .map(|node| SpanNodeSer { arena: &self.arena, node, }) .collect_vec(); s.serialize_field("detached", &detached)?; s.end() } } } impl std::fmt::Display for Tree { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { fn fmt_node( f: &mut std::fmt::Formatter<'_>, arena: &Arena, node: NodeId, depth: usize, current: NodeId, ) -> std::fmt::Result { f.write_str(&" ".repeat(depth * 2))?; let inner = arena[node].get(); f.write_str(&inner.span.name)?; let elapsed: std::time::Duration = inner.start_time.elapsed().into(); write!( f, " [{}{:.3?}]", if !inner.span.is_long_running && elapsed.as_secs() >= 10 { "!!! " } else { "" }, elapsed )?; if depth > 0 && node == current { f.write_str(" <== current")?; } f.write_char('\n')?; for child in node .children(arena) .sorted_by_key(|&id| arena[id].get().start_time) { fmt_node(f, arena, child, depth + 1, current)?; } Ok(()) } fmt_node(f, &self.arena, self.root, 0, self.current)?; // Format all detached spans. for node in self.detached_roots() { writeln!(f, "[Detached {node}]")?; fmt_node(f, &self.arena, node, 1, self.current)?; } Ok(()) } } impl Tree { /// Get the count of active span nodes in this context. #[cfg(test)] pub(crate) fn active_node_count(&self) -> usize { self.arena.iter().filter(|n| !n.is_removed()).count() } /// Get the count of active detached span nodes in this context. #[cfg(test)] pub(crate) fn detached_node_count(&self) -> usize { self.arena .iter() .filter(|n| { !n.is_removed() && n.parent().is_none() && self.arena.get_node_id(n).unwrap() != self.root }) .count() } /// Returns an iterator over the root nodes of detached subtrees. fn detached_roots(&self) -> impl Iterator + '_ { self.arena .iter() .filter(|n| !n.is_removed()) // still valid .map(|node| self.arena.get_node_id(node).unwrap()) // get id .filter(|&id| id != self.root && self.arena[id].parent().is_none()) // no parent but non-root } /// Push a new span as a child of current span, used for future firstly polled. /// /// Returns the new current span. pub(crate) fn push(&mut self, span: Span) -> NodeId { let child = self.arena.new_node(SpanNode::new(span)); self.current.prepend(child, &mut self.arena); self.current = child; child } /// Step in the current span to the given child, used for future polled again. /// /// If the child is not actually a child of the current span, it means we are using a new future /// to poll it, so we need to detach it from the previous parent, and attach it to the current /// span. pub(crate) fn step_in(&mut self, child: NodeId) { if !self.current.children(&self.arena).contains(&child) { // Actually we can always call this even if `child` is already a child of `current`. But // checking first performs better. self.current.prepend(child, &mut self.arena); } self.current = child; } /// Pop the current span to the parent, used for future ready. /// /// Note that there might still be some children of this node, like `select_stream.next()`. /// The children might be polled again later, and will be attached as the children of a new /// span. pub(crate) fn pop(&mut self) { let parent = self.arena[self.current] .parent() .expect("the root node should not be popped"); self.remove_and_detach(self.current); self.current = parent; } /// Step out the current span to the parent, used for future pending. pub(crate) fn step_out(&mut self) { let parent = self.arena[self.current] .parent() .expect("the root node should not be stepped out"); self.current = parent; } /// Remove the current span and detach the children, used for future aborting. /// /// The children might be polled again later, and will be attached as the children of a new /// span. pub(crate) fn remove_and_detach(&mut self, node: NodeId) { node.detach(&mut self.arena); // Removing detached `node` makes children detached. node.remove(&mut self.arena); } /// Get the current span node id. pub(crate) fn current(&self) -> NodeId { self.current } } /// The task-local await-tree context. #[derive(Debug)] pub(crate) struct TreeContext { /// The id of the context. id: ContextId, /// Whether to include the "verbose" span in the tree. verbose: bool, /// The await-tree. tree: Mutex, } impl TreeContext { /// Create a new context. pub(crate) fn new(root_span: Span, verbose: bool) -> Self { static ID: AtomicU64 = AtomicU64::new(0); let id = ID.fetch_add(1, Ordering::Relaxed); // Always make the root span long-running. let root_span = root_span.long_running(); let mut arena = Arena::new(); let root = arena.new_node(SpanNode::new(root_span)); Self { id: ContextId(id), verbose, tree: Tree { arena, root, current: root, } .into(), } } /// Get the context id. pub(crate) fn id(&self) -> ContextId { self.id } /// Returns the locked guard of the tree. pub(crate) fn tree(&self) -> MutexGuard<'_, Tree> { self.tree.lock() } /// Whether the verbose span should be included. pub(crate) fn verbose(&self) -> bool { self.verbose } } /// Get the await-tree of current task. Returns `None` if we're not instrumented. /// /// This is useful if you want to check which component or runtime task is calling this function. pub fn current_tree() -> Option { current_context().map(|c| c.tree().clone()) } ================================================ FILE: src/future.rs ================================================ // Copyright 2023 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. use std::future::Future; use std::pin::Pin; use std::task::Poll; use indextree::NodeId; use pin_project::{pin_project, pinned_drop}; use crate::context::ContextId; use crate::root::current_context; use crate::Span; enum State { Initial(Span), Polled { this_node: NodeId, this_context_id: ContextId, }, Ready, /// This span is disabled due to `verbose` configuration. Disabled, } /// The future for [`InstrumentAwait`][ia]. /// /// [ia]: crate::InstrumentAwait #[pin_project(PinnedDrop)] pub struct Instrumented { #[pin] inner: F, state: State, } impl Instrumented { pub(crate) fn new(inner: F, span: Span) -> Self { Self { inner, state: State::Initial(span), } } } impl Future for Instrumented { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { let this = self.project(); let context = current_context(); let (context, this_node) = match this.state { State::Initial(span) => { match context { Some(c) => { if !c.verbose() && span.is_verbose { // The tracing for this span is disabled according to the verbose // configuration. *this.state = State::Disabled; return this.inner.poll(cx); } // First polled, push a new span to the context. let node = c.tree().push(std::mem::take(span)); *this.state = State::Polled { this_node: node, this_context_id: c.id(), }; (c, node) } // Not in a context None => return this.inner.poll(cx), } } State::Polled { this_node, this_context_id: this_context, } => { match context { // Context correct Some(c) if c.id() == *this_context => { // Polled before, just step in. c.tree().step_in(*this_node); (c, *this_node) } // Context changed Some(_) => { tracing::warn!( "future polled in a different context as it was first polled" ); return this.inner.poll(cx); } // Out of context None => { tracing::warn!( "future polled not in a context, while it was when first polled" ); return this.inner.poll(cx); } } } State::Ready => unreachable!("the instrumented future should always be fused"), State::Disabled => return this.inner.poll(cx), }; // The current node must be the this_node. debug_assert_eq!(this_node, context.tree().current()); match this.inner.poll(cx) { // The future is ready, clean-up this span by popping from the context. Poll::Ready(output) => { context.tree().pop(); *this.state = State::Ready; Poll::Ready(output) } // Still pending, just step out. Poll::Pending => { context.tree().step_out(); Poll::Pending } } } } #[pinned_drop] impl PinnedDrop for Instrumented { fn drop(self: Pin<&mut Self>) { let this = self.project(); match this.state { State::Polled { this_node, this_context_id, } => match current_context() { // Context correct Some(c) if c.id() == *this_context_id => { c.tree().remove_and_detach(*this_node); } // Context changed Some(_) => { tracing::warn!("future is dropped in a different context as it was first polled, cannot clean up!"); } // Out of context None => { tracing::warn!("future is not in a context, while it was when first polled, cannot clean up!"); } }, State::Initial(_) | State::Ready | State::Disabled => {} } } } ================================================ FILE: src/global.rs ================================================ // Copyright 2023 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. use std::sync::OnceLock; use crate::{Config, Registry}; static GLOBAL_REGISTRY: OnceLock = OnceLock::new(); /// Initialize the global registry with the given configuration. /// Panics if the global registry has already been initialized. /// /// This is **optional** and only needed if you want to use the global registry. /// You can always create a new registry with [`Registry::new`] and pass it around to achieve /// better encapsulation. pub fn init_global_registry(config: Config) { if let Err(_r) = GLOBAL_REGISTRY.set(Registry::new(config)) { panic!("global registry already initialized") } } pub(crate) fn global_registry() -> Option { GLOBAL_REGISTRY.get().cloned() } ================================================ FILE: src/lib.rs ================================================ // Copyright 2023 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. //! Generate accurate and informative tree dumps of asynchronous tasks. //! //! # Example //! //! Below is a basic example of how to trace asynchronous tasks with the global registry of the //! `await-tree` crate. //! //! ```rust //! # use std::time::Duration; //! # use tokio::time::sleep; //! # use await_tree::{InstrumentAwait, Registry, span}; //! # use futures::future::join; //! # //! # async fn work() { futures::future::pending::<()>().await } //! # //! async fn bar(i: i32) { //! // `&'static str` span //! baz(i).instrument_await("baz in bar").await //! } //! //! async fn baz(i: i32) { //! // runtime `String` span is also supported //! work().instrument_await(span!("working in baz {i}")).await //! } //! //! async fn foo() { //! // spans of joined futures will be siblings in the tree //! join( //! bar(3).instrument_await("bar"), //! baz(2).instrument_await("baz"), //! ) //! .await; //! } //! //! # #[tokio::main] //! # async fn main() { //! // Init the global registry to start tracing the tasks. //! await_tree::init_global_registry(Default::default()); //! // Spawn a task with root span "foo" and key "foo". //! await_tree::spawn("foo", "foo", foo()); //! // Let the tasks run for a while. //! sleep(Duration::from_secs(1)).await; //! // Get the tree of the task with key "foo". //! let tree = Registry::current().get("foo").unwrap(); //! //! // foo [1.006s] //! // bar [1.006s] //! // baz in bar [1.006s] //! // working in baz 3 [1.006s] //! // baz [1.006s] //! // working in baz 2 [1.006s] //! println!("{tree}"); //! # } //! ``` #![forbid(missing_docs)] use std::future::Future; mod context; mod future; mod global; mod obj_utils; mod registry; mod root; mod span; #[cfg(feature = "tokio")] mod spawn; pub use context::{current_tree, Tree}; pub use future::Instrumented; pub use global::init_global_registry; pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry, ToRootSpan}; pub use root::TreeRoot; pub use span::{Span, SpanExt}; #[cfg(feature = "tokio")] pub use spawn::{spawn, spawn_anonymous, spawn_derived_root}; #[cfg(feature = "attributes")] pub use await_tree_attributes::instrument; #[doc(hidden)] pub mod __private { pub use crate::span::fmt_span; } /// Attach spans to a future to be traced in the await-tree. pub trait InstrumentAwait: Future + Sized { /// Instrument the future with a span. fn instrument_await(self, span: impl Into) -> Instrumented { Instrumented::new(self, span.into()) } } impl InstrumentAwait for F where F: Future {} #[cfg(test)] mod tests; ================================================ FILE: src/obj_utils.rs ================================================ // Copyright 2023 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. //! Utilities for using `Any` as the key of a `HashMap`. // Adopted from: // https://github.com/bevyengine/bevy/blob/56bcbb097552b45e3ff48c48947ed8ee4e2c24b1/crates/bevy_utils/src/label.rs use std::any::Any; use std::hash::{Hash, Hasher}; /// An object safe version of [`Eq`]. This trait is automatically implemented /// for any `'static` type that implements `Eq`. pub(crate) trait DynEq: Any { /// Casts the type to `dyn Any`. fn as_any(&self) -> &dyn Any; /// This method tests for `self` and `other` values to be equal. /// /// Implementers should avoid returning `true` when the underlying types are /// not the same. fn dyn_eq(&self, other: &dyn DynEq) -> bool; } impl DynEq for T where T: Any + Eq, { fn as_any(&self) -> &dyn Any { self } fn dyn_eq(&self, other: &dyn DynEq) -> bool { if let Some(other) = other.as_any().downcast_ref::() { return self == other; } false } } /// An object safe version of [`Hash`]. This trait is automatically implemented /// for any `'static` type that implements `Hash`. pub(crate) trait DynHash: DynEq { /// Casts the type to `dyn Any`. fn as_dyn_eq(&self) -> &dyn DynEq; /// Feeds this value into the given [`Hasher`]. fn dyn_hash(&self, state: &mut dyn Hasher); } impl DynHash for T where T: DynEq + Hash, { fn as_dyn_eq(&self) -> &dyn DynEq { self } fn dyn_hash(&self, mut state: &mut dyn Hasher) { T::hash(self, &mut state); self.type_id().hash(&mut state); } } ================================================ FILE: src/registry.rs ================================================ // Copyright 2023 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. use std::any::Any; use std::fmt::{Debug, Display}; use std::hash::Hash; use std::sync::{Arc, Weak}; use derive_builder::Builder; use parking_lot::RwLock; use weak_table::WeakValueHashMap; use crate::context::{ContextId, Tree, TreeContext}; use crate::obj_utils::{DynEq, DynHash}; use crate::{span, Span, TreeRoot}; /// Configuration for an await-tree registry, which affects the behavior of all await-trees in the /// registry. #[derive(Debug, Clone, Builder)] #[builder(default)] pub struct Config { /// Whether to include the **verbose** span in the await-tree. verbose: bool, } #[allow(clippy::derivable_impls)] impl Default for Config { fn default() -> Self { Self { verbose: false } } } /// A key that can be used to identify a task and its await-tree in the [`Registry`]. /// /// All thread-safe types that can be used as a key of a hash map are automatically implemented with /// this trait. pub trait Key: Hash + Eq + Debug + Send + Sync + 'static {} impl Key for T where T: Hash + Eq + Debug + Send + Sync + 'static {} /// The object-safe version of [`Key`], automatically implemented. trait ObjKey: DynHash + DynEq + Debug + Send + Sync + 'static {} impl ObjKey for T where T: DynHash + DynEq + Debug + Send + Sync + 'static {} /// A trait for types that can be converted to a [`Span`] that can be used as the root of an /// await-tree. pub trait ToRootSpan { /// Convert the type to a [`Span`] that can be used as the root of an await-tree. fn to_root_span(&self) -> Span; } impl ToRootSpan for T { fn to_root_span(&self) -> Span { span!("{self}") } } /// Key type for anonymous await-trees. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] struct AnonymousKey(ContextId); impl Display for AnonymousKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "Anonymous #{}", self.0 .0) } } /// Type-erased key for the [`Registry`]. #[derive(Clone)] pub struct AnyKey(Arc); impl PartialEq for AnyKey { fn eq(&self, other: &Self) -> bool { self.0.dyn_eq(other.0.as_dyn_eq()) } } impl Eq for AnyKey {} impl Hash for AnyKey { fn hash(&self, state: &mut H) { self.0.dyn_hash(state); } } impl Debug for AnyKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { self.0.fmt(f) } } impl Display for AnyKey { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { // TODO: for all `impl Display`? macro_rules! delegate_to_display { ($($t:ty),* $(,)?) => { $( if let Some(k) = self.as_any().downcast_ref::<$t>() { return write!(f, "{}", k); } )* }; } delegate_to_display!(String, &str, AnonymousKey); write!(f, "{:?}", self) } } impl AnyKey { fn new(key: impl ObjKey) -> Self { Self(Arc::new(key)) } /// Cast the key to `dyn Any`. pub fn as_any(&self) -> &dyn Any { self.0.as_ref().as_any() } /// Returns whether the key is of type `K`. /// /// Equivalent to `self.as_any().is::()`. pub fn is(&self) -> bool { self.as_any().is::() } /// Returns whether the key corresponds to an anonymous await-tree. pub fn is_anonymous(&self) -> bool { self.as_any().is::() } /// Returns the key as a reference to type `K`, if it is of type `K`. /// /// Equivalent to `self.as_any().downcast_ref::()`. pub fn downcast_ref(&self) -> Option<&K> { self.as_any().downcast_ref() } } type Contexts = RwLock>>; struct RegistryCore { contexts: Contexts, config: Config, } /// The registry of multiple await-trees. /// /// Can be cheaply cloned to share the same registry. pub struct Registry(Arc); impl Debug for Registry { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Registry") .field("config", self.config()) .finish_non_exhaustive() } } impl Clone for Registry { fn clone(&self) -> Self { Self(Arc::clone(&self.0)) } } impl Registry { fn contexts(&self) -> &Contexts { &self.0.contexts } fn config(&self) -> &Config { &self.0.config } } impl Registry { /// Create a new registry with given `config`. pub fn new(config: Config) -> Self { Self( RegistryCore { contexts: Default::default(), config, } .into(), ) } /// Returns the current registry, if exists. /// /// 1. If the current task is registered with a registry, returns the registry. /// 2. If the global registry is initialized with /// [`init_global_registry`](crate::global::init_global_registry), returns the global /// registry. /// 3. Otherwise, returns `None`. pub fn try_current() -> Option { crate::root::current_registry() } /// Returns the current registry, panics if not exists. /// /// See [`Registry::try_current`] for more information. pub fn current() -> Self { Self::try_current().expect("no current registry") } fn register_inner(&self, key: impl Key, context: Arc) -> TreeRoot { self.contexts() .write() .insert(AnyKey::new(key), Arc::clone(&context)); TreeRoot { context, registry: WeakRegistry(Arc::downgrade(&self.0)), } } /// Register with given key and root span. Returns a [`TreeRoot`] that can be used to instrument /// a future. /// /// If the key already exists, a new [`TreeRoot`] is returned and the reference to the old /// [`TreeRoot`] is dropped. pub fn register(&self, key: impl Key, root_span: impl Into) -> TreeRoot { let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose)); self.register_inner(key, context) } /// Derive the root span from the given key and register with it. /// /// This is a convenience method for `self.register(key, key.to_root_span())`. See /// [`Registry::register`] for more details. pub fn register_derived_root(&self, key: impl Key + ToRootSpan) -> TreeRoot { let root_span = key.to_root_span(); self.register(key, root_span) } /// Register an anonymous await-tree without specifying a key. Returns a [`TreeRoot`] that can /// be used to instrument a future. /// /// Anonymous await-trees are not able to be retrieved through the [`Registry::get`] method. Use /// [`Registry::collect_anonymous`] or [`Registry::collect_all`] to collect them. // TODO: we have keyed and anonymous, should we also have a typed-anonymous (for classification // only)? pub fn register_anonymous(&self, root_span: impl Into) -> TreeRoot { let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose)); self.register_inner(AnonymousKey(context.id()), context) // use the private id as the key } /// Get a clone of the await-tree with given key. /// /// Returns `None` if the key does not exist or the tree root has been dropped. pub fn get(&self, key: impl Key) -> Option { self.contexts() .read() .get(&AnyKey::new(key)) // TODO: accept ref can? .map(|v| v.tree().clone()) } /// Remove all the registered await-trees. pub fn clear(&self) { self.contexts().write().clear(); } /// Collect the snapshots of all await-trees with the key of type `K`. pub fn collect(&self) -> Vec<(K, Tree)> { self.contexts() .read() .iter() .filter_map(|(k, v)| { k.0.as_ref() .as_any() .downcast_ref::() .map(|k| (k.clone(), v.tree().clone())) }) .collect() } /// Collect the snapshots of all await-trees registered with [`Registry::register_anonymous`]. pub fn collect_anonymous(&self) -> Vec { self.contexts() .read() .iter() .filter_map(|(k, v)| { if k.is_anonymous() { Some(v.tree().clone()) } else { None } }) .collect() } /// Collect the snapshots of all await-trees regardless of the key type. pub fn collect_all(&self) -> Vec<(AnyKey, Tree)> { self.contexts() .read() .iter() .map(|(k, v)| (k.clone(), v.tree().clone())) .collect() } } pub(crate) struct WeakRegistry(Weak); impl WeakRegistry { pub fn upgrade(&self) -> Option { self.0.upgrade().map(Registry) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_registry() { let registry = Registry::new(Config::default()); let _0_i32 = registry.register(0_i32, "0"); let _1_i32 = registry.register(1_i32, "1"); let _2_i32 = registry.register(2_i32, "2"); let _0_str = registry.register("0", "0"); let _1_str = registry.register("1", "1"); let _unit = registry.register((), "()"); let _unit_replaced = registry.register((), "[]"); let _anon = registry.register_anonymous("anon"); let _anon = registry.register_anonymous("anon"); let i32s = registry.collect::(); assert_eq!(i32s.len(), 3); let strs = registry.collect::<&'static str>(); assert_eq!(strs.len(), 2); let units = registry.collect::<()>(); assert_eq!(units.len(), 1); let anons = registry.collect_anonymous(); assert_eq!(anons.len(), 2); let all = registry.collect_all(); assert_eq!(all.len(), 8); } } ================================================ FILE: src/root.rs ================================================ // Copyright 2023 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. use std::future::Future; use std::sync::Arc; use crate::context::TreeContext; use crate::global::global_registry; use crate::registry::WeakRegistry; use crate::Registry; /// The root of an await-tree. pub struct TreeRoot { pub(crate) context: Arc, pub(crate) registry: WeakRegistry, } task_local::task_local! { static ROOT: TreeRoot } pub(crate) fn current_context() -> Option> { ROOT.try_with(|r| r.context.clone()).ok() } pub(crate) fn current_registry() -> Option { let local = || ROOT.try_with(|r| r.registry.upgrade()).ok().flatten(); let global = global_registry; local().or_else(global) } impl TreeRoot { /// Instrument the given future with the context of this tree root. pub async fn instrument(self, future: F) -> F::Output { ROOT.scope(self, future).await } } ================================================ FILE: src/span.rs ================================================ // 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. type SpanName = flexstr::SharedStr; #[doc(hidden)] pub fn fmt_span(args: std::fmt::Arguments<'_>) -> Span { let name = if let Some(str) = args.as_str() { SpanName::from_ref(str) } else { flexstr::flex_fmt(args) }; Span::new(name) } /// Creates a new span with formatted name. /// /// [`instrument_await`] accepts any type that implements [`AsRef`] as the span name. /// This macro provides similar functionality to [`format!`], but with improved performance /// by creating the span name on the stack when possible, avoiding unnecessary allocations. /// /// [`instrument_await`]: crate::InstrumentAwait::instrument_await #[macro_export] // XXX: Without this extra binding (`let res = ..`), it will make the future `!Send`. // This is also how `std::format!` behaves. But why? macro_rules! span { ($($fmt_arg:tt)*) => {{ let res = $crate::__private::fmt_span(format_args!($($fmt_arg)*)); res }}; } /// A cheaply cloneable span in the await-tree. #[cfg_attr(feature = "serde", derive(serde::Serialize))] #[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)] pub struct Span { pub(crate) name: SpanName, pub(crate) is_verbose: bool, pub(crate) is_long_running: bool, } impl Span { fn new(name: SpanName) -> Self { Self { name, is_verbose: false, is_long_running: false, } } } impl Span { /// Set the verbose attribute of the span. /// /// When a span is marked as verbose, it will be included in the output /// only if the `verbose` flag in the [`Config`] is set. /// /// [`Config`]: crate::Config pub fn verbose(mut self) -> Self { self.is_verbose = true; self } /// Set the long-running attribute of the span. /// /// When a span is marked as long-running, it will not be marked as "!!!" /// in the formatted [`Tree`] if it takes too long to complete. The root /// span is always marked as long-running. /// /// [`Tree`]: crate::Tree pub fn long_running(mut self) -> Self { self.is_long_running = true; self } } /// Convert a value into a span and set attributes. #[easy_ext::ext(SpanExt)] impl> T { /// Convert `self` into a span and set the verbose attribute. /// /// See [`Span::verbose`] for more details. pub fn verbose(self) -> Span { self.into().verbose() } /// Convert `self` into a span and set the long-running attribute. /// /// See [`Span::long_running`] for more details. pub fn long_running(self) -> Span { self.into().long_running() } } impl> From for Span { fn from(value: S) -> Self { Self::new(SpanName::from_ref(value)) } } impl std::fmt::Display for Span { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { self.name.fmt(f) } } ================================================ FILE: src/spawn.rs ================================================ // Copyright 2023 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. // TODO: should we consider exposing `current_registry` // so that users can not only spawn tasks but also get and collect trees? use std::future::Future; use tokio::task::JoinHandle; use crate::{Key, Registry, Span, ToRootSpan}; impl Registry { /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. /// /// The spawned task will be registered in the this registry. pub fn spawn( &self, key: impl Key, root_span: impl Into, future: T, ) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { tokio::spawn(self.register(key, root_span).instrument(future)) } /// Spawns a new asynchronous task instrumented with the root span derived from the given key, /// returning a [`JoinHandle`] for it. /// /// The spawned task will be registered in the this registry. /// /// This is a convenience method for `self.spawn(key, key.to_root_span(), future)`. pub fn spawn_derived_root( &self, key: impl Key + ToRootSpan, future: T, ) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { let root_span = key.to_root_span(); self.spawn(key, root_span, future) } /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. /// /// The spawned task will be registered in the this registry anonymously. pub fn spawn_anonymous(&self, root_span: impl Into, future: T) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { tokio::spawn(self.register_anonymous(root_span).instrument(future)) } } /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. /// /// The spawned task will be registered in the current [`Registry`](crate::Registry) returned by /// [`Registry::try_current`] with the given [`Key`], if it exists. Otherwise, this is equivalent to /// [`tokio::spawn`]. pub fn spawn(key: impl Key, root_span: impl Into, future: T) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { if let Some(registry) = Registry::try_current() { registry.spawn(key, root_span, future) } else { tokio::spawn(future) } } /// Spawns a new asynchronous task instrumented with the root span derived from the given key, /// returning a [`JoinHandle`] for it. /// /// The spawned task will be registered in the current [`Registry`](crate::Registry) returned by /// [`Registry::try_current`] with the given [`Key`], if it exists. Otherwise, this is equivalent to /// [`tokio::spawn`]. /// /// This is a convenience function for `spawn(key, key.to_root_span(), future)`. See [`spawn`] for /// more details. pub fn spawn_derived_root(key: impl Key + ToRootSpan, future: T) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { let root_span = key.to_root_span(); spawn(key, root_span, future) } /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a /// [`JoinHandle`] for it. /// /// The spawned task will be registered in the current [`Registry`](crate::Registry) returned by /// [`Registry::try_current`] anonymously, if it exists. Otherwise, this is equivalent to /// [`tokio::spawn`]. pub fn spawn_anonymous(root_span: impl Into, future: T) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, { if let Some(registry) = Registry::try_current() { registry.spawn_anonymous(root_span, future) } else { tokio::spawn(future) } } ================================================ FILE: src/tests/functionality.rs ================================================ // Copyright 2023 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. use futures::future::{join_all, poll_fn, select_all}; use futures::{pin_mut, FutureExt, Stream, StreamExt}; use itertools::Itertools; use crate::root::current_context; use crate::{span, Config, InstrumentAwait, Registry}; async fn sleep(time: u64) { tokio::time::sleep(std::time::Duration::from_millis(time)).await; println!("slept {time}ms"); } async fn sleep_nested() { join_all([ sleep(1500).instrument_await("sleep nested 1500"), sleep(2500).instrument_await("sleep nested 2500"), ]) .await; } async fn multi_sleep() { sleep(400).await; sleep(800) .instrument_await("sleep another in multi sleep") .await; } fn stream1() -> impl Stream { use futures::stream::{iter, once}; iter(std::iter::repeat_with(|| { once(async { sleep(150).await; }) })) .flatten() } fn stream2() -> impl Stream { use futures::stream::{iter, once}; iter([ once(async { sleep(444).await; }) .boxed(), once(async { join_all([ sleep(400).instrument_await("sleep nested 400"), sleep(600).instrument_await("sleep nested 600"), ]) .await; }) .boxed(), ]) .flatten() } async fn hello() { async move { // Join join_all([ sleep(1000) .boxed() .instrument_await(span!("sleep {}", 1000)), sleep(2000).boxed().instrument_await("sleep 2000"), sleep_nested().boxed().instrument_await("sleep nested"), multi_sleep().boxed().instrument_await("multi sleep"), ]) .await; // Join another join_all([ sleep(1200).instrument_await("sleep 1200"), sleep(2200).instrument_await("sleep 2200"), ]) .await; // Cancel select_all([ sleep(666).boxed().instrument_await("sleep 666"), sleep_nested() .boxed() .instrument_await("sleep nested (should be cancelled)"), ]) .await; // Check whether cleaned up sleep(233).instrument_await("sleep 233").await; // Check stream next drop { let mut stream1 = stream1().fuse().boxed(); let mut stream2 = stream2().fuse().boxed(); let mut count = 0; 'outer: loop { tokio::select! { _ = stream1.next().instrument_await(span!("stream1 next {count}")) => {}, r = stream2.next().instrument_await(span!("stream2 next {count}")) => { if r.is_none() { break 'outer } }, } sleep(50) .instrument_await(span!("sleep before next stream poll: {count}")) .await; count += 1; } } // Check whether cleaned up sleep(233).instrument_await("sleep 233").await; // TODO: add tests on sending the future to another task or context. } .instrument_await("hello") .await; // Aborted futures have been cleaned up. There should only be a single active node of root. assert_eq!(current_context().unwrap().tree().active_node_count(), 1); } #[tokio::test] async fn test_await_tree() { let registry = Registry::new(Config::default()); let root = registry.register((), "actor 233"); let fut = root.instrument(hello()); pin_mut!(fut); let expected_counts = vec![ (1, 0), (8, 0), (9, 0), (8, 0), (6, 0), (5, 0), (4, 0), (4, 0), (3, 0), (6, 0), (3, 0), (4, 0), (3, 0), (4, 0), (3, 0), (4, 0), (3, 0), (6, 0), (5, 2), (6, 0), (5, 2), (6, 0), (5, 0), (4, 1), (5, 0), (3, 0), (3, 0), ]; let mut actual_counts = vec![]; poll_fn(|cx| { let tree = registry .collect::<()>() .into_iter() .exactly_one() .ok() .unwrap() .1; println!("{tree}"); actual_counts.push((tree.active_node_count(), tree.detached_node_count())); fut.poll_unpin(cx) }) .await; assert_eq!(actual_counts, expected_counts); } ================================================ FILE: src/tests/spawn.rs ================================================ // Copyright 2023 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. #![cfg(feature = "tokio")] use std::time::Duration; use futures::future::pending; use tokio::time::sleep; use crate::{Config, InstrumentAwait, Registry}; #[tokio::test] async fn main() { let registry = Registry::new(Config::default()); tokio::spawn(registry.register((), "root").instrument(async { crate::spawn_anonymous("child", async { crate::spawn_anonymous("grandson", async { pending::<()>().await; }) .instrument_await("wait for grandson") .await .unwrap() }) .instrument_await("wait for child") .await .unwrap() })); sleep(Duration::from_secs(1)).await; assert_eq!(registry.collect::<()>().len(), 1); assert_eq!(registry.collect_anonymous().len(), 2); assert_eq!(registry.collect_all().len(), 3); } ================================================ FILE: src/tests.rs ================================================ // Copyright 2023 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. mod functionality; mod spawn;