From d8fc96033eac8e3a518b97edd2cbe6f4b1e9d751 Mon Sep 17 00:00:00 2001 From: CatWithAHat Date: Sat, 29 Mar 2025 08:59:12 +0000 Subject: [PATCH] Upload files to "src" --- src/progress.rs | 77 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 src/progress.rs diff --git a/src/progress.rs b/src/progress.rs new file mode 100644 index 0000000..9382bbe --- /dev/null +++ b/src/progress.rs @@ -0,0 +1,77 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::time::sleep; +use tokio::sync::Mutex; +use std::io::{stdout, Write}; + +pub struct ProgressTracker { + pub lines_processed: AtomicUsize, + pub batches_written: AtomicUsize, + pub batches_queued: AtomicUsize, + pub total_lines: Option, +} + +impl ProgressTracker { + /// Create a new ProgressTracker. + + pub fn new(total_lines: Option) -> Self { + ProgressTracker { + lines_processed: AtomicUsize::new(0), + batches_written: AtomicUsize::new(0), + batches_queued: AtomicUsize::new(0), + total_lines, + } + } + + /// Increment the processed lines counter. + pub fn increment_lines(&self) { + self.lines_processed.fetch_add(1, Ordering::Relaxed); + } + + /// Increment the batches written counter. + pub fn increment_batches_written(&self) { + self.batches_written.fetch_add(1, Ordering::Relaxed); + } + + /// Update the number of batches currently queued. + pub fn set_batches_queued(&self, value: usize) { + self.batches_queued.store(value, Ordering::Relaxed); + } + + /// Returns a snapshot of current progress: + /// (lines processed, batches written, batches queued, percentage complete) + pub fn snapshot(&self) -> (usize, usize, usize, Option) { + let lines = self.lines_processed.load(Ordering::Relaxed); + let written = self.batches_written.load(Ordering::Relaxed); + let queued = self.batches_queued.load(Ordering::Relaxed); + let percent = self.total_lines.map(|total| { + if total > 0 { + (lines as f64 / total as f64) * 100.0 + } else { + 0.0 + } + }); + (lines, written, queued, percent) + } +} + +/// Spawns a background task that updates a single progress line every second. +pub fn start_progress_tracker(progress: Arc>) { + tokio::spawn(async move { + loop { + let (lines, written, queued, percent_opt) = progress.lock().await.snapshot(); + let percent_str = if let Some(percent) = percent_opt { + format!("{:.2}%", percent) + } else { + "N/A".to_string() + }; + print!( + "\rProcessed: {} lines | Batches Written: {} | Batches Queued: {} | % Complete: {}", + lines, written, queued, percent_str + ); + stdout().flush().unwrap(); + sleep(Duration::from_secs(1)).await; + } + }); +} \ No newline at end of file