From 97e07920dc4da8df1b45cee0534f17dd3a864f84 Mon Sep 17 00:00:00 2001 From: Yoo1tic <137816438+Yoo1tic@users.noreply.github.com> Date: Sun, 10 Aug 2025 21:25:09 +0800 Subject: [PATCH] feat: add tracing and progress indicators for key validation process --- Cargo.lock | 212 +++++++++++++++++++++++++-- Cargo.toml | 2 + src/main.rs | 27 +++- src/validation/validation_service.rs | 28 +++- 4 files changed, 247 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0e9cd18..115456f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -76,6 +76,12 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + [[package]] name = "async-stream" version = "0.3.6" @@ -136,7 +142,7 @@ dependencies = [ "miniz_oxide", "object", "rustc-demangle", - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -236,6 +242,19 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "console" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e09ced7ebbccb63b4c65413d821f2e00ce54c5ca4514ddc6b3c892fdbcbc69d" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width 0.2.1", + "windows-sys 0.60.2", +] + [[package]] name = "displaydoc" version = "0.2.5" @@ -247,6 +266,12 @@ dependencies = [ "syn", ] +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + [[package]] name = "equivalent" version = "1.0.2" @@ -386,6 +411,7 @@ dependencies = [ "clap", "figment", "futures", + "indicatif", "mimalloc", "regex", "reqwest", @@ -395,6 +421,7 @@ dependencies = [ "tokio", "toml 0.9.5", "tracing", + "tracing-indicatif", "tracing-subscriber", "url", ] @@ -693,6 +720,20 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "indicatif" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70a646d946d06bedbbc4cac4c218acf4bbf2d87757a784857025f4d447e4e1cd" +dependencies = [ + "console", + "portable-atomic", + "unicode-width 0.2.1", + "unit-prefix", + "vt100", + "web-time", +] + [[package]] name = "inlinable_string" version = "0.1.15" @@ -910,6 +951,12 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "portable-atomic" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483" + [[package]] name = "potential_utf" version = "0.1.2" @@ -1626,6 +1673,18 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-indicatif" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1983afead46ff13a3c93581e0cec31d20b29efdd22cbdaa8b9f850eccf2c352" +dependencies = [ + "indicatif", + "tracing", + "tracing-core", + "tracing-subscriber", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -1676,6 +1735,24 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "unicode-width" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" + +[[package]] +name = "unit-prefix" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "323402cff2dd658f39ca17c789b502021b3f18707c91cdf22e3838e1b4023817" + [[package]] name = "untrusted" version = "0.9.0" @@ -1718,6 +1795,39 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vt100" +version = "0.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84cd863bf0db7e392ba3bd04994be3473491b31e66340672af5d11943c6274de" +dependencies = [ + "itoa", + "log", + "unicode-width 0.1.14", + "vte", +] + +[[package]] +name = "vte" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5022b5fbf9407086c180e9557be968742d839e68346af7792b8592489732197" +dependencies = [ + "arrayvec", + "utf8parse", + "vte_generate_state_changes", +] + +[[package]] +name = "vte_generate_state_changes" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e369bee1b05d510a7b4ed645f5faa90619e05437111783ea5848f28d97d3c2e" +dependencies = [ + "proc-macro2", + "quote", +] + [[package]] name = "want" version = "0.3.1" @@ -1864,13 +1974,19 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -1879,7 +1995,16 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.3", ] [[package]] @@ -1888,14 +2013,31 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.0", + "windows_aarch64_msvc 0.53.0", + "windows_i686_gnu 0.53.0", + "windows_i686_gnullvm 0.53.0", + "windows_i686_msvc 0.53.0", + "windows_x86_64_gnu 0.53.0", + "windows_x86_64_gnullvm 0.53.0", + "windows_x86_64_msvc 0.53.0", ] [[package]] @@ -1904,48 +2046,96 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c" + [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" +[[package]] +name = "windows_i686_gnu" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3" + [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11" + [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_i686_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" + [[package]] name = "winnow" version = "0.7.12" diff --git a/Cargo.toml b/Cargo.toml index 9c074a1..f3bd2d7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,3 +25,5 @@ toml = "0.9.5" thiserror = "2.0" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] } +tracing-indicatif = "0.3.12" +indicatif = "0.18.0" diff --git a/src/main.rs b/src/main.rs index 5fbcf75..3366948 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,10 @@ use gemini_keychecker::error::ValidatorError; use gemini_keychecker::{BANNER, config::KeyCheckerConfig, validation::start_validation}; use mimalloc::MiMalloc; use tracing::info; +use tracing_indicatif::IndicatifLayer; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -10,15 +14,24 @@ static GLOBAL: MiMalloc = MiMalloc; #[tokio::main] async fn main() -> Result<(), ValidatorError> { let config = KeyCheckerConfig::load_config()?; - // Initialize tracing with professional format using configured log level - tracing_subscriber::fmt() - .with_target(false) - .with_level(true) - .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(&config.log_level)), + + let indicatif_layer = IndicatifLayer::new(); + + tracing_subscriber::registry() + .with( + tracing_subscriber::fmt::layer() + .with_target(false) + .with_level(true) + .with_writer(indicatif_layer.get_stderr_writer()) + .with_ansi(false) + .with_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(&config.log_level)), + ), ) + .with(indicatif_layer) .init(); + // Display banner and configuration status at startup println!("{BANNER}"); info!("Configuration loaded: {}", config); diff --git a/src/validation/validation_service.rs b/src/validation/validation_service.rs index 21b8eea..e3363d1 100644 --- a/src/validation/validation_service.rs +++ b/src/validation/validation_service.rs @@ -6,10 +6,12 @@ use crate::types::GeminiKey; use crate::utils::client_builder; use async_stream::stream; use futures::{pin_mut, stream::StreamExt}; +use indicatif::ProgressStyle; use reqwest::Client; use std::time::Instant; use tokio::{fs, io::AsyncWriteExt, sync::mpsc}; -use tracing::error; +use tracing::{Span, error, info_span}; +use tracing_indicatif::span_ext::IndicatifSpanExt; pub struct ValidationService { config: KeyCheckerConfig, @@ -29,7 +31,7 @@ impl ValidationService { pub async fn validate_keys(&self, keys: Vec) -> Result<(), ValidatorError> { let start_time = Instant::now(); - + let total_keys = keys.len(); // Create channel for streaming keys from producer to consumer let (tx, mut rx) = mpsc::unbounded_channel::(); let stream = stream! { @@ -46,17 +48,32 @@ impl ValidationService { } } }); + // Create a progress bar to track validation progress + let progress_span = info_span!("progress_bar"); + progress_span.pb_set_style( + &ProgressStyle::with_template( + "[{bar:60.cyan/blue}] {pos}/{len} ({percent}%) [{elapsed_precise}] ETA:{eta}", + ) + .unwrap(), + ); + progress_span.pb_set_length(total_keys as u64); + progress_span.pb_set_message("Validating keys..."); + progress_span.pb_set_finish_message("All items processed"); + let progress_span_enter = progress_span.enter(); // Create stream to validate keys concurrently (two-stage pipeline) let cache_api_url = self.config.cache_api_url(); let valid_keys_stream = stream - .map(|key| { - test_generate_content_api( + .map(|key| async move { + let result = test_generate_content_api( self.client.clone(), self.full_url.clone(), key, self.config.clone(), ) + .await; + Span::current().pb_inc(1); + result }) .buffer_unordered(self.config.concurrency) .filter_map(|result| async { result.ok() }) @@ -93,6 +110,9 @@ impl ValidationService { free_buffer_writer.flush().await?; paid_buffer_writer.flush().await?; + std::mem::drop(progress_span_enter); + std::mem::drop(progress_span); + println!("Total Elapsed Time: {:?}", start_time.elapsed()); Ok(()) }