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 URLs being downloaded to their corresponding progress bars.
+    file_progress_bars: HashMap<String, ProgressBar>,
 }
 
 impl DownloadTracker {
@@ -28,22 +28,35 @@ 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,
+                url,
+            ))) => {
+                if let Some(url) = url {
+                    self.content_length_received(content_len, url);
+                }
                 true
             }
-            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
-                self.data_received(data.len());
+            Notification::Install(In::Utils(Un::DownloadDataReceived(data, url))) => {
+                if let Some(url) = url {
+                    self.data_received(data.len(), url);
+                }
                 true
             }
-            Notification::Install(In::Utils(Un::DownloadFinished)) => {
-                self.download_finished();
+            Notification::Install(In::Utils(Un::DownloadFinished(url))) => {
+                if let Some(url) = url {
+                    self.download_finished(url);
+                }
+                true
+            }
+            Notification::Install(In::DownloadingComponent(component, _, _, url)) => {
+                self.create_progress_bar(component.to_owned(), url.to_owned());
                 true
             }
             Notification::Install(In::Utils(Un::DownloadPushUnit(_))) => true,
@@ -53,30 +66,44 @@ 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(
+    /// Creates a new ProgressBar for the given component.
+    pub(crate) fn create_progress_bar(&mut self, component: String, url: String) {
+        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);
+        self.multi_progress_bars.add(pb.clone());
+        self.file_progress_bars.insert(url, pb);
+    }
+
+    /// Sets the length for a new ProgressBar and gives it a style.
+    pub(crate) fn content_length_received(&mut self, content_len: u64, url: &str) {
+        if let Some(pb) = self.file_progress_bars.get(url) {
+            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, url: &str) {
+        if let Some(pb) = self.file_progress_bars.get(url) {
+            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, url: &str) {
+        let Some(pb) = self.file_progress_bars.get(url) else {
+            return;
+        };
+        pb.set_style(
+            ProgressStyle::with_template("{msg:>12.bold}  downloaded {total_bytes} in {elapsed}")
+                .unwrap(),
+        );
+        pb.finish();
     }
 }

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(None));
             handler(Notification::DownloadPushUnit(Unit::IO));
             handler(Notification::DownloadContentLengthReceived(
                 prev_files as u64,
+                None,
             ));
         }
         if prev_files > 50 {
@@ -284,12 +285,15 @@ 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],
+                    None,
+                ));
             }
         }
         self.pool.join();
         if let Some(handler) = self.notify_handler {
-            handler(Notification::DownloadFinished);
+            handler(Notification::DownloadFinished(None));
             handler(Notification::DownloadPopUnit);
         }
         // close the feedback channel so that blocking reads on it can

src/dist/manifestation.rs

@@ -7,7 +7,8 @@ mod tests;
 use std::path::Path;
 
 use anyhow::{Context, Result, anyhow, bail};
-use tokio_retry::{RetryIf, strategy::FixedInterval};
+use futures_util::stream::StreamExt;
+use tracing::info;
 
 use crate::dist::component::{
     Components, Package, TarGzPackage, TarXzPackage, TarZStdPackage, Transaction,
@@ -153,6 +154,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 +164,40 @@ 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, _, url, _) in components.clone() {
             (download_cfg.notify_handler)(Notification::DownloadingComponent(
                 &component.short_name(new_manifest),
                 &self.target_triple,
                 component.target.as_ref(),
+                &url,
             ));
-            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)| {
+                self.download_component(
+                    component,
+                    format,
+                    url,
+                    hash,
+                    altered,
+                    tmp_cx,
+                    download_cfg,
+                    max_retries,
+                    new_manifest,
+                )
+            });
+        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
@@ -452,6 +453,7 @@ impl Manifestation {
             "rust",
             &self.target_triple,
             Some(&self.target_triple),
+            &url,
         ));
 
         use std::path::PathBuf;
@@ -531,6 +533,50 @@ impl Manifestation {
 
         Ok(tx)
     }
+
+    #[allow(clippy::too_many_arguments)]
+    async fn download_component(
+        &self,
+        component: Component,
+        format: CompressionKind,
+        url: String,
+        hash: String,
+        altered: bool,
+        tmp_cx: &temp::Context,
+        download_cfg: &DownloadCfg<'_>,
+        max_retries: usize,
+        new_manifest: &Manifest,
+    ) -> Result<(Component, CompressionKind, File, String)> {
+        use tokio_retry::{RetryIf, strategy::FixedInterval};
+
+        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((component, format, downloaded_file, hash))
+    }
 }
 
 #[derive(Debug)]

src/dist/notifications.rs

@@ -23,7 +23,8 @@ pub enum Notification<'a> {
     ExtensionNotInstalled(&'a str),
     NonFatalError(&'a anyhow::Error),
     MissingInstalledComponent(&'a str),
-    DownloadingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
+    /// The URL of the download is passed as the last argument, to allow us to track concurrent downloads.
+    DownloadingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>, &'a str),
     InstallingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
     RemovingComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
     RemovingOldComponent(&'a str, &'a TargetTriple, Option<&'a TargetTriple>),
@@ -61,7 +62,7 @@ impl Notification<'_> {
             | FileAlreadyDownloaded
             | DownloadingLegacyManifest => NotificationLevel::Debug,
             Extracting(_, _)
-            | DownloadingComponent(_, _, _)
+            | DownloadingComponent(_, _, _, _)
             | InstallingComponent(_, _, _)
             | RemovingComponent(_, _, _)
             | RemovingOldComponent(_, _, _)
@@ -107,7 +108,7 @@ impl Display for Notification<'_> {
             MissingInstalledComponent(c) => {
                 write!(f, "during uninstall component {c} was not found")
             }
-            DownloadingComponent(c, h, t) => {
+            DownloadingComponent(c, h, t, _) => {
                 if Some(h) == t.as_ref() || t.is_none() {
                     write!(f, "downloading component '{c}'")
                 } else {

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,
+                    Some(url.as_str()),
+                ));
             }
             Event::DownloadDataReceived(data) => {
-                notify_handler(Notification::DownloadDataReceived(data));
+                notify_handler(Notification::DownloadDataReceived(data, Some(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(Some(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, None));
 
         let fh = BufReader::with_capacity(8 * 1024 * 1024, fh);
 
@@ -525,10 +525,13 @@ 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],
+                        None,
+                    ));
                 }
                 if (nbytes == 0) || (self.flen == self.nbytes) {
-                    (self.notify_handler)(Notification::DownloadFinished);
+                    (self.notify_handler)(Notification::DownloadFinished(None));
                 }
                 Ok(nbytes)
             }