From 0c5bf64e62ed3b6a84fc4b25d6bdb7fc7f572520 Mon Sep 17 00:00:00 2001 From: spongecaptain Date: Sat, 8 Jul 2023 14:17:55 +0800 Subject: [PATCH 1/6] add time warn threshold and colored feature --- Cargo.toml | 1 + examples/color.rs | 58 +++++++++++++++++++++++++ src/context.rs | 105 +++++++++++++++++++++++++++------------------- src/registry.rs | 17 ++++++-- 4 files changed, 135 insertions(+), 46 deletions(-) create mode 100644 examples/color.rs 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..0073fe0 --- /dev/null +++ b/examples/color.rs @@ -0,0 +1,58 @@ +// 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 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..13ee3f6 100644 --- a/src/context.rs +++ b/src/context.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::fmt::{Debug, Write}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +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,54 @@ 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 +231,7 @@ 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..bdbdde4 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 coloring the terminal + pub colored : bool, + + /// if the time of execution is beyond it, warn it + 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,7 @@ 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); From ef1fcc37e86bc6789200eb4c9910af026d68110c Mon Sep 17 00:00:00 2001 From: spongecaptain Date: Sat, 8 Jul 2023 14:20:48 +0800 Subject: [PATCH 2/6] update comment --- .idea/workspace.xml | 145 ++++++++++++++++++++++++++++++++++++++++++++ examples/color.rs | 2 +- 2 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 .idea/workspace.xml diff --git a/.idea/workspace.xml b/.idea/workspace.xml new file mode 100644 index 0000000..0a605e4 --- /dev/null +++ b/.idea/workspace.xml @@ -0,0 +1,145 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + true + + \ No newline at end of file diff --git a/examples/color.rs b/examples/color.rs index 0073fe0..53dedc6 100644 --- a/examples/color.rs +++ b/examples/color.rs @@ -12,7 +12,7 @@ // 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. +//! This example shows how to configure warn threshold and color the warn use std::time::Duration; From ede78043e2441dff0a0027549efc963ac9ba2ce1 Mon Sep 17 00:00:00 2001 From: wathenjiang Date: Mon, 10 Jul 2023 14:04:16 +0800 Subject: [PATCH 3/6] rm irrelevant file --- .idea/workspace.xml | 145 -------------------------------------------- 1 file changed, 145 deletions(-) delete mode 100644 .idea/workspace.xml diff --git a/.idea/workspace.xml b/.idea/workspace.xml deleted file mode 100644 index 0a605e4..0000000 --- a/.idea/workspace.xml +++ /dev/null @@ -1,145 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - true - - \ No newline at end of file From 79756bae0deb80a638613467c07e44a27bef71ff Mon Sep 17 00:00:00 2001 From: wathen Date: Mon, 10 Jul 2023 14:09:07 +0800 Subject: [PATCH 4/6] update the comment of verbose field in Config Co-authored-by: Bugen Zhao --- src/registry.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/registry.rs b/src/registry.rs index bdbdde4..9c5ca74 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -32,7 +32,7 @@ pub struct Config { /// Whether to include the **verbose** span in the await-tree. pub verbose: bool, - /// Whether to coloring the terminal + /// Whether to enable coloring the terminal. pub colored : bool, /// if the time of execution is beyond it, warn it From 456f0d375d5b75204558e08aff57955639dcfd8a Mon Sep 17 00:00:00 2001 From: wathen Date: Mon, 10 Jul 2023 14:09:31 +0800 Subject: [PATCH 5/6] update the comment of warn_threshold field in Config Co-authored-by: Bugen Zhao --- src/registry.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/registry.rs b/src/registry.rs index 9c5ca74..f74712c 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -35,7 +35,7 @@ pub struct Config { /// Whether to enable coloring the terminal. pub colored : bool, - /// if the time of execution is beyond it, warn it + /// The threshold of the execution time to attach the warning sign (`!!!`) after the span name. pub warn_threshold : time::Duration, } From 7d4829572f738a2ce0111e7be934be522e8a1b89 Mon Sep 17 00:00:00 2001 From: wathenjiang Date: Mon, 10 Jul 2023 14:13:57 +0800 Subject: [PATCH 6/6] cargo fmt and clippy --- examples/color.rs | 17 ++++++++++------- src/context.rs | 22 +++++++++++----------- src/registry.rs | 13 +++++++++---- 3 files changed, 30 insertions(+), 22 deletions(-) diff --git a/examples/color.rs b/examples/color.rs index 53dedc6..55c18e8 100644 --- a/examples/color.rs +++ b/examples/color.rs @@ -14,22 +14,26 @@ //! 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 short_work() { + sleep(Duration::from_millis(500)) + .instrument_await("short") + .await } -async fn long_work(){ - sleep(Duration::from_millis(5000)).instrument_await("long").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{ + let mut registry = Registry::new(Config { verbose: true, colored: true, warn_threshold: Duration::from_millis(1000).into(), @@ -48,7 +52,6 @@ async fn main() { // short [105.606ms] println!("{tree}"); - sleep(Duration::from_millis(2000)).await; let tree = registry.get(&()).unwrap().to_string(); diff --git a/src/context.rs b/src/context.rs index 13ee3f6..d5d36b2 100644 --- a/src/context.rs +++ b/src/context.rs @@ -13,8 +13,8 @@ // limitations under the License. use std::fmt::{Debug, Write}; -use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; use std::time; use colored::Colorize; @@ -185,19 +185,15 @@ impl Tree { if depth == 0 { "".to_string() } else if elapsed.lt(&self.warn_threshold) { - format!(" [{:.3?}]", elapsed) + format!(" [{:.3?}]", elapsed) } else if self.colored { - format!(" [{:.3?}]", elapsed).red().to_string() + format!(" [{:.3?}]", elapsed).red().to_string() } else { - format!("!!! [{:.3?}]", elapsed) + format!("!!! [{:.3?}]", elapsed) } }; - write!( - f, - "{}", - elapsed_str - )?; + write!(f, "{}", elapsed_str)?; if depth > 0 && node == self.current { f.write_str(" <== current")?; @@ -213,7 +209,6 @@ impl Tree { Ok(()) } - } /// The task-local await-tree context. @@ -231,7 +226,12 @@ pub struct TreeContext { impl TreeContext { /// Create a new context. - pub(crate) fn new(root_span: Span, verbose: bool, colored: bool, warn_threshold: time::Duration) -> 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); diff --git a/src/registry.rs b/src/registry.rs index f74712c..2a9c945 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -33,10 +33,10 @@ pub struct Config { pub verbose: bool, /// Whether to enable coloring the terminal. - pub colored : bool, + pub colored: bool, /// The threshold of the execution time to attach the warning sign (`!!!`) after the span name. - pub warn_threshold : time::Duration, + pub warn_threshold: time::Duration, } #[allow(clippy::derivable_impls)] @@ -45,7 +45,7 @@ impl Default for Config { Self { verbose: false, colored: false, - warn_threshold : time::Duration::from_secs(10) + warn_threshold: time::Duration::from_secs(10), } } } @@ -91,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, self.config.colored, self.config.warn_threshold)); + 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);