risingwavelabs / await-tree

Generate accurate and informative tree dumps of asynchronous tasks.

await-tree

The Futures 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 Futures.

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 directory.

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(format!("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;
}

// 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}");

Compared to async-backtrace

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 for more details.

• await-tree maintains the tree structure with an arena-based data structure, 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, 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 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 for more information.

await-tree's Issues

Maybe add throttling for the "GC" operation while tracing new root futures?

I have noticed there is a TODO in Registery::register:

await-tree/src/registry.rs

Lines 79 to 81 in 334b563

	pub fn register(&mut self, key: K, root_span: impl Into<Span>) -> TreeRoot {
	// TODO: make this more efficient
	self.contexts.retain(|_, v| v.upgrade().is_some());

I guess this will make spawning an task become O(m) operation where m is the current running task. Perhaps we can add a new config like gc_throttle_duration: Option<Duration> (default None). Once it become Some, the registry will maintain the coarse instant of last GC, new register callls within this period again won't trigger GC anymore.

I have written a simple bench to test that spawns many background tasks(Even not really usual in real word, this case was usually used to show how coroutines are lighter that threads):

Details

    fn traced_multi_thread(size: usize) {
        let pool = tokio::runtime::Builder::new_multi_thread()
            .enable_time()
            .build()
            .unwrap();
        let futs = (0..size)
            .map(|_| {
                pool.spawn(root!(async {
                    tokio::time::sleep(std::time::Duration::from_millis(100)).await
                }))
            })
            .collect::<Vec<_>>();
        pool.block_on(futures::future::join_all(futs));
    }

Where root is a simple macro registers the current line to the register function. What it did like:

    pub fn root<T>(
        fut: impl Future<Output = T>,
    ) -> impl Future<Output = T> {
        let id = TID.fetch_add(1, SeqCst);
        REG.lock().unwrap().register(id, concat!(file!(), ":", line!(), ",", column!())).instrument(fut)
    }

The result shows comparing to the baseline with 10,000 tasks , there was a 3x performance regression, which is huge I think.

running 2 tests
test bench_async_trace::multi_thread::baseline_multi_thread_10000 ... bench: 101,856,955 ns/iter (+/- 345,245)
test bench_async_trace::multi_thread::traced_multi_thread_10000   ... bench: 308,596,986 ns/iter (+/- 3,973,652)

I haven't do further works now. If somewhere I'm wrong during benching or something else please let me know. If the new option is acceptable I'm glad to do further perfing to test whether the solution is effect and providing a patch.

How to add instrument for functions returns `Poll<T>`?

Hi, OpenDAL is currently developing an AwaitTreeLayer to give our users the ability to dump the execution tree at runtime: apache/opendal#2623

The missing part here is that how can we add instrument for functions returns Poll<T>? For example, AsyncRead will have API like the following:

pub trait AsyncRead {
    // Required method
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut [u8]
    ) -> Poll<Result<usize, Error>>;

Can we add instrument for every poll operations?