From e589f89922f634334ddf71853102c4271843efcc Mon Sep 17 00:00:00 2001
From: Dustin Carlino <dabreegster@gmail.com>
Date: Thu, 26 Nov 2020 16:12:16 -0800
Subject: [PATCH] Add a simple progress bar for updater downloads. #410

---
 Cargo.lock           |  1 +
 abstutil/src/time.rs |  4 ++--
 geom/src/percent.rs  |  6 +++++-
 updater/Cargo.toml   |  1 +
 updater/src/main.rs  | 31 ++++++++++++++++++++++---------
 5 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 09332730ab..81515666db 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3263,6 +3263,7 @@ version = "0.1.0"
 dependencies = [
  "abstutil",
  "flate2",
+ "geom",
  "md5",
  "reqwest",
  "tokio",
diff --git a/abstutil/src/time.rs b/abstutil/src/time.rs
index c04aae6153..31e64fa7fc 100644
--- a/abstutil/src/time.rs
+++ b/abstutil/src/time.rs
@@ -489,7 +489,7 @@ pub fn prettyprint_time(seconds: f64) -> String {
 }
 
 #[cfg(unix)]
-pub(crate) fn clear_current_line() {
+pub fn clear_current_line() {
     // Fails in the test runner.
     if let Ok((terminal_width, _)) = termion::terminal_size() {
         print!(
@@ -503,7 +503,7 @@ pub(crate) fn clear_current_line() {
 }
 
 #[cfg(not(unix))]
-pub(crate) fn clear_current_line() {
+pub fn clear_current_line() {
     print!("\r");
 }
 
diff --git a/geom/src/percent.rs b/geom/src/percent.rs
index 1e83a08f66..4d245eff5f 100644
--- a/geom/src/percent.rs
+++ b/geom/src/percent.rs
@@ -15,10 +15,14 @@ impl Percent {
         }
         Percent((x as f64) / 100.0)
     }
+
+    pub fn of(x: usize, total: usize) -> Percent {
+        Percent((x as f64) / (total as f64))
+    }
 }
 
 impl fmt::Display for Percent {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        writeln!(f, "{:.2}%", self.0 * 100.0)
+        write!(f, "{:.2}%", self.0 * 100.0)
     }
 }
diff --git a/updater/Cargo.toml b/updater/Cargo.toml
index 5179faeb03..7c5ae18adc 100644
--- a/updater/Cargo.toml
+++ b/updater/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2018"
 [dependencies]
 abstutil = { path = "../abstutil" }
 flate2 = "1.0.19"
+geom = { path = "../geom" }
 md5 = "0.7.0"
 reqwest = { version = "0.10.8", default-features=false, features=["rustls-tls"] }
 tokio = { version = "0.2.22", features = ["full"] }
diff --git a/updater/src/main.rs b/updater/src/main.rs
index e5e46a1a20..cc77b42fdb 100644
--- a/updater/src/main.rs
+++ b/updater/src/main.rs
@@ -6,7 +6,8 @@ use std::process::Command;
 
 use walkdir::WalkDir;
 
-use abstutil::{CmdArgs, DataPacks, Entry, Manifest, Timer};
+use abstutil::{prettyprint_usize, CmdArgs, DataPacks, Entry, Manifest, Timer};
+use geom::Percent;
 
 const MD5_BUF_READ_SIZE: usize = 4096;
 
@@ -168,7 +169,9 @@ fn generate_manifest() -> Manifest {
             continue;
         }
 
-        println!("> compute md5sum of {}", path);
+        abstutil::clear_current_line();
+        print!("> compute md5sum of {}", path);
+        std::io::stdout().flush().unwrap();
 
         // since these files can be very large, computes the md5 hash in chunks
         let mut file = File::open(&orig_path).unwrap();
@@ -191,6 +194,7 @@ fn generate_manifest() -> Manifest {
             },
         );
     }
+    println!();
     Manifest { entries: kv }
 }
 
@@ -228,23 +232,32 @@ async fn curl(version: &str, path: &str) -> Result<Vec<u8>, Box<dyn Error>> {
         "http://abstreet.s3-website.us-east-2.amazonaws.com/{}/{}.gz",
         version, path
     );
-
-    let mut bytes = Vec::new();
-
     println!("> download {}", src);
 
     let mut resp = reqwest::get(&src).await.unwrap();
-
     match resp.error_for_status_ref() {
         Ok(_) => {}
         Err(err) => {
             let err = format!("error getting {}: {}", src, err);
             return Err(err.into());
         }
-    };
-    while let Some(chunk) = resp.chunk().await.unwrap() {
-        bytes.write_all(&chunk).unwrap();
     }
 
+    let total_size = resp.content_length().map(|x| x as usize);
+    let mut bytes = Vec::new();
+    while let Some(chunk) = resp.chunk().await.unwrap() {
+        if let Some(n) = total_size {
+            abstutil::clear_current_line();
+            print!(
+                "{} ({} / {} bytes)",
+                Percent::of(bytes.len(), n),
+                prettyprint_usize(bytes.len()),
+                prettyprint_usize(n)
+            );
+        }
+
+        bytes.write_all(&chunk).unwrap();
+    }
+    println!();
     Ok(bytes)
 }