replay.rs

#!/usr/bin/env rust-script
//! ```cargo
//! [dependencies]
//! indicatif="0.17.3"
//! ```

use std::env;
use std::fs::File;
use std::io::{self, BufRead}; // BufRead provides .lines() trait for BufReader
use std::path::Path;
use std::thread;
use std::time::{Duration, Instant};
use std::collections::VecDeque;

use indicatif::{ProgressBar, ProgressStyle, HumanBytes};
use indicatif::style::ProgressTracker;

fn read_file<P>(filepath: P, pb: Option<ProgressBar>) -> u64
where
    P: AsRef<Path>,
{
    let mut file_size: u64 = 0;
    let mut last_read_size: u64 = 0;
    let mut time = Instant::now();
    let file = File::open(filepath).unwrap();
    for line in io::BufReader::new(file).lines() {
        let line = line.unwrap();
        let mut line_iter = line.split(" ");
        let delay: u64 = line_iter.next().unwrap().parse().unwrap();
        time += Duration::from_micros(delay);
        if let Some(read_size_str) = line_iter.next() {
            last_read_size = read_size_str.parse().unwrap();
        }
        file_size += last_read_size;
        if let Some(ref progress_bar) = pb {
            thread::sleep(time - Instant::now());
            progress_bar.set_position(file_size);
        }
    }
    file_size
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 || args.len() > 3 {
        println!("Usage: indicatif-replay <logfile> [file size]");
        return;
    }

    let file_size: u64;

    if args.len() == 3 {
        file_size = args[2].parse().unwrap();
    } else {
        file_size = read_file(&args[1], None);
        println!("Automatically detected file size: {}", file_size);
    }

    let pb = ProgressBar::new(file_size);
    pb.set_style(
        ProgressStyle::default_bar()
//            .template("[{elapsed_precise}] {percent}% [{wide_bar}] {bytes} ({bytes_per_sec}) eta {eta}")
            .template("[{elapsed_precise}] {percent}% [{wide_bar}] {bytes} ({smoothed_per_sec}) eta {eta}")
            .unwrap()
            .with_key("smoothed_per_sec", MovingAvgRate::default())
            .progress_chars("#>-"),
    );
    pb.enable_steady_tick(Duration::from_millis(250));

    read_file(&args[1], Some(pb));
}


#[derive(Clone, Default)]
struct MovingAvgRate {
    samples: VecDeque<(std::time::Instant, u64)>,
}

impl ProgressTracker for MovingAvgRate {
    fn clone_box(&self) -> Box<dyn ProgressTracker> {
        Box::new(self.clone())
    }

    fn tick(&mut self, state: &indicatif::ProgressState, now: std::time::Instant) {
        // sample at most every 20ms
        if self
            .samples
            .back()
            .map_or(true, |(prev, _)| (now - *prev) > Duration::from_millis(20))
        {
            self.samples.push_back((now, state.pos()));
        }

        while let Some(first) = self.samples.front() {
            if now - first.0 > Duration::from_secs(4) {
                self.samples.pop_front();
            } else {
                break;
            }
        }
    }

    fn reset(&mut self, _state: &indicatif::ProgressState, _now: std::time::Instant) {
        self.samples = Default::default();
    }

    fn write(&self, _state: &indicatif::ProgressState, w: &mut dyn std::fmt::Write) {
        match (self.samples.front(), self.samples.back()) {
            (Some((t0, p0)), Some((t1, p1))) if self.samples.len() > 1 => {
                let elapsed = (*t1 - *t0).as_secs_f64();
                let rate = (p1 - p0) as f64 / elapsed;

                write!(w, "{}/s", HumanBytes(rate as u64)).unwrap()
            }
            _ => write!(w, "-").unwrap(),
        }
    }
}