Add Progress Bar For Large Streamed Downloads (#6096)

* add progress bar

* fix progress bar

* fix

* remove unused lifetime

* lint

* update to latest indicatif

---------

Co-authored-by: Paweł Buchowski <pawel.buchowski@enso.org>

Cargo.lock

@@ -4425,13 +4425,15 @@ dependencies = [
 
 [[package]]
 name = "indicatif"
-version = "0.17.3"
+version = "0.17.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cef509aa9bc73864d6756f0d34d35504af3cf0844373afe9b8669a5b8005a729"
+checksum = "8ff8cc23a7393a397ed1d7f56e6365cba772aba9f9912ab968b03043c395d057"
 dependencies = [
  "console",
+ "instant",
  "number_prefix",
  "portable-atomic",
+ "tokio",
  "unicode-width",
 ]
 
@@ -5479,9 +5481,9 @@ checksum = "f6519412c9e0d4be579b9f0618364d19cb434b324fc6ddb1b27b1e682c7105ed"
 
 [[package]]
 name = "portable-atomic"
-version = "0.3.19"
+version = "1.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "26f6a7b87c2e435a3241addceeeff740ff8b7e76b74c13bf9acb17fa454ea00b"
+checksum = "767eb9f07d4a5ebcb39bbf2d452058a93c011373abf6832e24194a1c3f004794"
 
 [[package]]
 name = "ppv-lite86"

build/ci_utils/Cargo.toml

@@ -35,7 +35,7 @@ headers = "0.3.7"
 heck = "0.4.0"
 http-serde = "1.1.0"
 indexmap = "1.7.0"
-indicatif = "0.17.1"
+indicatif = { version = "0.17.1", features = ["tokio"] }
 itertools = { workspace = true }
 lazy_static = { workspace = true }
 log = "0.4.14"

build/ci_utils/src/global.rs

@@ -32,6 +32,10 @@ const REFRESHES_PER_SECOND: u32 = 100;
 #[derive(derivative::Derivative)]
 #[derivative(Debug)]
 struct GlobalState {
+    /// A globally-shared reference to the multi-progress bar.
+    ///
+    /// All progress bars must be added to this multi-progress bar. This ensures that the progress
+    /// bars are displayed in a way that does not interfere with tracing log output.
     mp:            MultiProgress,
     #[derivative(Debug = "ignore")]
     bars:          Vec<WeakProgressBar>,
@@ -72,6 +76,11 @@ impl Default for GlobalState {
 
 static GLOBAL: LazyLock<Mutex<GlobalState>> = LazyLock::new(default);
 
+/// Returns a reference to the global multi-progress bar.
+pub fn multi_progress_bar() -> MultiProgress {
+    GLOBAL.lock().unwrap().mp.clone()
+}
+
 pub fn progress_bar(f: impl FnOnce() -> ProgressBar) -> ProgressBar {
     let ret = f();
     let ret = GLOBAL.lock().unwrap().mp.add(ret);

build/ci_utils/src/io/web.rs

@@ -2,6 +2,7 @@ use crate::prelude::*;
 
 use crate::fs::tokio::copy_to_file;
 use crate::fs::tokio::create_parent_dir_if_missing;
+use crate::global::progress_bar;
 
 use anyhow::Context;
 use reqwest::Client;
@@ -67,12 +68,30 @@ pub async fn download_file(url: impl IntoUrl, output: impl AsRef<Path>) -> Resul
 ), err)]
 pub async fn stream_response_to_file(response: Response, output: impl AsRef<Path>) -> Result {
     trace!("Streamed response: {:#?}", response);
+    let bar = response_progress_bar(&response);
     let response = handle_error_response(response).await?;
     let reader = async_reader(response);
+    let reader = &mut bar.wrap_async_read(reader);
     copy_to_file(reader, output).await?;
     Ok(())
 }
 
+/// Creates a progress bar for the response, with the length from the `Content-Length` header.
+///
+/// The progress bar provides a visual indication of the download progress. Visual indication is
+/// helpful for large files, as the build script does not print any output for a long time. Without
+/// the progress bar, the user might think that the build script is stuck.
+pub fn response_progress_bar(response: &Response) -> indicatif::ProgressBar {
+    let url = response.url().to_string();
+    let len = response.content_length();
+    let draw_target = indicatif::ProgressDrawTarget::stderr();
+    let bar = progress_bar(|| indicatif::ProgressBar::with_draw_target(len, draw_target));
+    let style = indicatif::ProgressStyle::default_bar();
+    bar.set_message(format!("Streaming response to file {url}."));
+    bar.set_style(style);
+    bar
+}
+
 pub fn async_reader(response: Response) -> impl AsyncBufRead + Unpin {
     tokio_util::io::StreamReader::new(response.bytes_stream().map_err(std::io::Error::other))
 }

build/ci_utils/src/log.rs

@@ -1,6 +1,8 @@
 use crate::prelude::*;
 use tracing_subscriber::prelude::*;
 
+use crate::global;
+use std::io;
 use std::sync::Once;
 use tracing::span::Attributes;
 use tracing::subscriber::Interest;
@@ -68,11 +70,15 @@ pub fn setup_logging() -> Result {
             .with_default_directive(LevelFilter::TRACE.into())
             .from_env_lossy();
 
+        let progress_bar = global::multi_progress_bar();
+        let progress_bar_writer = IndicatifWriter::new(progress_bar);
+
         tracing::subscriber::set_global_default(
             Registry::default().with(MyLayer).with(
                 tracing_subscriber::fmt::layer()
                     .without_time()
                     .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
+                    .with_writer(progress_bar_writer)
                     .with_filter(filter),
             ),
         )
@@ -80,3 +86,51 @@ pub fn setup_logging() -> Result {
     });
     Ok(())
 }
+
+
+
+// =======================
+// === IndicatifWriter ===
+// =======================
+
+/// A writer that writes to stderr, but suspends any active progress bars while doing so.
+///
+/// Progress bars use `stderr` to draw themselves. If we log to `stderr` while there is a progress
+/// bar visible, the progress bar will be overwritten by the log message. This results in the
+/// previous states of the progress bar remaining visible on the screen, which is not desirable.
+///
+/// To avoid this, the writer suspends the progress bars when writing logs.
+#[derive(Clone, Debug)]
+struct IndicatifWriter {
+    progress_bar: indicatif::MultiProgress,
+}
+
+
+// === Main `impl` ===
+
+impl IndicatifWriter {
+    pub fn new(progress_bar: indicatif::MultiProgress) -> Self {
+        Self { progress_bar }
+    }
+}
+
+
+// === Trait `impl`s ===
+
+impl std::io::Write for IndicatifWriter {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.progress_bar.suspend(|| io::stderr().write(buf))
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.progress_bar.suspend(|| io::stderr().flush())
+    }
+}
+
+impl tracing_subscriber::fmt::MakeWriter<'_> for IndicatifWriter {
+    type Writer = Self;
+
+    fn make_writer(&self) -> Self::Writer {
+        self.clone()
+    }
+}