Concurrently download components of a toolchain

src/cli/download_tracker.rs

@@ -1,5 +1,5 @@
 use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
-use std::time::Duration;
+use std::collections::HashMap;
 
 use crate::dist::Notification as In;
 use crate::notifications::Notification;
@@ -13,8 +13,8 @@ use crate::utils::Notification as Un;
 pub(crate) struct DownloadTracker {
     /// MultiProgress bar for the downloads.
     multi_progress_bars: MultiProgress,
-    /// ProgressBar for the current download.
-    progress_bar: ProgressBar,
+    /// Mapping of files to their corresponding progress bars.
+    file_progress_bars: HashMap<String, ProgressBar>,
 }
 
 impl DownloadTracker {
@@ -28,22 +28,29 @@ impl DownloadTracker {
 
         Self {
             multi_progress_bars,
-            progress_bar: ProgressBar::hidden(),
+            file_progress_bars: HashMap::new(),
         }
     }
 
     pub(crate) fn handle_notification(&mut self, n: &Notification<'_>) -> bool {
         match *n {
-            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
-                self.content_length_received(content_len);
+            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(
+                content_len,
+                file,
+            ))) => {
+                self.content_length_received(content_len, file);
                 true
             }
-            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
-                self.data_received(data.len());
+            Notification::Install(In::Utils(Un::DownloadDataReceived(data, file))) => {
+                self.data_received(data.len(), file);
                 true
             }
-            Notification::Install(In::Utils(Un::DownloadFinished)) => {
-                self.download_finished();
+            Notification::Install(In::Utils(Un::DownloadFinished(file))) => {
+                self.download_finished(file);
+                true
+            }
+            Notification::Install(In::DownloadingComponent(component, _, _)) => {
+                self.create_progress_bar(component);
                 true
             }
             Notification::Install(In::Utils(Un::DownloadPushUnit(_))) => true,
@@ -53,30 +60,60 @@ impl DownloadTracker {
         }
     }
 
-    /// Sets the length for a new ProgressBar and gives it a style.
-    pub(crate) fn content_length_received(&mut self, content_len: u64) {
-        self.progress_bar.set_length(content_len);
-        self.progress_bar.set_style(
+    /// Helper function to find the progress bar for a given file.
+    fn progress_bar(&mut self, file: &str) -> Option<&ProgressBar> {
+        // During the installation this function can be called with an empty file/URL.
+        if file.is_empty() {
+            return None;
+        }
+        let component = self
+            .file_progress_bars
+            .keys()
+            .find(|comp| file.contains(*comp))?;
+
+        self.file_progress_bars.get(component)
+    }
+
+    /// Creates a new ProgressBar for the given component.
+    pub(crate) fn create_progress_bar(&mut self, component: &str) {
+        let pb = ProgressBar::hidden();
+        pb.set_style(
             ProgressStyle::with_template(
-                "[{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, ETA: {eta})",
+                "{msg:>12.bold}  [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, ETA: {eta})",
             )
             .unwrap()
             .progress_chars("## "),
         );
+        pb.set_message(component.to_string());
+        self.multi_progress_bars.add(pb.clone());
+        self.file_progress_bars.insert(component.to_string(), pb);
+    }
+
+    /// Sets the length for a new ProgressBar and gives it a style.
+    pub(crate) fn content_length_received(&mut self, content_len: u64, file: &str) {
+        if let Some(pb) = self.progress_bar(file) {
+            pb.set_length(content_len);
+        }
     }
 
     /// Notifies self that data of size `len` has been received.
-    pub(crate) fn data_received(&mut self, len: usize) {
-        if self.progress_bar.is_hidden() && self.progress_bar.elapsed() >= Duration::from_secs(1) {
-            self.multi_progress_bars.add(self.progress_bar.clone());
+    pub(crate) fn data_received(&mut self, len: usize, file: &str) {
+        if let Some(pb) = self.progress_bar(file) {
+            pb.inc(len as u64);
         }
-        self.progress_bar.inc(len as u64);
     }
 
     /// Notifies self that the download has finished.
-    pub(crate) fn download_finished(&mut self) {
-        self.progress_bar.finish_and_clear();
-        self.multi_progress_bars.remove(&self.progress_bar);
-        self.progress_bar = ProgressBar::hidden();
+    pub(crate) fn download_finished(&mut self, file: &str) {
+        if let Some(pb) = self.progress_bar(file) {
+            pb.set_style(
+                ProgressStyle::with_template(
+                    "{msg:>12.bold}  downloaded {total_bytes} in {elapsed}",
+                )
+                .unwrap(),
+            );
+            let msg = pb.message();
+            pb.finish_with_message(msg);
+        }
     }
 }

src/cli/self_update/windows.rs

@@ -274,7 +274,10 @@ pub(crate) async fn try_install_msvc(
     let download_tracker = Arc::new(Mutex::new(DownloadTracker::new_with_display_progress(
         true, process,
     )));
-    download_tracker.lock().unwrap().download_finished();
+    download_tracker
+        .lock()
+        .unwrap()
+        .download_finished(visual_studio_url.as_str());
 
     info!("downloading Visual Studio installer");
     download_file(

src/diskio/threaded.rs

@@ -263,10 +263,11 @@ impl Executor for Threaded<'_> {
         // pretend to have bytes to deliver.
         let mut prev_files = self.n_files.load(Ordering::Relaxed);
         if let Some(handler) = self.notify_handler {
-            handler(Notification::DownloadFinished);
+            handler(Notification::DownloadFinished(""));
             handler(Notification::DownloadPushUnit(Unit::IO));
             handler(Notification::DownloadContentLengthReceived(
                 prev_files as u64,
+                "",
             ));
         }
         if prev_files > 50 {
@@ -284,12 +285,12 @@ impl Executor for Threaded<'_> {
             current_files = self.n_files.load(Ordering::Relaxed);
             let step_count = prev_files - current_files;
             if let Some(handler) = self.notify_handler {
-                handler(Notification::DownloadDataReceived(&buf[0..step_count]));
+                handler(Notification::DownloadDataReceived(&buf[0..step_count], ""));
             }
         }
         self.pool.join();
         if let Some(handler) = self.notify_handler {
-            handler(Notification::DownloadFinished);
+            handler(Notification::DownloadFinished(""));
             handler(Notification::DownloadPopUnit);
         }
         // close the feedback channel so that blocking reads on it can

src/dist/manifestation.rs

@@ -6,8 +6,10 @@ mod tests;
 
 use std::path::Path;
 
-use anyhow::{Context, Result, anyhow, bail};
+use anyhow::{Context, Error, Result, anyhow, bail};
+use futures_util::stream::StreamExt;
 use tokio_retry::{RetryIf, strategy::FixedInterval};
+use tracing::info;
 
 use crate::dist::component::{
     Components, Package, TarGzPackage, TarXzPackage, TarZStdPackage, Transaction,
@@ -153,6 +155,7 @@ impl Manifestation {
         let mut things_to_install: Vec<(Component, CompressionKind, File)> = Vec::new();
         let mut things_downloaded: Vec<String> = Vec::new();
         let components = update.components_urls_and_hashes(new_manifest)?;
+        let components_len = components.len();
 
         const DEFAULT_MAX_RETRIES: usize = 3;
         let max_retries: usize = download_cfg
@@ -162,41 +165,60 @@ impl Manifestation {
             .and_then(|s| s.parse().ok())
             .unwrap_or(DEFAULT_MAX_RETRIES);
 
-        for (component, format, url, hash) in components {
+        info!("downloading component(s)");
+        for (component, _, _, _) in components.clone() {
             (download_cfg.notify_handler)(Notification::DownloadingComponent(
                 &component.short_name(new_manifest),
                 &self.target_triple,
                 component.target.as_ref(),
             ));
-            let url = if altered {
-                url.replace(DEFAULT_DIST_SERVER, tmp_cx.dist_server.as_str())
-            } else {
-                url
-            };
-
-            let url_url = utils::parse_url(&url)?;
-
-            let downloaded_file = RetryIf::spawn(
-                FixedInterval::from_millis(0).take(max_retries),
-                || download_cfg.download(&url_url, &hash),
-                |e: &anyhow::Error| {
-                    // retry only known retriable cases
-                    match e.downcast_ref::<RustupError>() {
-                        Some(RustupError::BrokenPartialFile)
-                        | Some(RustupError::DownloadingFile { .. }) => {
-                            (download_cfg.notify_handler)(Notification::RetryingDownload(&url));
-                            true
-                        }
-                        _ => false,
-                    }
-                },
-            )
-            .await
-            .with_context(|| RustupError::ComponentDownloadFailed(component.name(new_manifest)))?;
-
-            things_downloaded.push(hash);
+        }
 
-            things_to_install.push((component, format, downloaded_file));
+        let component_stream =
+            tokio_stream::iter(components.into_iter()).map(|(component, format, url, hash)| {
+                async move {
+                    let url = if altered {
+                        url.replace(DEFAULT_DIST_SERVER, tmp_cx.dist_server.as_str())
+                    } else {
+                        url
+                    };
+
+                    let url_url = utils::parse_url(&url)?;
+
+                    let downloaded_file = RetryIf::spawn(
+                        FixedInterval::from_millis(0).take(max_retries),
+                        || download_cfg.download(&url_url, &hash),
+                        |e: &anyhow::Error| {
+                            // retry only known retriable cases
+                            match e.downcast_ref::<RustupError>() {
+                                Some(RustupError::BrokenPartialFile)
+                                | Some(RustupError::DownloadingFile { .. }) => {
+                                    (download_cfg.notify_handler)(Notification::RetryingDownload(
+                                        &url,
+                                    ));
+                                    true
+                                }
+                                _ => false,
+                            }
+                        },
+                    )
+                    .await
+                    .with_context(|| {
+                        RustupError::ComponentDownloadFailed(component.name(new_manifest))
+                    })?;
+                    Ok::<_, Error>((component, format, downloaded_file, hash))
+                }
+            });
+        if components_len > 0 {
+            let results = component_stream
+                .buffered(components_len)
+                .collect::<Vec<_>>()
+                .await;
+            for result in results {
+                let (component, format, downloaded_file, hash) = result?;
+                things_downloaded.push(hash);
+                things_to_install.push((component, format, downloaded_file));
+            }
         }
 
         // Begin transaction

src/download/mod.rs

@@ -108,10 +108,13 @@ async fn download_file_(
 
         match msg {
             Event::DownloadContentLengthReceived(len) => {
-                notify_handler(Notification::DownloadContentLengthReceived(len));
+                notify_handler(Notification::DownloadContentLengthReceived(
+                    len,
+                    url.as_str(),
+                ));
             }
             Event::DownloadDataReceived(data) => {
-                notify_handler(Notification::DownloadDataReceived(data));
+                notify_handler(Notification::DownloadDataReceived(data, url.as_str()));
             }
             Event::ResumingPartialDownload => {
                 notify_handler(Notification::ResumingPartialDownload);
@@ -205,7 +208,7 @@ async fn download_file_(
         .download_to_path(url, path, resume_from_partial, Some(callback))
         .await;
 
-    notify_handler(Notification::DownloadFinished);
+    notify_handler(Notification::DownloadFinished(url.as_str()));
 
     res
 }

src/utils/mod.rs

@@ -506,7 +506,7 @@ impl<'a> FileReaderWithProgress<'a> {
 
         // Inform the tracker of the file size
         let flen = fh.metadata()?.len();
-        (notify_handler)(Notification::DownloadContentLengthReceived(flen));
+        (notify_handler)(Notification::DownloadContentLengthReceived(flen, ""));
 
         let fh = BufReader::with_capacity(8 * 1024 * 1024, fh);
 
@@ -525,10 +525,10 @@ impl io::Read for FileReaderWithProgress<'_> {
             Ok(nbytes) => {
                 self.nbytes += nbytes as u64;
                 if nbytes != 0 {
-                    (self.notify_handler)(Notification::DownloadDataReceived(&buf[0..nbytes]));
+                    (self.notify_handler)(Notification::DownloadDataReceived(&buf[0..nbytes], ""));
                 }
                 if (nbytes == 0) || (self.flen == self.nbytes) {
-                    (self.notify_handler)(Notification::DownloadFinished);
+                    (self.notify_handler)(Notification::DownloadFinished(""));
                 }
                 Ok(nbytes)
             }

src/utils/notifications.rs

@@ -13,12 +13,13 @@ pub enum Notification<'a> {
     CopyingDirectory(&'a Path, &'a Path),
     RemovingDirectory(&'a str, &'a Path),
     DownloadingFile(&'a Url, &'a Path),
-    /// Received the Content-Length of the to-be downloaded data.
-    DownloadContentLengthReceived(u64),
+    /// Received the Content-Length of the to-be downloaded data with
+    /// the respective URL of the download (for tracking concurrent downloads).
+    DownloadContentLengthReceived(u64, &'a str),
     /// Received some data.
-    DownloadDataReceived(&'a [u8]),
+    DownloadDataReceived(&'a [u8], &'a str),
     /// Download has finished.
-    DownloadFinished,
+    DownloadFinished(&'a str),
     /// The things we're tracking that are not counted in bytes.
     /// Must be paired with a pop-units; our other calls are not
     /// setup to guarantee this any better.
@@ -52,11 +53,11 @@ impl Notification<'_> {
             | LinkingDirectory(_, _)
             | CopyingDirectory(_, _)
             | DownloadingFile(_, _)
-            | DownloadContentLengthReceived(_)
-            | DownloadDataReceived(_)
+            | DownloadContentLengthReceived(_, _)
+            | DownloadDataReceived(_, _)
             | DownloadPushUnit(_)
             | DownloadPopUnit
-            | DownloadFinished
+            | DownloadFinished(_)
             | ResumingPartialDownload
             | UsingCurl
             | UsingReqwest => NotificationLevel::Debug,
@@ -92,11 +93,11 @@ impl Display for Notification<'_> {
                 units::Size::new(*size, units::Unit::B)
             ),
             DownloadingFile(url, _) => write!(f, "downloading file from: '{url}'"),
-            DownloadContentLengthReceived(len) => write!(f, "download size is: '{len}'"),
-            DownloadDataReceived(data) => write!(f, "received some data of size {}", data.len()),
+            DownloadContentLengthReceived(len, _) => write!(f, "download size is: '{len}'"),
+            DownloadDataReceived(data, _) => write!(f, "received some data of size {}", data.len()),
             DownloadPushUnit(_) => Ok(()),
             DownloadPopUnit => Ok(()),
-            DownloadFinished => write!(f, "download finished"),
+            DownloadFinished(_) => write!(f, "download finished"),
             NoCanonicalPath(path) => write!(f, "could not canonicalize path: '{}'", path.display()),
             ResumingPartialDownload => write!(f, "resuming partial download"),
             UsingCurl => write!(f, "downloading with curl"),