use crate::ProgressCallback;
use std::io::Read;
use std::path::Path;
pub struct ProgressTracker<'a> {
progress: &'a mut dyn ProgressCallback,
current_entry: usize,
total_entries: usize,
}
impl<'a> ProgressTracker<'a> {
#[must_use]
pub fn new(progress: &'a mut dyn ProgressCallback, total_entries: usize) -> Self {
Self {
progress,
current_entry: 0,
total_entries,
}
}
pub fn on_entry_start(&mut self, path: &Path) {
self.current_entry += 1;
self.progress
.on_entry_start(path, self.total_entries, self.current_entry);
}
pub fn on_entry_complete(&mut self, path: &Path) {
self.progress.on_entry_complete(path);
}
pub fn on_complete(&mut self) {
self.progress.on_complete();
}
}
pub struct ProgressReader<'a, R> {
inner: R,
progress: &'a mut dyn ProgressCallback,
bytes_since_last_update: u64,
batch_threshold: u64,
}
impl<'a, R> ProgressReader<'a, R> {
#[must_use]
pub fn new(inner: R, progress: &'a mut dyn ProgressCallback) -> Self {
Self {
inner,
progress,
bytes_since_last_update: 0,
batch_threshold: 1024 * 1024, }
}
#[must_use]
pub fn with_batch_threshold(
inner: R,
progress: &'a mut dyn ProgressCallback,
batch_threshold: u64,
) -> Self {
Self {
inner,
progress,
bytes_since_last_update: 0,
batch_threshold,
}
}
pub fn flush_progress(&mut self) {
if self.bytes_since_last_update > 0 {
self.progress.on_bytes_written(self.bytes_since_last_update);
self.bytes_since_last_update = 0;
}
}
}
impl<R: Read> Read for ProgressReader<'_, R> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let bytes_read = self.inner.read(buf)?;
if bytes_read > 0 {
self.bytes_since_last_update += bytes_read as u64;
if self.bytes_since_last_update >= self.batch_threshold {
self.progress.on_bytes_written(self.bytes_since_last_update);
self.bytes_since_last_update = 0;
}
}
Ok(bytes_read)
}
}
impl<R> Drop for ProgressReader<'_, R> {
fn drop(&mut self) {
self.flush_progress();
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::unused_io_amount)]
mod tests {
use super::*;
use std::io::Cursor;
#[derive(Debug, Default)]
struct TestProgress {
entries_started: Vec<String>,
entries_completed: Vec<String>,
bytes_written: u64,
completed: bool,
}
impl ProgressCallback for TestProgress {
fn on_entry_start(&mut self, path: &Path, _total: usize, _current: usize) {
self.entries_started
.push(path.to_string_lossy().to_string());
}
fn on_bytes_written(&mut self, bytes: u64) {
self.bytes_written += bytes;
}
fn on_entry_complete(&mut self, path: &Path) {
self.entries_completed
.push(path.to_string_lossy().to_string());
}
fn on_complete(&mut self) {
self.completed = true;
}
}
#[test]
fn test_progress_tracker_entry_counting() {
let mut progress = TestProgress::default();
let mut tracker = ProgressTracker::new(&mut progress, 3);
tracker.on_entry_start(Path::new("file1.txt"));
tracker.on_entry_complete(Path::new("file1.txt"));
tracker.on_entry_start(Path::new("file2.txt"));
tracker.on_entry_complete(Path::new("file2.txt"));
tracker.on_complete();
assert_eq!(progress.entries_started.len(), 2);
assert_eq!(progress.entries_completed.len(), 2);
assert_eq!(progress.entries_started[0], "file1.txt");
assert_eq!(progress.entries_started[1], "file2.txt");
assert!(progress.completed);
}
#[test]
fn test_progress_reader_reports_bytes() {
let data = b"Hello, World!";
let reader = Cursor::new(data);
let mut progress = TestProgress::default();
let mut tracking_reader = ProgressReader::new(reader, &mut progress);
let mut buffer = vec![0u8; 5];
let bytes_read = tracking_reader.read(&mut buffer).unwrap();
drop(tracking_reader);
assert_eq!(bytes_read, 5);
assert_eq!(progress.bytes_written, 5);
assert_eq!(&buffer[..bytes_read], b"Hello");
}
#[test]
fn test_progress_reader_batching() {
let data = vec![0u8; 2 * 1024 * 1024]; let reader = Cursor::new(data);
let mut progress = TestProgress::default();
let batch_threshold = 64 * 1024; let mut tracking_reader =
ProgressReader::with_batch_threshold(reader, &mut progress, batch_threshold);
let mut buffer = vec![0u8; 32 * 1024];
for _ in 0..4 {
tracking_reader.read(&mut buffer).unwrap();
}
drop(tracking_reader);
assert!(progress.bytes_written > 0);
}
#[test]
fn test_progress_reader_handles_eof() {
let data = b"";
let reader = Cursor::new(data);
let mut progress = TestProgress::default();
let mut tracking_reader = ProgressReader::new(reader, &mut progress);
let mut buffer = vec![0u8; 10];
let bytes_read = tracking_reader.read(&mut buffer).unwrap();
drop(tracking_reader);
assert_eq!(bytes_read, 0);
assert_eq!(progress.bytes_written, 0);
}
#[test]
fn test_progress_reader_manual_flush() {
let data = b"test data";
let reader = Cursor::new(data);
let mut progress = TestProgress::default();
let mut buffer = vec![0u8; 4];
{
let mut tracking_reader = ProgressReader::new(reader, &mut progress);
tracking_reader.read(&mut buffer).unwrap();
tracking_reader.flush_progress();
tracking_reader.read(&mut buffer).unwrap();
tracking_reader.flush_progress();
}
assert_eq!(progress.bytes_written, 8); }
}