feat(downloads): concurrently download components

Some notifications needed to be updated to include the download URL,
enabling the identification of the component being downloaded. This was
necessary for accurate progress reporting of each component.

Author: FranciscoTGouveia

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,61 @@ 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 find_progress_bar(&mut self, file: &str) -> Option<&mut 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))
+            .cloned()?;
+
+        self.file_progress_bars.get_mut(&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.find_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.find_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.find_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,69 @@ 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,
+        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::<(Component, CompressionKind, File, String), Error>((
+                        component,
+                        format,
+                        downloaded_file,
+                        hash,
+                    ))
+                }
+            });
+        if components_len > 0 {
+            let results = component_stream
+                .buffered(components_len)
+                .collect::<Vec<_>>()
+                .await;
+            for result in results {
+                match result {
+                    Ok((component, format, downloaded_file, hash)) => {
+                        things_downloaded.push(hash);
+                        things_to_install.push((component, format, downloaded_file));
                     }
-                },
-            )
-            .await
-            .with_context(|| RustupError::ComponentDownloadFailed(component.name(new_manifest)))?;
-
-            things_downloaded.push(hash);
-
-            things_to_install.push((component, format, downloaded_file));
+                    Err(e) => return Err(e),
+                }
+            }
         }
 
         // 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

@@ -14,11 +14,11 @@ pub enum Notification<'a> {
     RemovingDirectory(&'a str, &'a Path),
     DownloadingFile(&'a Url, &'a Path),
     /// Received the Content-Length of the to-be downloaded data.
-    DownloadContentLengthReceived(u64),
+    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 +52,11 @@ impl Notification<'_> {
             | LinkingDirectory(_, _)
             | CopyingDirectory(_, _)
             | DownloadingFile(_, _)
-            | DownloadContentLengthReceived(_)
-            | DownloadDataReceived(_)
+            | DownloadContentLengthReceived(_, _)
+            | DownloadDataReceived(_, _)
             | DownloadPushUnit(_)
             | DownloadPopUnit
-            | DownloadFinished
+            | DownloadFinished(_)
             | ResumingPartialDownload
             | UsingCurl
             | UsingReqwest => NotificationLevel::Debug,
@@ -92,11 +92,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"),