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