feat: configurable threshold and colorful output

Cargo.toml

@@ -19,6 +19,7 @@ parking_lot = "0.12"
 pin-project = "1"
 tokio = { version = "1", features = ["rt"] }
 tracing = "0.1"
+colored = "2"
 
 [dev-dependencies]
 criterion = { version = "0.4", features = ["async", "async_tokio"] }

examples/color.rs

@@ -0,0 +1,61 @@
+// 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 configure warn threshold and color the warn
+
+use std::time::Duration;
+
+use await_tree::{Config, InstrumentAwait, Registry};
+use tokio::time::sleep;
+
+async fn short_work() {
+    sleep(Duration::from_millis(500))
+        .instrument_await("short")
+        .await
+}
+
+async fn long_work() {
+    sleep(Duration::from_millis(5000))
+        .instrument_await("long")
+        .await
+}
+
+#[tokio::main]
+async fn main() {
+    let mut registry = Registry::new(Config {
+        verbose: true,
+        colored: true,
+        warn_threshold: Duration::from_millis(1000).into(),
+    });
+
+    let root = registry.register((), "work");
+    tokio::spawn(root.instrument(async {
+        short_work().await;
+        long_work().await;
+    }));
+
+    sleep(Duration::from_millis(100)).await;
+    let tree = registry.get(&()).unwrap().to_string();
+
+    // work
+    //   short [105.606ms]
+    println!("{tree}");
+
+    sleep(Duration::from_millis(2000)).await;
+    let tree = registry.get(&()).unwrap().to_string();
+
+    // work
+    //   long [1.609s]
+    println!("{tree}");
+}

src/context.rs

@@ -15,7 +15,9 @@
 use std::fmt::{Debug, Write};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
+use std::time;
 
+use colored::Colorize;
 use indextree::{Arena, NodeId};
 use itertools::Itertools;
 use parking_lot::{Mutex, MutexGuard};
@@ -57,50 +59,17 @@ pub struct Tree {
 
     /// The current span node. This is the node that is currently being polled.
     current: NodeId,
+
+    /// Whether to coloring the terminal
+    colored: bool,
+
+    /// if the time of execution is beyond it, warn it
+    warn_threshold: time::Duration,
 }
 
 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.as_str())?;
-
-            let elapsed: std::time::Duration = inner.start_time.elapsed().into();
-            write!(
-                f,
-                " [{}{:.3?}]",
-                if depth > 0 && 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)?;
+        self.fmt_node(f, &self.arena, self.root, 0)?;
 
         // Format all detached spans.
         for node in self.arena.iter().filter(|n| !n.is_removed()) {
@@ -110,7 +79,7 @@ impl std::fmt::Display for Tree {
             }
             if node.parent().is_none() {
                 writeln!(f, "[Detached {id}]")?;
-                fmt_node(f, &self.arena, id, 1, self.current)?;
+                self.fmt_node(f, &self.arena, id, 1)?;
             }
         }
 
@@ -197,6 +166,49 @@ impl Tree {
     pub(crate) fn current(&self) -> NodeId {
         self.current
     }
+
+    fn fmt_node(
+        &self,
+        f: &mut std::fmt::Formatter<'_>,
+        arena: &Arena<SpanNode>,
+        node: NodeId,
+        depth: usize,
+    ) -> std::fmt::Result {
+        f.write_str(&" ".repeat(depth * 2))?;
+
+        let inner = arena[node].get();
+        f.write_str(inner.span.as_str())?;
+
+        let elapsed: time::Duration = inner.start_time.elapsed().into();
+
+        let elapsed_str = {
+            if depth == 0 {
+                "".to_string()
+            } else if elapsed.lt(&self.warn_threshold) {
+                format!(" [{:.3?}]", elapsed)
+            } else if self.colored {
+                format!(" [{:.3?}]", elapsed).red().to_string()
+            } else {
+                format!("!!! [{:.3?}]", elapsed)
+            }
+        };
+
+        write!(f, "{}", elapsed_str)?;
+
+        if depth > 0 && node == self.current {
+            f.write_str("  <== current")?;
+        }
+
+        f.write_char('\n')?;
+        for child in node
+            .children(arena)
+            .sorted_by_key(|&id| arena[id].get().start_time)
+        {
+            self.fmt_node(f, arena, child, depth + 1)?;
+        }
+
+        Ok(())
+    }
 }
 
 /// The task-local await-tree context.
@@ -214,7 +226,12 @@ pub struct TreeContext {
 
 impl TreeContext {
     /// Create a new context.
-    pub(crate) fn new(root_span: Span, verbose: bool) -> Self {
+    pub(crate) fn new(
+        root_span: Span,
+        verbose: bool,
+        colored: bool,
+        warn_threshold: time::Duration,
+    ) -> Self {
         static ID: AtomicU64 = AtomicU64::new(0);
         let id = ID.fetch_add(1, Ordering::Relaxed);
 
@@ -225,6 +242,8 @@ impl TreeContext {
             id,
             verbose,
             tree: Tree {
+                colored,
+                warn_threshold,
                 arena,
                 root,
                 current: root,

src/registry.rs

@@ -17,6 +17,7 @@ use std::collections::HashMap;
 use std::future::Future;
 use std::hash::Hash;
 use std::sync::{Arc, Weak};
+use std::time;
 
 use derive_builder::Builder;
 
@@ -29,13 +30,23 @@ use crate::Span;
 #[builder(default)]
 pub struct Config {
     /// Whether to include the **verbose** span in the await-tree.
-    verbose: bool,
+    pub verbose: bool,
+
+    /// Whether to enable coloring the terminal.
+    pub colored: bool,
+
+    /// The threshold of the execution time to attach the warning sign (`!!!`) after the span name.
+    pub warn_threshold: time::Duration,
 }
 
 #[allow(clippy::derivable_impls)]
 impl Default for Config {
     fn default() -> Self {
-        Self { verbose: false }
+        Self {
+            verbose: false,
+            colored: false,
+            warn_threshold: time::Duration::from_secs(10),
+        }
     }
 }
 
@@ -80,7 +91,12 @@ where
         // TODO: make this more efficient
         self.contexts.retain(|_, v| v.upgrade().is_some());
 
-        let context = Arc::new(TreeContext::new(root_span.into(), self.config.verbose));
+        let context = Arc::new(TreeContext::new(
+            root_span.into(),
+            self.config.verbose,
+            self.config.colored,
+            self.config.warn_threshold,
+        ));
         let weak = Arc::downgrade(&context);
         self.contexts.insert(key, weak);