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