Display download progress for self-updating and packages (#2031)

1 files changed, 211 insertions, 0 deletions

diff --git a/crates/typst-cli/src/download.rs b/crates/typst-cli/src/download.rs
new file mode 100644
index 00000000..416f8c68
--- /dev/null
+++ b/crates/typst-cli/src/download.rs
@@ -0,0 +1,211 @@
+use std::collections::VecDeque;
+use std::io::{self, ErrorKind, Read, Stderr, Write};
+use std::time::{Duration, Instant};
+
+use ureq::Response;
+
+// Acknowledgement:
+// Closely modelled after rustup's [`DownloadTracker`].
+// https://github.com/rust-lang/rustup/blob/master/src/cli/download_tracker.rs
+
+/// Keep track of this many download speed samples.
+const SPEED_SAMPLES: usize = 5;
+
+/// Download binary data and display its progress.
+#[allow(clippy::result_large_err)]
+pub fn download_with_progress(url: &str) -> Result<Vec<u8>, ureq::Error> {
+    let response = ureq::get(url).call()?;
+    Ok(RemoteReader::from_response(response).download()?)
+}
+
+/// A wrapper around [`ureq::Response`] that reads the response body in chunks
+/// over a websocket and displays statistics about its progress.
+///
+/// Downloads will _never_ fail due to statistics failing to print, print errors
+/// are silently ignored.
+struct RemoteReader {
+    reader: Box<dyn Read + Send + Sync + 'static>,
+    content_len: Option<usize>,
+    total_downloaded: usize,
+    downloaded_this_sec: usize,
+    downloaded_last_few_secs: VecDeque<usize>,
+    start_time: Instant,
+    last_print: Option<Instant>,
+    displayed_charcount: Option<usize>,
+    stderr: Stderr,
+}
+
+impl RemoteReader {
+    /// Wraps a [`ureq::Response`] and prepares it for downloading.
+    ///
+    /// The 'Content-Length' header is used as a size hint for read
+    /// optimization, if present.
+    pub fn from_response(response: Response) -> Self {
+        let content_len: Option<usize> = response
+            .header("Content-Length")
+            .and_then(|header| header.parse().ok());
+
+        Self {
+            reader: response.into_reader(),
+            content_len,
+            total_downloaded: 0,
+            downloaded_this_sec: 0,
+            downloaded_last_few_secs: VecDeque::with_capacity(SPEED_SAMPLES),
+            start_time: Instant::now(),
+            last_print: None,
+            displayed_charcount: None,
+            stderr: io::stderr(),
+        }
+    }
+
+    /// Download the bodies content as raw bytes while attempting to print
+    /// download statistics to standard error. Download progress gets displayed
+    /// and updated every second.
+    ///
+    /// These statistics will never prevent a download from completing, errors
+    /// are silently ignored.
+    pub fn download(mut self) -> io::Result<Vec<u8>> {
+        let mut buffer = vec![0; 8192];
+        let mut data = match self.content_len {
+            Some(content_len) => Vec::with_capacity(content_len),
+            None => Vec::with_capacity(8192),
+        };
+
+        loop {
+            let read = match self.reader.read(&mut buffer) {
+                Ok(0) => break,
+                Ok(n) => n,
+                // If the data is not yet ready but will be available eventually
+                // keep trying until we either get an actual error, receive data
+                // or an Ok(0).
+                Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
+                Err(e) => return Err(e),
+            };
+
+            data.extend(&buffer[..read]);
+
+            let last_printed = match self.last_print {
+                Some(prev) => prev,
+                None => {
+                    let current_time = Instant::now();
+                    self.last_print = Some(current_time);
+                    current_time
+                }
+            };
+            let elapsed = Instant::now().saturating_duration_since(last_printed);
+
+            self.total_downloaded += read;
+            self.downloaded_this_sec += read;
+
+            if elapsed >= Duration::from_secs(1) {
+                if self.downloaded_last_few_secs.len() == SPEED_SAMPLES {
+                    self.downloaded_last_few_secs.pop_back();
+                }
+
+                self.downloaded_last_few_secs.push_front(self.downloaded_this_sec);
+                self.downloaded_this_sec = 0;
+
+                if let Some(n) = self.displayed_charcount {
+                    self.erase_chars(n);
+                }
+
+                self.display();
+                let _ = write!(self.stderr, "\r");
+                self.last_print = Some(Instant::now());
+            }
+        }
+
+        self.display();
+        let _ = writeln!(self.stderr);
+
+        Ok(data)
+    }
+
+    /// Compile and format several download statistics and make an attempt at
+    /// displaying them on standard error.
+    fn display(&mut self) {
+        let sum: usize = self.downloaded_last_few_secs.iter().sum();
+        let len = self.downloaded_last_few_secs.len();
+        let speed = if len > 0 { sum / len } else { self.content_len.unwrap_or(0) };
+
+        let total = as_time_unit(self.total_downloaded, false);
+        let speed_h = as_time_unit(speed, true);
+        let elapsed =
+            time_suffix(Instant::now().saturating_duration_since(self.start_time));
+
+        let output = match self.content_len {
+            Some(content_len) => {
+                let percent = (self.total_downloaded as f64 / content_len as f64) * 100.;
+                let remaining = content_len - self.total_downloaded;
+
+                format!(
+                    "{} / {} ({:3.0} %) {} in {} ETA: {}",
+                    total,
+                    as_time_unit(content_len, false),
+                    percent,
+                    speed_h,
+                    elapsed,
+                    time_suffix(Duration::from_secs(if speed == 0 {
+                        0
+                    } else {
+                        (remaining / speed) as u64
+                    }))
+                )
+            }
+            None => format!("Total: {} Speed: {} Elapsed: {}", total, speed_h, elapsed,),
+        };
+
+        let _ = write!(self.stderr, "{output}");
+
+        self.displayed_charcount = Some(output.chars().count());
+    }
+
+    /// Erase each previously printed character and add a carriage return
+    /// character, clearing the line for the next `display()` update.
+    fn erase_chars(&mut self, count: usize) {
+        let _ = write!(self.stderr, "{}", " ".repeat(count));
+        let _ = write!(self.stderr, "\r");
+    }
+}
+
+/// Append a unit-of-time suffix.
+fn time_suffix(duration: Duration) -> String {
+    let secs = duration.as_secs();
+    match format_dhms(secs) {
+        (0, 0, 0, s) => format!("{s:2.0}s"),
+        (0, 0, m, s) => format!("{m:2.0}m {s:2.0}s"),
+        (0, h, m, s) => format!("{h:2.0}h {m:2.0}m {s:2.0}s"),
+        (d, h, m, s) => format!("{d:3.0}d {h:2.0}h {m:2.0}m {s:2.0}s"),
+    }
+}
+
+/// Format the total amount of seconds into the amount of days, hours, minutes
+/// and seconds.
+fn format_dhms(sec: u64) -> (u64, u8, u8, u8) {
+    let (mins, sec) = (sec / 60, (sec % 60) as u8);
+    let (hours, mins) = (mins / 60, (mins % 60) as u8);
+    let (days, hours) = (hours / 24, (hours % 24) as u8);
+    (days, hours, mins, sec)
+}
+
+/// Format a given size as a unit of time. Setting `include_suffix` to true
+/// appends a '/s' (per second) suffix.
+fn as_time_unit(size: usize, include_suffix: bool) -> String {
+    const KI: f64 = 1024.0;
+    const MI: f64 = KI * KI;
+    const GI: f64 = KI * KI * KI;
+
+    let size = size as f64;
+
+    let suffix = if include_suffix { "/s" } else { "" };
+
+    if size >= GI {
+        format!("{:5.1} GiB{}", size / GI, suffix)
+    } else if size >= MI {
+        format!("{:5.1} MiB{}", size / MI, suffix)
+    } else if size >= KI {
+        format!("{:5.1} KiB{}", size / KI, suffix)
+    } else {
+        format!("{size:3.0} B{}", suffix)
+    }
+}