diff --git a/Cargo.toml b/Cargo.toml index 1e7ef40..101c4c7 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/examples/color.rs b/examples/color.rs new file mode 100644 index 0000000..55c18e8 --- /dev/null +++ b/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}"); +} diff --git a/src/context.rs b/src/context.rs index 2e628e9..d5d36b2 100644 --- a/src/context.rs +++ b/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, - 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, + 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, diff --git a/src/registry.rs b/src/registry.rs index 4d173a5..2a9c945 100644 --- a/src/registry.rs +++ b/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);