await_tree/
lib.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Generate accurate and informative tree dumps of asynchronous tasks.
16//!
17//! # Example
18//!
19//! Below is a basic example of how to trace asynchronous tasks with the global registry of the
20//! `await-tree` crate.
21//!
22//! ```rust
23//! # use std::time::Duration;
24//! # use tokio::time::sleep;
25//! # use await_tree::{InstrumentAwait, Registry, span};
26//! # use futures::future::join;
27//! #
28//! # async fn work() { futures::future::pending::<()>().await }
29//! #
30//! async fn bar(i: i32) {
31//!     // `&'static str` span
32//!     baz(i).instrument_await("baz in bar").await
33//! }
34//!
35//! async fn baz(i: i32) {
36//!     // runtime `String` span is also supported
37//!     work().instrument_await(span!("working in baz {i}")).await
38//! }
39//!
40//! async fn foo() {
41//!     // spans of joined futures will be siblings in the tree
42//!     join(
43//!         bar(3).instrument_await("bar"),
44//!         baz(2).instrument_await("baz"),
45//!     )
46//!     .await;
47//! }
48//!
49//! # #[tokio::main]
50//! # async fn main() {
51//! // Init the global registry to start tracing the tasks.
52//! await_tree::init_global_registry(Default::default());
53//! // Spawn a task with root span "foo" and key "foo".
54//! await_tree::spawn("foo", "foo", foo());
55//! // Let the tasks run for a while.
56//! sleep(Duration::from_secs(1)).await;
57//! // Get the tree of the task with key "foo".
58//! let tree = Registry::current().get("foo").unwrap();
59//!
60//! // foo [1.006s]
61//! //   bar [1.006s]
62//! //     baz in bar [1.006s]
63//! //       working in baz 3 [1.006s]
64//! //   baz [1.006s]
65//! //     working in baz 2 [1.006s]
66//! println!("{tree}");
67//! # }
68//! ```
69
70#![forbid(missing_docs)]
71
72use std::future::Future;
73
74mod context;
75mod future;
76mod global;
77mod obj_utils;
78mod registry;
79mod root;
80mod span;
81#[cfg(feature = "tokio")]
82mod spawn;
83
84pub use context::{current_tree, Tree};
85pub use future::Instrumented;
86pub use global::init_global_registry;
87pub use registry::{AnyKey, Config, ConfigBuilder, ConfigBuilderError, Key, Registry};
88pub use root::TreeRoot;
89pub use span::{Span, SpanExt};
90#[cfg(feature = "tokio")]
91pub use spawn::{spawn, spawn_anonymous};
92
93#[doc(hidden)]
94pub mod __private {
95    pub use crate::span::fmt_span;
96}
97
98/// Attach spans to a future to be traced in the await-tree.
99pub trait InstrumentAwait: Future + Sized {
100    /// Instrument the future with a span.
101    fn instrument_await(self, span: impl Into<Span>) -> Instrumented<Self> {
102        Instrumented::new(self, span.into())
103    }
104}
105impl<F> InstrumentAwait for F where F: Future {}
106
107#[cfg(test)]
108mod tests;
await_tree/lib.rs

await_tree/
lib.rs
