add progress bar to client downloads

bytecodealliance · Sep 24, 2024 · c2fcab8 · c2fcab8
1 parent 66639c1
commit c2fcab8
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 16 deletions.
diff --git a/crates/client/src/api.rs b/crates/client/src/api.rs
@@ -457,6 +457,7 @@ impl Client {
         &self,
         registry_domain: Option<&RegistryDomain>,
         digest: &AnyHash,
+        mut with_progress: impl FnMut(u64, u64) + Send + Sync + 'static,
     ) -> Result<impl Stream<Item = Result<Bytes>>, ClientError> {
         let ContentSourcesResponse { content_sources } =
             self.content_sources(registry_domain, digest).await?;
@@ -479,10 +480,16 @@ impl Client {
                 continue;
             }
 
-            return Ok(validate_stream(
-                digest,
-                response.bytes_stream().map_err(|e| anyhow!(e)),
-            ));
+            let total_bytes = response.content_length().unwrap_or(0);
+            let stream = response
+                .bytes_stream()
+                .map_err(|e| anyhow!(e))
+                .inspect(move |r| {
+                    if let Ok(b) = r {
+                        with_progress(b.len() as u64, total_bytes);
+                    }
+                });
+            return Ok(validate_stream(digest, stream));
         }
 
         Err(ClientError::AllSourcesFailed(digest.clone()))

diff --git a/crates/client/src/depsolve.rs b/crates/client/src/depsolve.rs
@@ -120,7 +120,7 @@ impl LockListBuilder {
                         }
                         self.lock_list.insert(import);
                     } else {
-                        client.download(&id, &VersionReq::STAR).await?;
+                        client.download(&id, &VersionReq::STAR, |_, _| ()).await?;
                         if let Some(info) = client
                             .registry()
                             .load_package(

diff --git a/crates/client/src/lib.rs b/crates/client/src/lib.rs
@@ -655,6 +655,7 @@ impl<R: RegistryStorage, C: ContentStorage, N: NamespaceMapStorage> Client<R, C,
         &self,
         package: &PackageName,
         requirement: &VersionReq,
+        with_progress: impl FnMut(u64, u64) + Send + Sync + 'static,
     ) -> Result<Option<PackageDownload>, ClientError> {
         let info = self.package(package).await?;
 
@@ -674,7 +675,7 @@ impl<R: RegistryStorage, C: ContentStorage, N: NamespaceMapStorage> Client<R, C,
                     .context("invalid state: not yanked but missing content")?
                     .clone();
                 let path = self
-                    .download_content(registry_domain.as_ref(), &digest)
+                    .download_content(registry_domain.as_ref(), &digest, with_progress)
                     .await?;
                 Ok(Some(PackageDownload {
                     version: release.version.clone(),
@@ -776,7 +777,7 @@ impl<R: RegistryStorage, C: ContentStorage, N: NamespaceMapStorage> Client<R, C,
             version: version.clone(),
             digest: digest.clone(),
             path: self
-                .download_content(registry_domain.as_ref(), digest)
+                .download_content(registry_domain.as_ref(), digest, |_, _| ())
                 .await?,
         })
     }
@@ -1310,6 +1311,7 @@ current_registry = registry_domain.map(|d| d.as_str()).unwrap_or(&self.url().saf
         &self,
         registry_domain: Option<&RegistryDomain>,
         digest: &AnyHash,
+        with_progress: impl FnMut(u64, u64) + Send + Sync + 'static,
     ) -> Result<PathBuf, ClientError> {
         match self.content.content_location(digest) {
             Some(path) => {
@@ -1319,7 +1321,11 @@ current_registry = registry_domain.map(|d| d.as_str()).unwrap_or(&self.url().saf
             None => {
                 self.content
                     .store_content(
-                        Box::pin(self.api.download_content(registry_domain, digest).await?),
+                        Box::pin(
+                            self.api
+                                .download_content(registry_domain, digest, with_progress)
+                                .await?,
+                        ),
                         Some(digest),
                     )
                     .await?;
@@ -1352,7 +1358,9 @@ current_registry = registry_domain.map(|d| d.as_str()).unwrap_or(&self.url().saf
                 Ok(ReaderStream::new(file).map_err(Into::into).boxed())
             }
             None => Ok(Box::pin(
-                self.api.download_content(registry_domain, digest).await?,
+                self.api
+                    .download_content(registry_domain, digest, |_, _| ())
+                    .await?,
             )),
         }
     }

diff --git a/src/commands/dependencies.rs b/src/commands/dependencies.rs
@@ -1,3 +1,5 @@
+use crate::progress::make_progress_handler;
+
 use super::CommonOptions;
 use anyhow::{bail, Result};
 use async_recursion::async_recursion;
@@ -47,9 +49,14 @@ impl DependenciesCommand {
         node: &mut TreeBuilder,
         parser: &mut DepsParser,
     ) -> Result<()> {
-        client.download(id, &version).await?;
-
-        if let Some(download) = client.download(id, &version).await? {
+        if let Some(download) = client
+            .download(
+                id,
+                &version,
+                make_progress_handler(format!("Downloading `{id}`")),
+            )
+            .await?
+        {
             let bytes = fs::read(download.path)?;
             let deps = parser.parse(&bytes)?;
             for dep in deps {
@@ -77,7 +84,11 @@ impl DependenciesCommand {
     async fn print_package_info(client: &FileSystemClient, info: &PackageInfo) -> Result<()> {
         let mut parser = DepsParser::new();
 
-        if let Some(download) = client.download(&info.name, &Default::default()).await? {
+        let progress = |_inc, _total| {};
+        if let Some(download) = client
+            .download(&info.name, &Default::default(), progress)
+            .await?
+        {
             let mut tree = new_tree(info.name.namespace(), info.name.name(), &download.version);
             let bytes = fs::read(&download.path)?;
             let deps = parser.parse(&bytes)?;

diff --git a/src/commands/download.rs b/src/commands/download.rs
@@ -1,3 +1,5 @@
+use crate::progress::make_progress_handler;
+
 use super::CommonOptions;
 use anyhow::Result;
 use clap::Args;
@@ -30,16 +32,18 @@ impl DownloadCommand {
         let config = self.common.read_config()?;
         let client = self.common.create_client(&config).await?;
 
-        println!("Downloading `{name}`...", name = self.name);
-
         // if user specifies exact verion, then set the `VersionReq` to exact match
         let version = match &self.version {
             Some(version) => VersionReq::parse(&format!("={}", version))?,
             None => VersionReq::STAR,
         };
 
         let download = client
-            .download(&self.name, &version)
+            .download(
+                &self.name,
+                &version,
+                make_progress_handler(format!("Downloading `{name}`...", name = self.name)),
+            )
             .await?
             .ok_or_else(|| ClientError::PackageVersionRequirementDoesNotExist {
                 name: self.name.clone(),