├── .github └── workflows │ ├── check-test.yaml │ └── publish.yml ├── .gitignore ├── .vscode └── settings.json ├── CHANGELOG.md ├── CODEOWNERS ├── Cargo.toml ├── LICENSE ├── README.md ├── benches └── basic.rs ├── examples ├── basic.rs ├── detach.rs ├── global.rs ├── long_running.rs ├── multiple.rs ├── serde.rs ├── spawn.rs └── verbose.rs ├── rust-toolchain ├── rustfmt.toml └── src ├── context.rs ├── future.rs ├── global.rs ├── lib.rs ├── obj_utils.rs ├── registry.rs ├── root.rs ├── span.rs ├── spawn.rs ├── tests.rs └── tests ├── functionality.rs └── spawn.rs /.github/workflows/check-test.yaml: -------------------------------------------------------------------------------- 1 | name: Check and Test 2 | 3 | on: 4 | push: 5 | branches: ["main"] 6 | pull_request: 7 | branches: ["main"] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | cargo-check-test: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/checkout@v3 17 | - name: Format 18 | run: cargo fmt --check 19 | 20 | # Check with all features 21 | - name: Build (all features) 22 | run: cargo build --all-targets --all-features 23 | - name: Clippy (all features) 24 | run: cargo clippy --all-targets --all-features 25 | 26 | # Check without any features 27 | - name: Build (no features) 28 | run: cargo build --all-targets 29 | - name: Clippy (no features) 30 | run: cargo clippy --all-targets 31 | 32 | # Run tests with all features 33 | - name: Run tests 34 | run: cargo test --all-features 35 | 36 | # Run examples 37 | - name: Run examples 38 | run: | 39 | for example in $(ls examples/ | sed 's/\.rs$//'); do 40 | echo "Running example $example with all features" 41 | cargo run --example $example --all-features 42 | done 43 | 44 | semver: 45 | runs-on: ubuntu-latest 46 | steps: 47 | - uses: actions/checkout@v3 48 | - uses: cargo-bins/cargo-binstall@main 49 | - name: Install cargo-semver-checks 50 | run: cargo binstall -y cargo-semver-checks 51 | - name: Check 52 | run: cargo semver-checks check-release -p await-tree 53 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Publish to crates.io 2 | 3 | on: 4 | push: 5 | tags: 6 | - "v*" 7 | 8 | env: 9 | CARGO_TERM_COLOR: always 10 | 11 | jobs: 12 | # Run the check-test workflow first 13 | check-test: 14 | uses: ./.github/workflows/check-test.yaml 15 | 16 | # Publish to crates.io after checks pass 17 | publish: 18 | name: Publish to crates.io 19 | needs: check-test 20 | runs-on: ubuntu-latest 21 | steps: 22 | - uses: actions/checkout@v3 23 | 24 | - name: Publish to crates.io 25 | run: cargo publish 26 | env: 27 | CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} 28 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | /Cargo.lock 3 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "rust-analyzer.cargo.features": "all" 3 | } 4 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 6 | and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 7 | 8 | ## [0.3.0] - 2025-04-09 9 | 10 | ### Added 11 | - Added `span!` macro as a replacement for `format!` with better performance ([#21](https://github.com/risingwavelabs/await-tree/pull/21)) 12 | - 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)) 13 | - 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)) 14 | - Added attributes `long_running` and `verbose` on span, removed `verbose_instrument_await` ([#20](https://github.com/risingwavelabs/await-tree/pull/20)) 15 | 16 | ### Changed 17 | - Only depend on `tokio` if spawn capability is desired ([#25](https://github.com/risingwavelabs/await-tree/pull/25)) 18 | - Added workflow for publishing crates ([#23](https://github.com/risingwavelabs/await-tree/pull/23)) 19 | 20 | ### Fixed 21 | - Fixed examples and added CI for running examples ([#19](https://github.com/risingwavelabs/await-tree/pull/19)) 22 | 23 | [0.3.0]: https://github.com/risingwavelabs/await-tree/compare/v0.2.1...v0.3.0 24 | -------------------------------------------------------------------------------- /CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @BugenZhao 2 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "await-tree" 3 | version = "0.3.0" 4 | edition = "2021" 5 | description = "Generate accurate and informative tree dumps of asynchronous tasks." 6 | repository = "https://github.com/risingwavelabs/await-tree" 7 | keywords = ["async", "tokio", "backtrace", "actor"] 8 | categories = ["development-tools::debugging"] 9 | license = "Apache-2.0" 10 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 11 | 12 | [features] 13 | serde = ["dep:serde", "flexstr/serde"] 14 | tokio = ["dep:tokio"] 15 | 16 | [dependencies] 17 | coarsetime = "0.1" 18 | derive_builder = "0.20" 19 | easy-ext = "1" 20 | flexstr = "0.9" 21 | indextree = "4" 22 | itertools = "0.12" 23 | parking_lot = "0.12" 24 | pin-project = "1" 25 | serde = { version = "1", features = ["derive"], optional = true } 26 | task-local = "0.1" 27 | tokio = { version = "1", features = ["rt"], optional = true } 28 | tracing = "0.1" 29 | weak-table = "0.3.2" 30 | 31 | [dev-dependencies] 32 | criterion = { version = "0.5", features = ["async", "async_tokio"] } 33 | futures = { version = "0.3", default-features = false, features = ["alloc"] } 34 | serde_json = "1" 35 | tokio = { version = "1", features = ["full"] } 36 | 37 | [[bench]] 38 | name = "basic" 39 | harness = false 40 | 41 | [profile.bench] 42 | opt-level = 3 43 | debug = false 44 | codegen-units = 1 45 | lto = 'fat' 46 | incremental = false 47 | debug-assertions = false 48 | overflow-checks = false 49 | rpath = false 50 | 51 | [[example]] 52 | name = "serde" 53 | required-features = ["serde"] 54 | 55 | [[example]] 56 | name = "spawn" 57 | required-features = ["tokio"] 58 | 59 | [[example]] 60 | name = "global" 61 | required-features = ["tokio"] 62 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # await-tree 2 | 3 | [![Crate](https://img.shields.io/crates/v/await-tree.svg)](https://crates.io/crates/await-tree) 4 | [![Docs](https://docs.rs/await-tree/badge.svg)](https://docs.rs/await-tree) 5 | 6 | The `Future`s in Async Rust can be arbitrarily composited or nested to achieve a variety of control flows. 7 | Assuming that the execution of each `Future` is represented as a node, 8 | then the asynchronous execution of an async task can be organized into a **logical tree**, 9 | which is constantly transformed over the polling, completion, and cancellation of `Future`s. 10 | 11 | `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. 12 | 13 | ```rust 14 | async fn bar(i: i32) { 15 | // static span 16 | baz(i).instrument_await("baz in bar").await 17 | } 18 | 19 | async fn baz(i: i32) { 20 | // runtime `String` span is also supported 21 | work().instrument_await(span!("working in baz {i}")).await 22 | } 23 | 24 | async fn foo() { 25 | // spans of joined futures will be siblings in the tree 26 | join( 27 | // attribute the span with `long_running` or `verbose` 28 | bar(3).instrument_await("bar".long_running()), 29 | baz(2).instrument_await("baz".verbose()), 30 | ) 31 | .await; 32 | } 33 | 34 | // Init the global registry to start tracing the tasks. 35 | await_tree::init_global_registry(Default::default()); 36 | // Spawn a task with root span "foo" and key "foo". 37 | // Note: The `spawn` function requires the `tokio` feature to be enabled. 38 | await_tree::spawn("foo", "foo", foo()); 39 | // Let the tasks run for a while. 40 | sleep(Duration::from_secs(1)).await; 41 | // Get the tree of the task with key "foo". 42 | let tree = Registry::current().get("foo").unwrap(); 43 | 44 | // foo [1.006s] 45 | // bar [1.006s] 46 | // baz in bar [1.006s] 47 | // working in baz 3 [1.006s] 48 | // baz [1.006s] 49 | // working in baz 2 [1.006s] 50 | println!("{tree}"); 51 | ``` 52 | 53 | ## Features 54 | 55 | `await-tree` provides the following optional features: 56 | 57 | - `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). 58 | 59 | ```rust 60 | // Enable the serde feature in Cargo.toml 61 | // await-tree = { version = "", features = ["serde"] } 62 | 63 | // Then you can serialize the tree 64 | let tree = Registry::current().get("foo").unwrap(); 65 | let json = serde_json::to_string_pretty(&tree).unwrap(); 66 | println!("{json}"); 67 | ``` 68 | 69 | - `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. 70 | 71 | ```rust 72 | // Enable the tokio feature in Cargo.toml 73 | // await-tree = { version = "", features = ["tokio"] } 74 | 75 | // Then you can spawn tasks with await-tree instrumentation 76 | await_tree::spawn("task-key", "root_span", async { 77 | // Your async code here 78 | work().instrument_await("work_span").await; 79 | }); 80 | ``` 81 | 82 | ## Compared to `async-backtrace` 83 | 84 | [`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`: 85 | 86 | **Pros of `await-tree`**: 87 | 88 | - `await-tree` support customizing the span with runtime `String`, while `async-backtrace` only supports function name and line number. 89 | 90 | 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. 91 | 92 | - `await-tree` support almost all kinds of async control flows with arbitrary `Future` topology, while `async-backtrace` fails to handle some of them. 93 | 94 | 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. 95 | 96 | - `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. 97 | 98 | 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. 99 | 100 | - `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. 101 | 102 | **Pros of `async-backtrace`**: 103 | 104 | - `async-backtrace` is under the Tokio organization. 105 | 106 | ## License 107 | 108 | `await-tree` is distributed under the Apache License (Version 2.0). Please refer to [LICENSE](./LICENSE) for more information. 109 | -------------------------------------------------------------------------------- /benches/basic.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::time::Duration; 16 | 17 | use await_tree::{Config, ConfigBuilder, InstrumentAwait, Registry}; 18 | use criterion::{black_box, criterion_group, criterion_main, Criterion}; 19 | use tokio::runtime::{Builder, Runtime}; 20 | use tokio::task::yield_now; 21 | 22 | fn runtime() -> Runtime { 23 | Builder::new_current_thread().enable_time().build().unwrap() 24 | } 25 | 26 | async fn test() { 27 | async fn test_inner() { 28 | futures::future::join( 29 | async { 30 | yield_now().await; 31 | black_box(1) 32 | } 33 | .instrument_await("fut1"), 34 | async { 35 | yield_now().await; 36 | yield_now().await; 37 | black_box(2) 38 | } 39 | .instrument_await("fut2"), 40 | ) 41 | .instrument_await("join") 42 | .await; 43 | } 44 | 45 | for _ in 0..10000 { 46 | test_inner().await; 47 | } 48 | } 49 | 50 | async fn test_baseline() { 51 | async fn test_inner() { 52 | futures::future::join( 53 | async { 54 | yield_now().await; 55 | black_box(1) 56 | }, 57 | async { 58 | yield_now().await; 59 | yield_now().await; 60 | black_box(2) 61 | }, 62 | ) 63 | .await; 64 | } 65 | 66 | for _ in 0..10000 { 67 | test_inner().await; 68 | } 69 | } 70 | 71 | async fn spawn_many(size: usize) { 72 | let registry = Registry::new(Config::default()); 73 | let mut handles = vec![]; 74 | for i in 0..size { 75 | let task = async { 76 | tokio::time::sleep(Duration::from_millis(10)).await; 77 | }; 78 | handles.push(tokio::spawn( 79 | registry.register(i, "new_task").instrument(task), 80 | )); 81 | } 82 | futures::future::try_join_all(handles) 83 | .await 84 | .expect("failed to join background task"); 85 | } 86 | 87 | async fn spawn_many_baseline(size: usize) { 88 | let mut handles = vec![]; 89 | for _ in 0..size { 90 | let task = async { 91 | tokio::time::sleep(Duration::from_millis(10)).await; 92 | }; 93 | handles.push(tokio::spawn(task)); 94 | } 95 | futures::future::try_join_all(handles) 96 | .await 97 | .expect("failed to join background task"); 98 | } 99 | 100 | // time: [6.5488 ms 6.5541 ms 6.5597 ms] 101 | // change: [+6.5978% +6.7838% +6.9299%] (p = 0.00 < 0.05) 102 | // Performance has regressed. 103 | fn bench_basic(c: &mut Criterion) { 104 | c.bench_function("basic", |b| { 105 | b.to_async(runtime()).iter(|| async { 106 | let config = ConfigBuilder::default().verbose(false).build().unwrap(); 107 | let registry = Registry::new(config); 108 | 109 | let root = registry.register(233, "root"); 110 | root.instrument(test()).await; 111 | }) 112 | }); 113 | } 114 | 115 | fn bench_basic_baseline(c: &mut Criterion) { 116 | c.bench_function("basic_baseline", |b| { 117 | b.to_async(runtime()).iter(|| async { 118 | let config = ConfigBuilder::default().verbose(false).build().unwrap(); 119 | let registry = Registry::new(config); 120 | 121 | let root = registry.register(233, "root"); 122 | black_box(root); 123 | test_baseline().await 124 | }) 125 | }); 126 | } 127 | 128 | criterion_group!(benches, bench_basic, bench_basic_baseline); 129 | 130 | // with_register_to_root time: [15.993 ms 16.122 ms 16.292 ms] 131 | // baseline time: [13.940 ms 13.961 ms 13.982 ms] 132 | 133 | fn bench_many_baseline(c: &mut Criterion) { 134 | c.bench_function("with_register_to_root_baseline", |b| { 135 | b.to_async(runtime()) 136 | .iter(|| async { black_box(spawn_many_baseline(10000)).await }) 137 | }); 138 | } 139 | 140 | fn bench_many_exp(c: &mut Criterion) { 141 | c.bench_function("with_register_to_root", |b| { 142 | b.to_async(runtime()) 143 | .iter(|| async { black_box(spawn_many(10000)).await }) 144 | }); 145 | } 146 | 147 | criterion_group!( 148 | name = bench_many; 149 | config = Criterion::default().sample_size(50); 150 | targets = bench_many_exp, bench_many_baseline 151 | ); 152 | 153 | criterion_main!(benches, bench_many); 154 | -------------------------------------------------------------------------------- /examples/basic.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows the basic usage of `await-tree`. 16 | 17 | use std::time::Duration; 18 | 19 | use await_tree::{span, Config, InstrumentAwait, Registry}; 20 | use futures::future::{join, pending}; 21 | use tokio::time::sleep; 22 | 23 | async fn bar(i: i32) { 24 | // `&'static str` span 25 | baz(i).instrument_await("baz in bar").await 26 | } 27 | 28 | async fn baz(i: i32) { 29 | // runtime `String` span is also supported 30 | pending() 31 | .instrument_await(span!("pending in baz {i}")) 32 | .await 33 | } 34 | 35 | async fn foo() { 36 | // spans of joined futures will be siblings in the tree 37 | join( 38 | bar(3).instrument_await("bar"), 39 | baz(2).instrument_await("baz"), 40 | ) 41 | .await; 42 | } 43 | 44 | #[tokio::main] 45 | async fn main() { 46 | let registry = Registry::new(Config::default()); 47 | let root = registry.register((), "foo"); 48 | tokio::spawn(root.instrument(foo())); 49 | 50 | sleep(Duration::from_secs(1)).await; 51 | let tree = registry.get(()).unwrap().to_string(); 52 | 53 | // foo [1.006s] 54 | // bar [1.006s] 55 | // baz in bar [1.006s] 56 | // pending in baz 3 [1.006s] 57 | // baz [1.006s] 58 | // pending in baz 2 [1.006s] 59 | println!("{tree}"); 60 | } 61 | -------------------------------------------------------------------------------- /examples/detach.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows how a span can be detached from and remounted to the tree. 16 | 17 | use std::time::Duration; 18 | 19 | use await_tree::{Config, InstrumentAwait, Registry}; 20 | use futures::channel::oneshot::{self, Receiver}; 21 | use futures::future::{pending, select}; 22 | use futures::FutureExt; 23 | use tokio::time::sleep; 24 | 25 | async fn work(rx: Receiver<()>) { 26 | let mut fut = pending().instrument_await("fut"); 27 | 28 | // poll `fut` under the `select` span 29 | let _ = select( 30 | sleep(Duration::from_millis(500)) 31 | .instrument_await("sleep") 32 | .boxed(), 33 | &mut fut, 34 | ) 35 | .instrument_await("select") 36 | .await; 37 | 38 | // `select` span closed so `fut` is detached 39 | // the elapsed time of `fut` should be preserved 40 | 41 | // wait for the signal to continue 42 | rx.instrument_await("rx").await.unwrap(); 43 | 44 | // poll `fut` under the root `work` span, and it'll be remounted 45 | fut.await 46 | } 47 | 48 | #[tokio::main] 49 | async fn main() { 50 | let registry = Registry::new(Config::default()); 51 | let root = registry.register((), "work"); 52 | let (tx, rx) = oneshot::channel(); 53 | tokio::spawn(root.instrument(work(rx))); 54 | 55 | sleep(Duration::from_millis(100)).await; 56 | let tree = registry.get(()).unwrap().to_string(); 57 | 58 | // work [106.290ms] 59 | // select [106.093ms] 60 | // sleep [106.093ms] 61 | // fut [106.093ms] 62 | println!("{tree}"); 63 | 64 | sleep(Duration::from_secs(1)).await; 65 | let tree = registry.get(()).unwrap().to_string(); 66 | 67 | // work [1.112s] 68 | // rx [606.944ms] 69 | // [Detached 4] 70 | // fut [1.112s] 71 | println!("{tree}"); 72 | 73 | tx.send(()).unwrap(); 74 | sleep(Duration::from_secs(1)).await; 75 | let tree = registry.get(()).unwrap().to_string(); 76 | 77 | // work [2.117s] 78 | // fut [2.117s] 79 | println!("{tree}"); 80 | } 81 | -------------------------------------------------------------------------------- /examples/global.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows the usage of the global registry. 16 | //! 17 | //! Note: This example requires the `tokio` feature to be enabled. 18 | //! Run with: `cargo run --example global --features tokio` 19 | 20 | #![cfg(feature = "tokio")] 21 | 22 | use std::time::Duration; 23 | 24 | use await_tree::{init_global_registry, Config, InstrumentAwait, Registry}; 25 | use futures::future::pending; 26 | 27 | async fn bar() { 28 | pending::<()>().instrument_await("pending").await; 29 | } 30 | 31 | async fn foo() { 32 | await_tree::spawn_anonymous("spawn bar", bar()); 33 | bar().instrument_await("bar").await; 34 | } 35 | 36 | async fn print() { 37 | tokio::time::sleep(Duration::from_secs(1)).await; 38 | 39 | // Access the registry anywhere and collect all trees. 40 | for (key, tree) in Registry::current().collect_all() { 41 | // [Actor 42] 42 | // foo [1.003s] 43 | // bar [1.003s] 44 | // pending [1.003s] 45 | // 46 | // [Anonymous #2] 47 | // spawn bar [1.003s] 48 | // pending [1.003s] 49 | // 50 | // [Print] 51 | // print [1.003s] 52 | println!("[{}]\n{}\n", key, tree); 53 | } 54 | } 55 | 56 | #[tokio::main] 57 | async fn main() { 58 | init_global_registry(Config::default()); 59 | 60 | // After global registry is initialized, the tasks can be spawned everywhere, being 61 | // registered in the global registry. 62 | await_tree::spawn("Actor 42", "foo", foo()); 63 | 64 | // The line above is a shorthand for the following: 65 | tokio::spawn( 66 | Registry::current() 67 | .register("Print", "print") 68 | .instrument(print()), 69 | ) 70 | .await 71 | .unwrap(); 72 | } 73 | -------------------------------------------------------------------------------- /examples/long_running.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2025 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows how to mark a span as "long_running", so that it will 16 | //! not be marked as "!!!" if it takes too long to complete. 17 | 18 | use std::time::Duration; 19 | 20 | use await_tree::{Config, InstrumentAwait, Registry, SpanExt}; 21 | use futures::future::{pending, select}; 22 | use futures::FutureExt; 23 | use tokio::time::sleep; 24 | 25 | async fn long_running_child() { 26 | pending() 27 | .instrument_await("long_running_child".long_running()) 28 | .await 29 | } 30 | 31 | async fn child() { 32 | pending().instrument_await("child").await 33 | } 34 | 35 | async fn foo() { 36 | select(long_running_child().boxed(), child().boxed()).await; 37 | } 38 | 39 | async fn work() -> String { 40 | let registry = Registry::new(Config::default()); 41 | let root = registry.register((), "foo"); 42 | tokio::spawn(root.instrument(foo())); 43 | 44 | // The default threshold is 10 seconds. 45 | sleep(Duration::from_secs(11)).await; 46 | registry.get(()).unwrap().to_string() 47 | } 48 | 49 | #[tokio::main] 50 | async fn main() { 51 | // foo [11.006s] 52 | // long_running_child [11.006s] 53 | // child [!!! 11.006s] 54 | println!("{}", work().await); 55 | } 56 | -------------------------------------------------------------------------------- /examples/multiple.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows how to use `await-tree` for multiple actors. 16 | 17 | use std::time::Duration; 18 | 19 | use await_tree::{span, Config, InstrumentAwait, Registry}; 20 | use futures::future::pending; 21 | use itertools::Itertools; 22 | use tokio::time::sleep; 23 | 24 | async fn work(i: i32) { 25 | foo().instrument_await(span!("actor work {i}")).await 26 | } 27 | 28 | async fn foo() { 29 | pending().instrument_await("pending").await 30 | } 31 | 32 | #[tokio::main] 33 | async fn main() { 34 | let registry = Registry::new(Config::default()); 35 | for i in 0_i32..3 { 36 | let root = registry.register(i, format!("actor {i}")); 37 | tokio::spawn(root.instrument(work(i))); 38 | } 39 | 40 | sleep(Duration::from_secs(1)).await; 41 | 42 | // actor 0 [1.007s] 43 | // actor work 0 [1.007s] 44 | // pending [1.007s] 45 | // 46 | // actor 1 [1.007s] 47 | // actor work 1 [1.007s] 48 | // pending [1.007s] 49 | // 50 | // actor 2 [1.007s] 51 | // actor work 2 [1.007s] 52 | // pending [1.007s] 53 | for (_, tree) in registry 54 | .collect::() 55 | .into_iter() 56 | .sorted_by_key(|(i, _)| *i) 57 | { 58 | println!("{tree}"); 59 | } 60 | } 61 | -------------------------------------------------------------------------------- /examples/serde.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2025 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows the serialization format of the tree. 16 | //! 17 | //! The execution flow is the same as `examples/detach.rs`. 18 | //! 19 | //! Note: This example requires the `serde` feature to be enabled. 20 | //! Run with: `cargo run --example serde --features serde` 21 | 22 | #![cfg(feature = "serde")] 23 | 24 | use std::time::Duration; 25 | 26 | use await_tree::{Config, InstrumentAwait, Registry}; 27 | use futures::channel::oneshot::{self, Receiver}; 28 | use futures::future::{pending, select}; 29 | use futures::FutureExt; 30 | use tokio::time::sleep; 31 | 32 | async fn work(rx: Receiver<()>) { 33 | let mut fut = pending().instrument_await("fut"); 34 | 35 | // poll `fut` under the `select` span 36 | let _ = select( 37 | sleep(Duration::from_millis(500)) 38 | .instrument_await("sleep") 39 | .boxed(), 40 | &mut fut, 41 | ) 42 | .instrument_await("select") 43 | .await; 44 | 45 | // `select` span closed so `fut` is detached 46 | // the elapsed time of `fut` should be preserved 47 | 48 | // wait for the signal to continue 49 | rx.instrument_await("rx").await.unwrap(); 50 | 51 | // poll `fut` under the root `work` span, and it'll be remounted 52 | fut.await 53 | } 54 | 55 | #[tokio::main] 56 | async fn main() { 57 | let registry = Registry::new(Config::default()); 58 | let root = registry.register((), "work"); 59 | let (tx, rx) = oneshot::channel(); 60 | tokio::spawn(root.instrument(work(rx))); 61 | 62 | sleep(Duration::from_millis(100)).await; 63 | let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); 64 | 65 | // { 66 | // "current": 1, 67 | // "tree": { 68 | // "id": 1, 69 | // "span": { 70 | // "name": "work", 71 | // "is_verbose": false, 72 | // "is_long_running": true 73 | // }, 74 | // "elapsed_ns": 105404875, 75 | // "children": [ 76 | // { 77 | // "id": 2, 78 | // "span": { 79 | // "name": "select", 80 | // "is_verbose": false, 81 | // "is_long_running": false 82 | // }, 83 | // "elapsed_ns": 105287624, 84 | // "children": [ 85 | // { 86 | // "id": 3, 87 | // "span": { 88 | // "name": "sleep", 89 | // "is_verbose": false, 90 | // "is_long_running": false 91 | // }, 92 | // "elapsed_ns": 105267874, 93 | // "children": [] 94 | // }, 95 | // { 96 | // "id": 4, 97 | // "span": { 98 | // "name": "fut", 99 | // "is_verbose": false, 100 | // "is_long_running": false 101 | // }, 102 | // "elapsed_ns": 105264874, 103 | // "children": [] 104 | // } 105 | // ] 106 | // } 107 | // ] 108 | // }, 109 | // "detached": [] 110 | // } 111 | println!("{tree}"); 112 | 113 | sleep(Duration::from_secs(1)).await; 114 | let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); 115 | 116 | // { 117 | // "current": 1, 118 | // "tree": { 119 | // "id": 1, 120 | // "span": { 121 | // "name": "work", 122 | // "is_verbose": false, 123 | // "is_long_running": true 124 | // }, 125 | // "elapsed_ns": 1108552791, 126 | // "children": [ 127 | // { 128 | // "id": 3, 129 | // "span": { 130 | // "name": "rx", 131 | // "is_verbose": false, 132 | // "is_long_running": false 133 | // }, 134 | // "elapsed_ns": 603081749, 135 | // "children": [] 136 | // } 137 | // ] 138 | // }, 139 | // "detached": [ 140 | // { 141 | // "id": 4, 142 | // "span": { 143 | // "name": "fut", 144 | // "is_verbose": false, 145 | // "is_long_running": false 146 | // }, 147 | // "elapsed_ns": 1108412791, 148 | // "children": [] 149 | // } 150 | // ] 151 | // } 152 | println!("{tree}"); 153 | 154 | tx.send(()).unwrap(); 155 | sleep(Duration::from_secs(1)).await; 156 | let tree = serde_json::to_string_pretty(®istry.get(()).unwrap()).unwrap(); 157 | 158 | // { 159 | // "current": 1, 160 | // "tree": { 161 | // "id": 1, 162 | // "span": { 163 | // "name": "work", 164 | // "is_verbose": false, 165 | // "is_long_running": true 166 | // }, 167 | // "elapsed_ns": 2114497458, 168 | // "children": [ 169 | // { 170 | // "id": 4, 171 | // "span": { 172 | // "name": "fut", 173 | // "is_verbose": false, 174 | // "is_long_running": false 175 | // }, 176 | // "elapsed_ns": 2114366458, 177 | // "children": [] 178 | // } 179 | // ] 180 | // }, 181 | // "detached": [] 182 | // } 183 | println!("{tree}"); 184 | } 185 | -------------------------------------------------------------------------------- /examples/spawn.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows how to spawn tasks with `await_tree::spawn` that are automatically registered 16 | //! to the current registry of the scope. 17 | //! 18 | //! Note: This example requires the `tokio` feature to be enabled. 19 | //! Run with: `cargo run --example spawn --features tokio` 20 | 21 | #![cfg(feature = "tokio")] 22 | 23 | use std::time::Duration; 24 | 25 | use await_tree::{Config, InstrumentAwait, Registry}; 26 | use futures::future::pending; 27 | use tokio::time::sleep; 28 | 29 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] 30 | struct Actor(usize); 31 | 32 | async fn actor(i: usize) { 33 | // Since we're already inside the scope of a registered/instrumented task, we can directly spawn 34 | // new tasks with `await_tree::spawn` to also register them in the same registry. 35 | await_tree::spawn_anonymous(format!("background task {i}"), async { 36 | pending::<()>().await; 37 | }) 38 | .instrument_await("waiting for background task") 39 | .await 40 | .unwrap(); 41 | } 42 | 43 | #[tokio::main] 44 | async fn main() { 45 | let registry = Registry::new(Config::default()); 46 | 47 | for i in 0..3 { 48 | let root = registry.register(Actor(i), format!("actor {i}")); 49 | tokio::spawn(root.instrument(actor(i))); 50 | } 51 | 52 | sleep(Duration::from_secs(1)).await; 53 | 54 | for (_actor, tree) in registry.collect::() { 55 | // actor 0 [1.004s] 56 | // waiting for background task [1.004s] 57 | println!("{tree}"); 58 | } 59 | for tree in registry.collect_anonymous() { 60 | // background task 0 [1.004s] 61 | println!("{tree}"); 62 | } 63 | } 64 | -------------------------------------------------------------------------------- /examples/verbose.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! This example shows how to mark a span as "verbose", so that it's conditionally 16 | //! enabled based on the config. 17 | 18 | use std::time::Duration; 19 | 20 | use await_tree::{ConfigBuilder, InstrumentAwait, Registry, SpanExt}; 21 | use futures::future::pending; 22 | use tokio::time::sleep; 23 | 24 | async fn foo() { 25 | // verbose span will be disabled if the `verbose` flag in the config is false 26 | pending().instrument_await("pending".verbose()).await 27 | } 28 | 29 | async fn work(verbose: bool) -> String { 30 | let config = ConfigBuilder::default().verbose(verbose).build().unwrap(); 31 | let registry = Registry::new(config); 32 | let root = registry.register((), "foo"); 33 | tokio::spawn(root.instrument(foo())); 34 | 35 | sleep(Duration::from_secs(1)).await; 36 | registry.get(()).unwrap().to_string() 37 | } 38 | 39 | #[tokio::main] 40 | async fn main() { 41 | // foo [1.001s] 42 | println!("{}", work(false).await); 43 | 44 | // foo [1.004s] 45 | // pending [1.004s] 46 | println!("{}", work(true).await); 47 | } 48 | -------------------------------------------------------------------------------- /rust-toolchain: -------------------------------------------------------------------------------- 1 | 1.83 2 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | comment_width = 120 2 | format_code_in_doc_comments = true 3 | format_macro_bodies = true 4 | format_macro_matchers = true 5 | normalize_comments = true 6 | normalize_doc_attributes = true 7 | imports_granularity = "Module" 8 | group_imports = "StdExternalCrate" 9 | reorder_impl_items = true 10 | reorder_imports = true 11 | tab_spaces = 4 12 | wrap_comments = true 13 | -------------------------------------------------------------------------------- /src/context.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::fmt::{Debug, Write}; 16 | use std::sync::atomic::{AtomicU64, Ordering}; 17 | 18 | use indextree::{Arena, NodeId}; 19 | use itertools::Itertools; 20 | use parking_lot::{Mutex, MutexGuard}; 21 | 22 | use crate::root::current_context; 23 | use crate::Span; 24 | 25 | /// Node in the span tree. 26 | #[derive(Debug, Clone)] 27 | struct SpanNode { 28 | /// The span value. 29 | span: Span, 30 | 31 | /// The time when this span was started, or the future was first polled. 32 | start_time: coarsetime::Instant, 33 | } 34 | 35 | impl SpanNode { 36 | /// Create a new node with the given value. 37 | fn new(span: Span) -> Self { 38 | Self { 39 | span, 40 | start_time: coarsetime::Instant::now(), 41 | } 42 | } 43 | } 44 | 45 | /// The id of an await-tree context. 46 | /// 47 | /// We will check the id recorded in the instrumented future against the current task-local context 48 | /// before trying to update the tree. 49 | /// 50 | /// Also used as the key for anonymous trees in the registry. Intentionally made private to prevent 51 | /// users from reusing the same id when registering a new tree. 52 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] 53 | pub(crate) struct ContextId(pub(crate) u64); 54 | 55 | /// An await-tree for a task. 56 | #[derive(Debug, Clone)] 57 | pub struct Tree { 58 | /// The arena for allocating span nodes in this context. 59 | arena: Arena, 60 | 61 | /// The root span node. 62 | root: NodeId, 63 | 64 | /// The current span node. This is the node that is currently being polled. 65 | current: NodeId, 66 | } 67 | 68 | #[cfg(feature = "serde")] 69 | mod serde_impl { 70 | use serde::ser::SerializeStruct as _; 71 | use serde::Serialize; 72 | 73 | use super::*; 74 | 75 | struct SpanNodeSer<'a> { 76 | arena: &'a Arena, 77 | node: NodeId, 78 | } 79 | 80 | impl Serialize for SpanNodeSer<'_> { 81 | fn serialize(&self, serializer: S) -> Result 82 | where 83 | S: serde::Serializer, 84 | { 85 | let inner = self.arena[self.node].get(); 86 | let mut s = serializer.serialize_struct("Span", 4)?; 87 | 88 | // Basic info. 89 | let id: usize = self.node.into(); 90 | s.serialize_field("id", &id)?; 91 | s.serialize_field("span", &inner.span)?; 92 | s.serialize_field("elapsed_ns", &inner.start_time.elapsed().as_nanos())?; 93 | 94 | // Children. 95 | let children = (self.node.children(self.arena)) 96 | .map(|node| SpanNodeSer { 97 | arena: self.arena, 98 | node, 99 | }) 100 | .sorted_by_key(|child| { 101 | let inner = self.arena[child.node].get(); 102 | inner.start_time 103 | }) 104 | .collect_vec(); 105 | s.serialize_field("children", &children)?; 106 | 107 | s.end() 108 | } 109 | } 110 | 111 | impl Serialize for Tree { 112 | fn serialize(&self, serializer: S) -> Result 113 | where 114 | S: serde::Serializer, 115 | { 116 | let mut s = serializer.serialize_struct("Tree", 3)?; 117 | 118 | let current_id: usize = self.current.into(); 119 | s.serialize_field("current", ¤t_id)?; 120 | 121 | // The main tree. 122 | s.serialize_field( 123 | "tree", 124 | &SpanNodeSer { 125 | arena: &self.arena, 126 | node: self.root, 127 | }, 128 | )?; 129 | 130 | // The detached subtrees. 131 | let detached = self 132 | .detached_roots() 133 | .map(|node| SpanNodeSer { 134 | arena: &self.arena, 135 | node, 136 | }) 137 | .collect_vec(); 138 | s.serialize_field("detached", &detached)?; 139 | 140 | s.end() 141 | } 142 | } 143 | } 144 | 145 | impl std::fmt::Display for Tree { 146 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 147 | fn fmt_node( 148 | f: &mut std::fmt::Formatter<'_>, 149 | arena: &Arena, 150 | node: NodeId, 151 | depth: usize, 152 | current: NodeId, 153 | ) -> std::fmt::Result { 154 | f.write_str(&" ".repeat(depth * 2))?; 155 | 156 | let inner = arena[node].get(); 157 | f.write_str(&inner.span.name)?; 158 | 159 | let elapsed: std::time::Duration = inner.start_time.elapsed().into(); 160 | write!( 161 | f, 162 | " [{}{:.3?}]", 163 | if !inner.span.is_long_running && elapsed.as_secs() >= 10 { 164 | "!!! " 165 | } else { 166 | "" 167 | }, 168 | elapsed 169 | )?; 170 | 171 | if depth > 0 && node == current { 172 | f.write_str(" <== current")?; 173 | } 174 | 175 | f.write_char('\n')?; 176 | for child in node 177 | .children(arena) 178 | .sorted_by_key(|&id| arena[id].get().start_time) 179 | { 180 | fmt_node(f, arena, child, depth + 1, current)?; 181 | } 182 | 183 | Ok(()) 184 | } 185 | 186 | fmt_node(f, &self.arena, self.root, 0, self.current)?; 187 | 188 | // Format all detached spans. 189 | for node in self.detached_roots() { 190 | writeln!(f, "[Detached {node}]")?; 191 | fmt_node(f, &self.arena, node, 1, self.current)?; 192 | } 193 | 194 | Ok(()) 195 | } 196 | } 197 | 198 | impl Tree { 199 | /// Get the count of active span nodes in this context. 200 | #[cfg(test)] 201 | pub(crate) fn active_node_count(&self) -> usize { 202 | self.arena.iter().filter(|n| !n.is_removed()).count() 203 | } 204 | 205 | /// Get the count of active detached span nodes in this context. 206 | #[cfg(test)] 207 | pub(crate) fn detached_node_count(&self) -> usize { 208 | self.arena 209 | .iter() 210 | .filter(|n| { 211 | !n.is_removed() 212 | && n.parent().is_none() 213 | && self.arena.get_node_id(n).unwrap() != self.root 214 | }) 215 | .count() 216 | } 217 | 218 | /// Returns an iterator over the root nodes of detached subtrees. 219 | fn detached_roots(&self) -> impl Iterator + '_ { 220 | self.arena 221 | .iter() 222 | .filter(|n| !n.is_removed()) // still valid 223 | .map(|node| self.arena.get_node_id(node).unwrap()) // get id 224 | .filter(|&id| id != self.root && self.arena[id].parent().is_none()) // no parent but non-root 225 | } 226 | 227 | /// Push a new span as a child of current span, used for future firstly polled. 228 | /// 229 | /// Returns the new current span. 230 | pub(crate) fn push(&mut self, span: Span) -> NodeId { 231 | let child = self.arena.new_node(SpanNode::new(span)); 232 | self.current.prepend(child, &mut self.arena); 233 | self.current = child; 234 | child 235 | } 236 | 237 | /// Step in the current span to the given child, used for future polled again. 238 | /// 239 | /// If the child is not actually a child of the current span, it means we are using a new future 240 | /// to poll it, so we need to detach it from the previous parent, and attach it to the current 241 | /// span. 242 | pub(crate) fn step_in(&mut self, child: NodeId) { 243 | if !self.current.children(&self.arena).contains(&child) { 244 | // Actually we can always call this even if `child` is already a child of `current`. But 245 | // checking first performs better. 246 | self.current.prepend(child, &mut self.arena); 247 | } 248 | self.current = child; 249 | } 250 | 251 | /// Pop the current span to the parent, used for future ready. 252 | /// 253 | /// Note that there might still be some children of this node, like `select_stream.next()`. 254 | /// The children might be polled again later, and will be attached as the children of a new 255 | /// span. 256 | pub(crate) fn pop(&mut self) { 257 | let parent = self.arena[self.current] 258 | .parent() 259 | .expect("the root node should not be popped"); 260 | self.remove_and_detach(self.current); 261 | self.current = parent; 262 | } 263 | 264 | /// Step out the current span to the parent, used for future pending. 265 | pub(crate) fn step_out(&mut self) { 266 | let parent = self.arena[self.current] 267 | .parent() 268 | .expect("the root node should not be stepped out"); 269 | self.current = parent; 270 | } 271 | 272 | /// Remove the current span and detach the children, used for future aborting. 273 | /// 274 | /// The children might be polled again later, and will be attached as the children of a new 275 | /// span. 276 | pub(crate) fn remove_and_detach(&mut self, node: NodeId) { 277 | node.detach(&mut self.arena); 278 | // Removing detached `node` makes children detached. 279 | node.remove(&mut self.arena); 280 | } 281 | 282 | /// Get the current span node id. 283 | pub(crate) fn current(&self) -> NodeId { 284 | self.current 285 | } 286 | } 287 | 288 | /// The task-local await-tree context. 289 | #[derive(Debug)] 290 | pub(crate) struct TreeContext { 291 | /// The id of the context. 292 | id: ContextId, 293 | 294 | /// Whether to include the "verbose" span in the tree. 295 | verbose: bool, 296 | 297 | /// The await-tree. 298 | tree: Mutex, 299 | } 300 | 301 | impl TreeContext { 302 | /// Create a new context. 303 | pub(crate) fn new(root_span: Span, verbose: bool) -> Self { 304 | static ID: AtomicU64 = AtomicU64::new(0); 305 | let id = ID.fetch_add(1, Ordering::Relaxed); 306 | 307 | // Always make the root span long-running. 308 | let root_span = root_span.long_running(); 309 | 310 | let mut arena = Arena::new(); 311 | let root = arena.new_node(SpanNode::new(root_span)); 312 | 313 | Self { 314 | id: ContextId(id), 315 | verbose, 316 | tree: Tree { 317 | arena, 318 | root, 319 | current: root, 320 | } 321 | .into(), 322 | } 323 | } 324 | 325 | /// Get the context id. 326 | pub(crate) fn id(&self) -> ContextId { 327 | self.id 328 | } 329 | 330 | /// Returns the locked guard of the tree. 331 | pub(crate) fn tree(&self) -> MutexGuard<'_, Tree> { 332 | self.tree.lock() 333 | } 334 | 335 | /// Whether the verbose span should be included. 336 | pub(crate) fn verbose(&self) -> bool { 337 | self.verbose 338 | } 339 | } 340 | 341 | /// Get the await-tree of current task. Returns `None` if we're not instrumented. 342 | /// 343 | /// This is useful if you want to check which component or runtime task is calling this function. 344 | pub fn current_tree() -> Option { 345 | current_context().map(|c| c.tree().clone()) 346 | } 347 | -------------------------------------------------------------------------------- /src/future.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::future::Future; 16 | use std::pin::Pin; 17 | use std::task::Poll; 18 | 19 | use indextree::NodeId; 20 | use pin_project::{pin_project, pinned_drop}; 21 | 22 | use crate::context::ContextId; 23 | use crate::root::current_context; 24 | use crate::Span; 25 | 26 | enum State { 27 | Initial(Span), 28 | Polled { 29 | this_node: NodeId, 30 | this_context_id: ContextId, 31 | }, 32 | Ready, 33 | /// This span is disabled due to `verbose` configuration. 34 | Disabled, 35 | } 36 | 37 | /// The future for [`InstrumentAwait`][ia]. 38 | /// 39 | /// [ia]: crate::InstrumentAwait 40 | #[pin_project(PinnedDrop)] 41 | pub struct Instrumented { 42 | #[pin] 43 | inner: F, 44 | state: State, 45 | } 46 | 47 | impl Instrumented { 48 | pub(crate) fn new(inner: F, span: Span) -> Self { 49 | Self { 50 | inner, 51 | state: State::Initial(span), 52 | } 53 | } 54 | } 55 | 56 | impl Future for Instrumented { 57 | type Output = F::Output; 58 | 59 | fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { 60 | let this = self.project(); 61 | let context = current_context(); 62 | 63 | let (context, this_node) = match this.state { 64 | State::Initial(span) => { 65 | match context { 66 | Some(c) => { 67 | if !c.verbose() && span.is_verbose { 68 | // The tracing for this span is disabled according to the verbose 69 | // configuration. 70 | *this.state = State::Disabled; 71 | return this.inner.poll(cx); 72 | } 73 | // First polled, push a new span to the context. 74 | let node = c.tree().push(std::mem::take(span)); 75 | *this.state = State::Polled { 76 | this_node: node, 77 | this_context_id: c.id(), 78 | }; 79 | (c, node) 80 | } 81 | // Not in a context 82 | None => return this.inner.poll(cx), 83 | } 84 | } 85 | State::Polled { 86 | this_node, 87 | this_context_id: this_context, 88 | } => { 89 | match context { 90 | // Context correct 91 | Some(c) if c.id() == *this_context => { 92 | // Polled before, just step in. 93 | c.tree().step_in(*this_node); 94 | (c, *this_node) 95 | } 96 | // Context changed 97 | Some(_) => { 98 | tracing::warn!( 99 | "future polled in a different context as it was first polled" 100 | ); 101 | return this.inner.poll(cx); 102 | } 103 | // Out of context 104 | None => { 105 | tracing::warn!( 106 | "future polled not in a context, while it was when first polled" 107 | ); 108 | return this.inner.poll(cx); 109 | } 110 | } 111 | } 112 | State::Ready => unreachable!("the instrumented future should always be fused"), 113 | State::Disabled => return this.inner.poll(cx), 114 | }; 115 | 116 | // The current node must be the this_node. 117 | debug_assert_eq!(this_node, context.tree().current()); 118 | 119 | match this.inner.poll(cx) { 120 | // The future is ready, clean-up this span by popping from the context. 121 | Poll::Ready(output) => { 122 | context.tree().pop(); 123 | *this.state = State::Ready; 124 | Poll::Ready(output) 125 | } 126 | // Still pending, just step out. 127 | Poll::Pending => { 128 | context.tree().step_out(); 129 | Poll::Pending 130 | } 131 | } 132 | } 133 | } 134 | 135 | #[pinned_drop] 136 | impl PinnedDrop for Instrumented { 137 | fn drop(self: Pin<&mut Self>) { 138 | let this = self.project(); 139 | 140 | match this.state { 141 | State::Polled { 142 | this_node, 143 | this_context_id, 144 | } => match current_context() { 145 | // Context correct 146 | Some(c) if c.id() == *this_context_id => { 147 | c.tree().remove_and_detach(*this_node); 148 | } 149 | // Context changed 150 | Some(_) => { 151 | tracing::warn!("future is dropped in a different context as it was first polled, cannot clean up!"); 152 | } 153 | // Out of context 154 | None => { 155 | tracing::warn!("future is not in a context, while it was when first polled, cannot clean up!"); 156 | } 157 | }, 158 | State::Initial(_) | State::Ready | State::Disabled => {} 159 | } 160 | } 161 | } 162 | -------------------------------------------------------------------------------- /src/global.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::sync::OnceLock; 16 | 17 | use crate::{Config, Registry}; 18 | 19 | static GLOBAL_REGISTRY: OnceLock = OnceLock::new(); 20 | 21 | /// Initialize the global registry with the given configuration. 22 | /// Panics if the global registry has already been initialized. 23 | /// 24 | /// This is **optional** and only needed if you want to use the global registry. 25 | /// You can always create a new registry with [`Registry::new`] and pass it around to achieve 26 | /// better encapsulation. 27 | pub fn init_global_registry(config: Config) { 28 | if let Err(_r) = GLOBAL_REGISTRY.set(Registry::new(config)) { 29 | panic!("global registry already initialized") 30 | } 31 | } 32 | 33 | pub(crate) fn global_registry() -> Option { 34 | GLOBAL_REGISTRY.get().cloned() 35 | } 36 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! Generate accurate and informative tree dumps of asynchronous tasks. 16 | //! 17 | //! # Example 18 | //! 19 | //! Below is a basic example of how to trace asynchronous tasks with the global registry of the 20 | //! `await-tree` crate. 21 | //! 22 | //! ```rust 23 | //! # use std::time::Duration; 24 | //! # use tokio::time::sleep; 25 | //! # use await_tree::{InstrumentAwait, Registry, span}; 26 | //! # use futures::future::join; 27 | //! # 28 | //! # async fn work() { futures::future::pending::<()>().await } 29 | //! # 30 | //! async fn bar(i: i32) { 31 | //! // `&'static str` span 32 | //! baz(i).instrument_await("baz in bar").await 33 | //! } 34 | //! 35 | //! async fn baz(i: i32) { 36 | //! // runtime `String` span is also supported 37 | //! work().instrument_await(span!("working in baz {i}")).await 38 | //! } 39 | //! 40 | //! async fn foo() { 41 | //! // spans of joined futures will be siblings in the tree 42 | //! join( 43 | //! bar(3).instrument_await("bar"), 44 | //! baz(2).instrument_await("baz"), 45 | //! ) 46 | //! .await; 47 | //! } 48 | //! 49 | //! # #[tokio::main] 50 | //! # async fn main() { 51 | //! // Init the global registry to start tracing the tasks. 52 | //! await_tree::init_global_registry(Default::default()); 53 | //! // Spawn a task with root span "foo" and key "foo". 54 | //! await_tree::spawn("foo", "foo", foo()); 55 | //! // Let the tasks run for a while. 56 | //! sleep(Duration::from_secs(1)).await; 57 | //! // Get the tree of the task with key "foo". 58 | //! let tree = Registry::current().get("foo").unwrap(); 59 | //! 60 | //! // foo [1.006s] 61 | //! // bar [1.006s] 62 | //! // baz in bar [1.006s] 63 | //! // working in baz 3 [1.006s] 64 | //! // baz [1.006s] 65 | //! // working in baz 2 [1.006s] 66 | //! println!("{tree}"); 67 | //! # } 68 | //! ``` 69 | 70 | #![forbid(missing_docs)] 71 | 72 | use std::future::Future; 73 | 74 | mod context; 75 | mod future; 76 | mod global; 77 | mod obj_utils; 78 | mod registry; 79 | mod root; 80 | mod span; 81 | #[cfg(feature = "tokio")] 82 | mod spawn; 83 | 84 | pub use context::{current_tree, Tree}; 85 | pub use future::Instrumented; 86 | pub use global::init_global_registry; 87 | pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry}; 88 | pub use root::TreeRoot; 89 | pub use span::{Span, SpanExt}; 90 | #[cfg(feature = "tokio")] 91 | pub use spawn::{spawn, spawn_anonymous}; 92 | 93 | #[doc(hidden)] 94 | pub mod __private { 95 | pub use crate::span::fmt_span; 96 | } 97 | 98 | /// Attach spans to a future to be traced in the await-tree. 99 | pub trait InstrumentAwait: Future + Sized { 100 | /// Instrument the future with a span. 101 | fn instrument_await(self, span: impl Into) -> Instrumented { 102 | Instrumented::new(self, span.into()) 103 | } 104 | } 105 | impl InstrumentAwait for F where F: Future {} 106 | 107 | #[cfg(test)] 108 | mod tests; 109 | -------------------------------------------------------------------------------- /src/obj_utils.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | //! Utilities for using `Any` as the key of a `HashMap`. 16 | 17 | // Adopted from: 18 | // https://github.com/bevyengine/bevy/blob/56bcbb097552b45e3ff48c48947ed8ee4e2c24b1/crates/bevy_utils/src/label.rs 19 | 20 | use std::any::Any; 21 | use std::hash::{Hash, Hasher}; 22 | 23 | /// An object safe version of [`Eq`]. This trait is automatically implemented 24 | /// for any `'static` type that implements `Eq`. 25 | pub(crate) trait DynEq: Any { 26 | /// Casts the type to `dyn Any`. 27 | fn as_any(&self) -> &dyn Any; 28 | 29 | /// This method tests for `self` and `other` values to be equal. 30 | /// 31 | /// Implementers should avoid returning `true` when the underlying types are 32 | /// not the same. 33 | fn dyn_eq(&self, other: &dyn DynEq) -> bool; 34 | } 35 | 36 | impl DynEq for T 37 | where 38 | T: Any + Eq, 39 | { 40 | fn as_any(&self) -> &dyn Any { 41 | self 42 | } 43 | 44 | fn dyn_eq(&self, other: &dyn DynEq) -> bool { 45 | if let Some(other) = other.as_any().downcast_ref::() { 46 | return self == other; 47 | } 48 | false 49 | } 50 | } 51 | 52 | /// An object safe version of [`Hash`]. This trait is automatically implemented 53 | /// for any `'static` type that implements `Hash`. 54 | pub(crate) trait DynHash: DynEq { 55 | /// Casts the type to `dyn Any`. 56 | fn as_dyn_eq(&self) -> &dyn DynEq; 57 | 58 | /// Feeds this value into the given [`Hasher`]. 59 | fn dyn_hash(&self, state: &mut dyn Hasher); 60 | } 61 | 62 | impl DynHash for T 63 | where 64 | T: DynEq + Hash, 65 | { 66 | fn as_dyn_eq(&self) -> &dyn DynEq { 67 | self 68 | } 69 | 70 | fn dyn_hash(&self, mut state: &mut dyn Hasher) { 71 | T::hash(self, &mut state); 72 | self.type_id().hash(&mut state); 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /src/registry.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::any::Any; 16 | use std::fmt::{Debug, Display}; 17 | use std::hash::Hash; 18 | use std::sync::{Arc, Weak}; 19 | 20 | use derive_builder::Builder; 21 | use parking_lot::RwLock; 22 | use weak_table::WeakValueHashMap; 23 | 24 | use crate::context::{ContextId, Tree, TreeContext}; 25 | use crate::obj_utils::{DynEq, DynHash}; 26 | use crate::{Span, TreeRoot}; 27 | 28 | /// Configuration for an await-tree registry, which affects the behavior of all await-trees in the 29 | /// registry. 30 | #[derive(Debug, Clone, Builder)] 31 | #[builder(default)] 32 | pub struct Config { 33 | /// Whether to include the **verbose** span in the await-tree. 34 | verbose: bool, 35 | } 36 | 37 | #[allow(clippy::derivable_impls)] 38 | impl Default for Config { 39 | fn default() -> Self { 40 | Self { verbose: false } 41 | } 42 | } 43 | 44 | /// A key that can be used to identify a task and its await-tree in the [`Registry`]. 45 | /// 46 | /// All thread-safe types that can be used as a key of a hash map are automatically implemented with 47 | /// this trait. 48 | pub trait Key: Hash + Eq + Debug + Send + Sync + 'static {} 49 | impl Key for T where T: Hash + Eq + Debug + Send + Sync + 'static {} 50 | 51 | /// The object-safe version of [`Key`], automatically implemented. 52 | trait ObjKey: DynHash + DynEq + Debug + Send + Sync + 'static {} 53 | impl ObjKey for T where T: DynHash + DynEq + Debug + Send + Sync + 'static {} 54 | 55 | /// Key type for anonymous await-trees. 56 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] 57 | struct AnonymousKey(ContextId); 58 | 59 | impl Display for AnonymousKey { 60 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 61 | write!(f, "Anonymous #{}", self.0 .0) 62 | } 63 | } 64 | 65 | /// Type-erased key for the [`Registry`]. 66 | #[derive(Clone)] 67 | pub struct AnyKey(Arc); 68 | 69 | impl PartialEq for AnyKey { 70 | fn eq(&self, other: &Self) -> bool { 71 | self.0.dyn_eq(other.0.as_dyn_eq()) 72 | } 73 | } 74 | 75 | impl Eq for AnyKey {} 76 | 77 | impl Hash for AnyKey { 78 | fn hash(&self, state: &mut H) { 79 | self.0.dyn_hash(state); 80 | } 81 | } 82 | 83 | impl Debug for AnyKey { 84 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 85 | self.0.fmt(f) 86 | } 87 | } 88 | 89 | impl Display for AnyKey { 90 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 91 | // TODO: for all `impl Display`? 92 | macro_rules! delegate_to_display { 93 | ($($t:ty),* $(,)?) => { 94 | $( 95 | if let Some(k) = self.as_any().downcast_ref::<$t>() { 96 | return write!(f, "{}", k); 97 | } 98 | )* 99 | }; 100 | } 101 | delegate_to_display!(String, &str, AnonymousKey); 102 | 103 | write!(f, "{:?}", self) 104 | } 105 | } 106 | 107 | impl AnyKey { 108 | fn new(key: impl ObjKey) -> Self { 109 | Self(Arc::new(key)) 110 | } 111 | 112 | /// Cast the key to `dyn Any`. 113 | pub fn as_any(&self) -> &dyn Any { 114 | self.0.as_ref().as_any() 115 | } 116 | 117 | /// Returns whether the key is of type `K`. 118 | /// 119 | /// Equivalent to `self.as_any().is::()`. 120 | pub fn is(&self) -> bool { 121 | self.as_any().is::() 122 | } 123 | 124 | /// Returns whether the key corresponds to an anonymous await-tree. 125 | pub fn is_anonymous(&self) -> bool { 126 | self.as_any().is::() 127 | } 128 | 129 | /// Returns the key as a reference to type `K`, if it is of type `K`. 130 | /// 131 | /// Equivalent to `self.as_any().downcast_ref::()`. 132 | pub fn downcast_ref(&self) -> Option<&K> { 133 | self.as_any().downcast_ref() 134 | } 135 | } 136 | 137 | type Contexts = RwLock>>; 138 | 139 | struct RegistryCore { 140 | contexts: Contexts, 141 | config: Config, 142 | } 143 | 144 | /// The registry of multiple await-trees. 145 | /// 146 | /// Can be cheaply cloned to share the same registry. 147 | pub struct Registry(Arc); 148 | 149 | impl Debug for Registry { 150 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 151 | f.debug_struct("Registry") 152 | .field("config", self.config()) 153 | .finish_non_exhaustive() 154 | } 155 | } 156 | 157 | impl Clone for Registry { 158 | fn clone(&self) -> Self { 159 | Self(Arc::clone(&self.0)) 160 | } 161 | } 162 | 163 | impl Registry { 164 | fn contexts(&self) -> &Contexts { 165 | &self.0.contexts 166 | } 167 | 168 | fn config(&self) -> &Config { 169 | &self.0.config 170 | } 171 | } 172 | 173 | impl Registry { 174 | /// Create a new registry with given `config`. 175 | pub fn new(config: Config) -> Self { 176 | Self( 177 | RegistryCore { 178 | contexts: Default::default(), 179 | config, 180 | } 181 | .into(), 182 | ) 183 | } 184 | 185 | /// Returns the current registry, if exists. 186 | /// 187 | /// 1. If the current task is registered with a registry, returns the registry. 188 | /// 2. If the global registry is initialized with 189 | /// [`init_global_registry`](crate::global::init_global_registry), returns the global 190 | /// registry. 191 | /// 3. Otherwise, returns `None`. 192 | pub fn try_current() -> Option { 193 | crate::root::current_registry() 194 | } 195 | 196 | /// Returns the current registry, panics if not exists. 197 | /// 198 | /// See [`Registry::try_current`] for more information. 199 | pub fn current() -> Self { 200 | Self::try_current().expect("no current registry") 201 | } 202 | 203 | fn register_inner(&self, key: impl Key, context: Arc) -> TreeRoot { 204 | self.contexts() 205 | .write() 206 | .insert(AnyKey::new(key), Arc::clone(&context)); 207 | 208 | TreeRoot { 209 | context, 210 | registry: WeakRegistry(Arc::downgrade(&self.0)), 211 | } 212 | } 213 | 214 | /// Register with given key. Returns a [`TreeRoot`] that can be used to instrument a future. 215 | /// 216 | /// If the key already exists, a new [`TreeRoot`] is returned and the reference to the old 217 | /// [`TreeRoot`] is dropped. 218 | pub fn register(&self, key: impl Key, root_span: impl Into) -> TreeRoot { 219 | let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose)); 220 | self.register_inner(key, context) 221 | } 222 | 223 | /// Register an anonymous await-tree without specifying a key. Returns a [`TreeRoot`] that can 224 | /// be used to instrument a future. 225 | /// 226 | /// Anonymous await-trees are not able to be retrieved through the [`Registry::get`] method. Use 227 | /// [`Registry::collect_anonymous`] or [`Registry::collect_all`] to collect them. 228 | // TODO: we have keyed and anonymous, should we also have a typed-anonymous (for classification 229 | // only)? 230 | pub fn register_anonymous(&self, root_span: impl Into) -> TreeRoot { 231 | let context = Arc::new(TreeContext::new(root_span.into(), self.config().verbose)); 232 | self.register_inner(AnonymousKey(context.id()), context) // use the private id as the key 233 | } 234 | 235 | /// Get a clone of the await-tree with given key. 236 | /// 237 | /// Returns `None` if the key does not exist or the tree root has been dropped. 238 | pub fn get(&self, key: impl Key) -> Option { 239 | self.contexts() 240 | .read() 241 | .get(&AnyKey::new(key)) // TODO: accept ref can? 242 | .map(|v| v.tree().clone()) 243 | } 244 | 245 | /// Remove all the registered await-trees. 246 | pub fn clear(&self) { 247 | self.contexts().write().clear(); 248 | } 249 | 250 | /// Collect the snapshots of all await-trees with the key of type `K`. 251 | pub fn collect(&self) -> Vec<(K, Tree)> { 252 | self.contexts() 253 | .read() 254 | .iter() 255 | .filter_map(|(k, v)| { 256 | k.0.as_ref() 257 | .as_any() 258 | .downcast_ref::() 259 | .map(|k| (k.clone(), v.tree().clone())) 260 | }) 261 | .collect() 262 | } 263 | 264 | /// Collect the snapshots of all await-trees registered with [`Registry::register_anonymous`]. 265 | pub fn collect_anonymous(&self) -> Vec { 266 | self.contexts() 267 | .read() 268 | .iter() 269 | .filter_map(|(k, v)| { 270 | if k.is_anonymous() { 271 | Some(v.tree().clone()) 272 | } else { 273 | None 274 | } 275 | }) 276 | .collect() 277 | } 278 | 279 | /// Collect the snapshots of all await-trees regardless of the key type. 280 | pub fn collect_all(&self) -> Vec<(AnyKey, Tree)> { 281 | self.contexts() 282 | .read() 283 | .iter() 284 | .map(|(k, v)| (k.clone(), v.tree().clone())) 285 | .collect() 286 | } 287 | } 288 | 289 | pub(crate) struct WeakRegistry(Weak); 290 | 291 | impl WeakRegistry { 292 | pub fn upgrade(&self) -> Option { 293 | self.0.upgrade().map(Registry) 294 | } 295 | } 296 | 297 | #[cfg(test)] 298 | mod tests { 299 | use super::*; 300 | 301 | #[test] 302 | fn test_registry() { 303 | let registry = Registry::new(Config::default()); 304 | 305 | let _0_i32 = registry.register(0_i32, "0"); 306 | let _1_i32 = registry.register(1_i32, "1"); 307 | let _2_i32 = registry.register(2_i32, "2"); 308 | 309 | let _0_str = registry.register("0", "0"); 310 | let _1_str = registry.register("1", "1"); 311 | 312 | let _unit = registry.register((), "()"); 313 | let _unit_replaced = registry.register((), "[]"); 314 | 315 | let _anon = registry.register_anonymous("anon"); 316 | let _anon = registry.register_anonymous("anon"); 317 | 318 | let i32s = registry.collect::(); 319 | assert_eq!(i32s.len(), 3); 320 | 321 | let strs = registry.collect::<&'static str>(); 322 | assert_eq!(strs.len(), 2); 323 | 324 | let units = registry.collect::<()>(); 325 | assert_eq!(units.len(), 1); 326 | 327 | let anons = registry.collect_anonymous(); 328 | assert_eq!(anons.len(), 2); 329 | 330 | let all = registry.collect_all(); 331 | assert_eq!(all.len(), 8); 332 | } 333 | } 334 | -------------------------------------------------------------------------------- /src/root.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use std::future::Future; 16 | use std::sync::Arc; 17 | 18 | use crate::context::TreeContext; 19 | use crate::global::global_registry; 20 | use crate::registry::WeakRegistry; 21 | use crate::Registry; 22 | 23 | /// The root of an await-tree. 24 | pub struct TreeRoot { 25 | pub(crate) context: Arc, 26 | pub(crate) registry: WeakRegistry, 27 | } 28 | 29 | task_local::task_local! { 30 | static ROOT: TreeRoot 31 | } 32 | 33 | pub(crate) fn current_context() -> Option> { 34 | ROOT.try_with(|r| r.context.clone()).ok() 35 | } 36 | 37 | pub(crate) fn current_registry() -> Option { 38 | let local = || ROOT.try_with(|r| r.registry.upgrade()).ok().flatten(); 39 | let global = global_registry; 40 | 41 | local().or_else(global) 42 | } 43 | 44 | impl TreeRoot { 45 | /// Instrument the given future with the context of this tree root. 46 | pub async fn instrument(self, future: F) -> F::Output { 47 | ROOT.scope(self, future).await 48 | } 49 | } 50 | -------------------------------------------------------------------------------- /src/span.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2025 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | type SpanName = flexstr::SharedStr; 16 | 17 | #[doc(hidden)] 18 | pub fn fmt_span(args: std::fmt::Arguments<'_>) -> Span { 19 | let name = if let Some(str) = args.as_str() { 20 | SpanName::from_ref(str) 21 | } else { 22 | flexstr::flex_fmt(args) 23 | }; 24 | Span::new(name) 25 | } 26 | 27 | /// Creates a new span with formatted name. 28 | /// 29 | /// [`instrument_await`] accepts any type that implements [`AsRef`] as the span name. 30 | /// This macro provides similar functionality to [`format!`], but with improved performance 31 | /// by creating the span name on the stack when possible, avoiding unnecessary allocations. 32 | /// 33 | /// [`instrument_await`]: crate::InstrumentAwait::instrument_await 34 | #[macro_export] 35 | // XXX: Without this extra binding (`let res = ..`), it will make the future `!Send`. 36 | // This is also how `std::format!` behaves. But why? 37 | macro_rules! span { 38 | ($($fmt_arg:tt)*) => {{ 39 | let res = $crate::__private::fmt_span(format_args!($($fmt_arg)*)); 40 | res 41 | }}; 42 | } 43 | 44 | /// A cheaply cloneable span in the await-tree. 45 | #[cfg_attr(feature = "serde", derive(serde::Serialize))] 46 | #[derive(Debug, Clone, Default, PartialEq, Eq, PartialOrd, Ord)] 47 | pub struct Span { 48 | pub(crate) name: SpanName, 49 | pub(crate) is_verbose: bool, 50 | pub(crate) is_long_running: bool, 51 | } 52 | 53 | impl Span { 54 | fn new(name: SpanName) -> Self { 55 | Self { 56 | name, 57 | is_verbose: false, 58 | is_long_running: false, 59 | } 60 | } 61 | } 62 | 63 | impl Span { 64 | /// Set the verbose attribute of the span. 65 | /// 66 | /// When a span is marked as verbose, it will be included in the output 67 | /// only if the `verbose` flag in the [`Config`] is set. 68 | /// 69 | /// [`Config`]: crate::Config 70 | pub fn verbose(mut self) -> Self { 71 | self.is_verbose = true; 72 | self 73 | } 74 | 75 | /// Set the long-running attribute of the span. 76 | /// 77 | /// When a span is marked as long-running, it will not be marked as "!!!" 78 | /// in the formatted [`Tree`] if it takes too long to complete. The root 79 | /// span is always marked as long-running. 80 | /// 81 | /// [`Tree`]: crate::Tree 82 | pub fn long_running(mut self) -> Self { 83 | self.is_long_running = true; 84 | self 85 | } 86 | } 87 | 88 | /// Convert a value into a span and set attributes. 89 | #[easy_ext::ext(SpanExt)] 90 | impl> T { 91 | /// Convert `self` into a span and set the verbose attribute. 92 | /// 93 | /// See [`Span::verbose`] for more details. 94 | pub fn verbose(self) -> Span { 95 | self.into().verbose() 96 | } 97 | 98 | /// Convert `self` into a span and set the long-running attribute. 99 | /// 100 | /// See [`Span::long_running`] for more details. 101 | pub fn long_running(self) -> Span { 102 | self.into().long_running() 103 | } 104 | } 105 | 106 | impl> From for Span { 107 | fn from(value: S) -> Self { 108 | Self::new(SpanName::from_ref(value)) 109 | } 110 | } 111 | 112 | impl std::fmt::Display for Span { 113 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { 114 | self.name.fmt(f) 115 | } 116 | } 117 | -------------------------------------------------------------------------------- /src/spawn.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | // TODO: should we consider exposing `current_registry` 16 | // so that users can not only spawn tasks but also get and collect trees? 17 | 18 | use std::future::Future; 19 | 20 | use tokio::task::JoinHandle; 21 | 22 | use crate::{Key, Registry, Span}; 23 | 24 | /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a 25 | /// [`JoinHandle`] for it. 26 | /// 27 | /// The spawned task will be registered in the current [`Registry`](crate::Registry) returned by 28 | /// [`Registry::try_current`] with the given [`Key`], if it exists. Otherwise, this is equivalent to 29 | /// [`tokio::spawn`]. 30 | pub fn spawn(key: impl Key, root_span: impl Into, future: T) -> JoinHandle 31 | where 32 | T: Future + Send + 'static, 33 | T::Output: Send + 'static, 34 | { 35 | if let Some(registry) = Registry::try_current() { 36 | tokio::spawn(registry.register(key, root_span).instrument(future)) 37 | } else { 38 | tokio::spawn(future) 39 | } 40 | } 41 | 42 | /// Spawns a new asynchronous task instrumented with the given root [`Span`], returning a 43 | /// [`JoinHandle`] for it. 44 | /// 45 | /// The spawned task will be registered in the current [`Registry`](crate::Registry) returned by 46 | /// [`Registry::try_current`] anonymously, if it exists. Otherwise, this is equivalent to 47 | /// [`tokio::spawn`]. 48 | pub fn spawn_anonymous(root_span: impl Into, future: T) -> JoinHandle 49 | where 50 | T: Future + Send + 'static, 51 | T::Output: Send + 'static, 52 | { 53 | if let Some(registry) = Registry::try_current() { 54 | tokio::spawn(registry.register_anonymous(root_span).instrument(future)) 55 | } else { 56 | tokio::spawn(future) 57 | } 58 | } 59 | -------------------------------------------------------------------------------- /src/tests.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | mod functionality; 16 | mod spawn; 17 | -------------------------------------------------------------------------------- /src/tests/functionality.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | use futures::future::{join_all, poll_fn, select_all}; 16 | use futures::{pin_mut, FutureExt, Stream, StreamExt}; 17 | use itertools::Itertools; 18 | 19 | use crate::root::current_context; 20 | use crate::{span, Config, InstrumentAwait, Registry}; 21 | 22 | async fn sleep(time: u64) { 23 | tokio::time::sleep(std::time::Duration::from_millis(time)).await; 24 | println!("slept {time}ms"); 25 | } 26 | 27 | async fn sleep_nested() { 28 | join_all([ 29 | sleep(1500).instrument_await("sleep nested 1500"), 30 | sleep(2500).instrument_await("sleep nested 2500"), 31 | ]) 32 | .await; 33 | } 34 | 35 | async fn multi_sleep() { 36 | sleep(400).await; 37 | 38 | sleep(800) 39 | .instrument_await("sleep another in multi sleep") 40 | .await; 41 | } 42 | 43 | fn stream1() -> impl Stream { 44 | use futures::stream::{iter, once}; 45 | 46 | iter(std::iter::repeat_with(|| { 47 | once(async { 48 | sleep(150).await; 49 | }) 50 | })) 51 | .flatten() 52 | } 53 | 54 | fn stream2() -> impl Stream { 55 | use futures::stream::{iter, once}; 56 | 57 | iter([ 58 | once(async { 59 | sleep(444).await; 60 | }) 61 | .boxed(), 62 | once(async { 63 | join_all([ 64 | sleep(400).instrument_await("sleep nested 400"), 65 | sleep(600).instrument_await("sleep nested 600"), 66 | ]) 67 | .await; 68 | }) 69 | .boxed(), 70 | ]) 71 | .flatten() 72 | } 73 | 74 | async fn hello() { 75 | async move { 76 | // Join 77 | join_all([ 78 | sleep(1000) 79 | .boxed() 80 | .instrument_await(span!("sleep {}", 1000)), 81 | sleep(2000).boxed().instrument_await("sleep 2000"), 82 | sleep_nested().boxed().instrument_await("sleep nested"), 83 | multi_sleep().boxed().instrument_await("multi sleep"), 84 | ]) 85 | .await; 86 | 87 | // Join another 88 | join_all([ 89 | sleep(1200).instrument_await("sleep 1200"), 90 | sleep(2200).instrument_await("sleep 2200"), 91 | ]) 92 | .await; 93 | 94 | // Cancel 95 | select_all([ 96 | sleep(666).boxed().instrument_await("sleep 666"), 97 | sleep_nested() 98 | .boxed() 99 | .instrument_await("sleep nested (should be cancelled)"), 100 | ]) 101 | .await; 102 | 103 | // Check whether cleaned up 104 | sleep(233).instrument_await("sleep 233").await; 105 | 106 | // Check stream next drop 107 | { 108 | let mut stream1 = stream1().fuse().boxed(); 109 | let mut stream2 = stream2().fuse().boxed(); 110 | let mut count = 0; 111 | 112 | 'outer: loop { 113 | tokio::select! { 114 | _ = stream1.next().instrument_await(span!("stream1 next {count}")) => {}, 115 | r = stream2.next().instrument_await(span!("stream2 next {count}")) => { 116 | if r.is_none() { break 'outer } 117 | }, 118 | } 119 | sleep(50) 120 | .instrument_await(span!("sleep before next stream poll: {count}")) 121 | .await; 122 | count += 1; 123 | } 124 | } 125 | 126 | // Check whether cleaned up 127 | sleep(233).instrument_await("sleep 233").await; 128 | 129 | // TODO: add tests on sending the future to another task or context. 130 | } 131 | .instrument_await("hello") 132 | .await; 133 | 134 | // Aborted futures have been cleaned up. There should only be a single active node of root. 135 | assert_eq!(current_context().unwrap().tree().active_node_count(), 1); 136 | } 137 | 138 | #[tokio::test] 139 | async fn test_await_tree() { 140 | let registry = Registry::new(Config::default()); 141 | let root = registry.register((), "actor 233"); 142 | 143 | let fut = root.instrument(hello()); 144 | pin_mut!(fut); 145 | 146 | let expected_counts = vec![ 147 | (1, 0), 148 | (8, 0), 149 | (9, 0), 150 | (8, 0), 151 | (6, 0), 152 | (5, 0), 153 | (4, 0), 154 | (4, 0), 155 | (3, 0), 156 | (6, 0), 157 | (3, 0), 158 | (4, 0), 159 | (3, 0), 160 | (4, 0), 161 | (3, 0), 162 | (4, 0), 163 | (3, 0), 164 | (6, 0), 165 | (5, 2), 166 | (6, 0), 167 | (5, 2), 168 | (6, 0), 169 | (5, 0), 170 | (4, 1), 171 | (5, 0), 172 | (3, 0), 173 | (3, 0), 174 | ]; 175 | let mut actual_counts = vec![]; 176 | 177 | poll_fn(|cx| { 178 | let tree = registry 179 | .collect::<()>() 180 | .into_iter() 181 | .exactly_one() 182 | .ok() 183 | .unwrap() 184 | .1; 185 | println!("{tree}"); 186 | actual_counts.push((tree.active_node_count(), tree.detached_node_count())); 187 | fut.poll_unpin(cx) 188 | }) 189 | .await; 190 | 191 | assert_eq!(actual_counts, expected_counts); 192 | } 193 | -------------------------------------------------------------------------------- /src/tests/spawn.rs: -------------------------------------------------------------------------------- 1 | // Copyright 2023 RisingWave Labs 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #![cfg(feature = "tokio")] 16 | 17 | use std::time::Duration; 18 | 19 | use futures::future::pending; 20 | use tokio::time::sleep; 21 | 22 | use crate::{Config, InstrumentAwait, Registry}; 23 | 24 | #[tokio::test] 25 | async fn main() { 26 | let registry = Registry::new(Config::default()); 27 | 28 | tokio::spawn(registry.register((), "root").instrument(async { 29 | crate::spawn_anonymous("child", async { 30 | crate::spawn_anonymous("grandson", async { 31 | pending::<()>().await; 32 | }) 33 | .instrument_await("wait for grandson") 34 | .await 35 | .unwrap() 36 | }) 37 | .instrument_await("wait for child") 38 | .await 39 | .unwrap() 40 | })); 41 | 42 | sleep(Duration::from_secs(1)).await; 43 | 44 | assert_eq!(registry.collect::<()>().len(), 1); 45 | assert_eq!(registry.collect_anonymous().len(), 2); 46 | assert_eq!(registry.collect_all().len(), 3); 47 | } 48 | --------------------------------------------------------------------------------