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 <[email protected]>

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,15 +70,67 @@ 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),
             ),
         )
         .unwrap()
     });
     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()
+    }
+}