feat: add tracing support

Signed-off-by: Wenxuan Zhang <[email protected]>

Author: wfxr

Cargo.lock

@@ -19,10 +19,10 @@ name = "rlt"
 path = "src/lib.rs"
 
 [features]
-default = ["http", "rate_limit"]
-http = ["dep:http"]
-log = ["dep:log", "dep:tui-logger"]
+default = ["log", "rate_limit", "http"]
+log = ["dep:log", "dep:tracing", "dep:tui-logger"]
 rate_limit = ["dep:governor"]
+http = ["dep:http"]
 
 [dependencies]
 anyhow = "1"
@@ -41,14 +41,18 @@ async-trait = "0.1"
 hdrhistogram = "7.5"
 itertools = "0.12"
 tabled = "0.15"
-http = { version = "1.1", optional = true }
-tui-logger = { version = "0.11", optional = true }
-log = { version = "0.4", optional = true }
 cfg-if = "1"
 parking_lot = "0.12"
-governor = { version = "0.6", optional = true }
 nonzero_ext = "0.3"
 
+http = { version = "1.1", optional = true }
+governor = { version = "0.6", optional = true }
+log = { version = "0.4", optional = true }
+tracing = { version = "0.1", optional = true }
+tui-logger = { version = "0.11", optional = true, features = [
+    "tracing-support",
+] }
+
 [dev-dependencies]
 tokio = { version = "1.36", features = ["rt-multi-thread"] }
 bytes = "1.6"
@@ -63,3 +67,4 @@ hyper-util = { version = "0.1", features = [
 ] }
 reqwest = "0.12"
 tokio-postgres = "0.7"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }

Cargo.toml

@@ -0,0 +1,57 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use clap::Parser;
+use rlt::{cli::BenchCli, IterInfo, IterReport, StatelessBenchSuite, Status, StatusKind};
+use tokio::time::{Duration, Instant};
+
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
+
+#[derive(Clone)]
+struct SimpleBench;
+
+#[async_trait]
+impl StatelessBenchSuite for SimpleBench {
+    async fn bench(&mut self, info: &IterInfo) -> Result<IterReport> {
+        let t = Instant::now();
+
+        // simulate some work
+        tokio::time::sleep(Duration::from_micros(info.runner_seq % 30)).await;
+        let duration = t.elapsed();
+
+        // simulate status code
+        let seq = info.runner_seq;
+        let status = match seq % 10 {
+            8..=10 => Status::server_error(500),
+            6..=7 => Status::client_error(400),
+            _ => Status::success(200),
+        };
+
+        match status.kind() {
+            StatusKind::Success => tracing::info!(?status, seq),
+            StatusKind::ClientError => tracing::warn!(?status, seq),
+            StatusKind::ServerError | StatusKind::Error => tracing::error!(?status, seq),
+        };
+
+        Ok(IterReport { duration, status, bytes: 0, items: 1 })
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let opt = BenchCli::parse();
+    match opt.collector() {
+        rlt::cli::Collector::Tui => {
+            tracing_subscriber::registry()
+                .with(EnvFilter::from_default_env())
+                .with(rlt::tui_tracing_subscriber_layer())
+                .init();
+        }
+        rlt::cli::Collector::Silent => {
+            tracing_subscriber::fmt()
+                .with_env_filter(EnvFilter::from_default_env())
+                .init();
+        }
+    }
+
+    rlt::cli::run(opt, SimpleBench).await
+}

examples/logging.rs

@@ -29,12 +29,6 @@ impl StatelessBenchSuite for SimpleBench {
         // simulate items processed in current iteration
         let items = info.worker_seq % 100;
 
-        #[cfg(feature = "log")]
-        // press `l` to see log output
-        if status.kind() != rlt::StatusKind::Success {
-            log::warn!("duration: {:?}, status: {}, items: {}", duration, status, items);
-        }
-
         Ok(IterReport { duration, status, bytes: items * 1024, items })
     }
 }

examples/simple_stateless.rs

@@ -166,7 +166,8 @@ impl BenchCli {
         }
     }
 
-    pub(crate) fn collector(&self) -> Collector {
+    /// Get the actual collector type.
+    pub fn collector(&self) -> Collector {
         match self.collector {
             Some(collector) => collector,
             None if self.quiet || !stdout().is_tty() => Collector::Silent,

src/cli.rs

@@ -653,7 +653,6 @@ mod tui_log {
     use super::*;
 
     use log::LevelFilter;
-    use std::str::FromStr;
     use tui_logger::{TuiLoggerLevelOutput, TuiLoggerSmartWidget, TuiWidgetState};
 
     pub(crate) struct LogState {
@@ -663,13 +662,8 @@ mod tui_log {
 
     impl LogState {
         pub(crate) fn from_env() -> Result<Self> {
-            let log_level = match std::env::var("RUST_LOG") {
-                Ok(log_level) => LevelFilter::from_str(&log_level).unwrap_or(LevelFilter::Info),
-                Err(_) => LevelFilter::Info,
-            };
-            tui_logger::init_logger(log_level).map_err(|e| anyhow::anyhow!(e))?;
-            tui_logger::set_default_level(log_level);
-            let state = TuiWidgetState::new().set_default_display_level(log_level);
+            tui_logger::set_default_level(LevelFilter::Trace);
+            let state = TuiWidgetState::new().set_default_display_level(LevelFilter::Info);
             Ok(Self { inner: state, display: false })
         }
     }

src/collector/tui.rs

@@ -72,3 +72,6 @@ pub use crate::{
     runner::{BenchSuite, StatelessBenchSuite},
     status::{Status, StatusKind},
 };
+
+#[cfg(feature = "log")]
+pub use tui_logger::tracing_subscriber_layer as tui_tracing_subscriber_layer;