From 3fb4be667c3d2c8d742ce87c1fffb0212737087f Mon Sep 17 00:00:00 2001 From: Valentin Haudiquet Date: Sat, 3 Oct 2026 00:02:11 +0200 Subject: [PATCH] refactor(download): report download progress through the sink The direct-tarball and OCI-blob download bars drew to stderr from inside the library; download_image now emits Progress events with the content length when the server provides one. --- crates/ecr/src/download.rs | 53 ++++++++++++++++++++------------------ crates/ecr/src/rootfs.rs | 2 +- 2 files changed, 29 insertions(+), 26 deletions(-) diff --git a/crates/ecr/src/download.rs b/crates/ecr/src/download.rs index 8266759..4cece4c 100644 --- a/crates/ecr/src/download.rs +++ b/crates/ecr/src/download.rs @@ -1,7 +1,7 @@ +use crate::progress::{Progress, Unit}; use crate::veprintln; use anyhow::{anyhow, Context, Result}; use futures_util::StreamExt; -use indicatif::{ProgressBar, ProgressStyle}; use reqwest::Client; use serde_json::Value; use std::path::{Path, PathBuf}; @@ -72,8 +72,14 @@ async fn fetch_oci_digest_async(registry: &str, repository: &str, tag: &str) -> .ok_or_else(|| anyhow!("Registry did not return a Docker-Content-Digest header")) } -/// Download a container image (either direct tarball or OCI image) -pub fn download_image(source: &ImageSource, dest: &Path, arch: &str) -> Result<()> { +/// Download a container image (either direct tarball or OCI image), +/// reporting transfer progress through `progress` +pub fn download_image( + source: &ImageSource, + dest: &Path, + arch: &str, + progress: &dyn Progress, +) -> Result<()> { // Create a single Tokio runtime for all async I/O in this download. // Previously each of the two sync wrappers (download_file_sync and // download_oci_image) created their own runtime; this consolidates them. @@ -86,7 +92,7 @@ pub fn download_image(source: &ImageSource, dest: &Path, arch: &str) -> Result<( if !url.starts_with("http://") && !url.starts_with("https://") { return Err(anyhow!("Unsupported URL scheme: {}", url)); } - rt.block_on(download_file_async(&url, dest)) + rt.block_on(download_file_async(&url, dest, progress)) } ImageSource::OciImage { registry, @@ -101,12 +107,13 @@ pub fn download_image(source: &ImageSource, dest: &Path, arch: &str) -> Result<( tag, architecture, dest, + progress, )) } } } -async fn download_file_async(url: &str, dest: &Path) -> Result<()> { +async fn download_file_async(url: &str, dest: &Path, progress: &dyn Progress) -> Result<()> { veprintln!("Downloading: {}", url); let client = Client::builder() @@ -133,12 +140,11 @@ async fn download_file_async(url: &str, dest: &Path) -> Result<()> { .context("Failed to create cache directory")?; } - // Setup progress bar - let pb = ProgressBar::new(total_size); - pb.set_style(ProgressStyle::default_bar() - .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})") - .unwrap() - .progress_chars("#>-")); + progress.begin( + "Downloading", + Unit::Bytes, + (total_size > 0).then_some(total_size), + ); let temp_dest = dest.with_extension("partial"); @@ -154,10 +160,10 @@ async fn download_file_async(url: &str, dest: &Path) -> Result<()> { while let Some(chunk) = stream.next().await { let chunk = chunk.context("Download interrupted")?; tokio::io::copy(&mut &chunk[..], &mut file).await?; - pb.inc(chunk.len() as u64); + progress.advance(chunk.len() as u64); } - pb.finish_and_clear(); + progress.finish(); tokio::fs::rename(&temp_dest, dest) .await @@ -182,13 +188,15 @@ async fn download_oci_image_async( tag: &str, arch: &str, dest: &Path, + progress: &dyn Progress, ) -> Result<()> { let client = Client::builder() .timeout(Duration::from_secs(300)) .build() .context("Failed to create HTTP client")?; - let digest = try_download_oci_image(&client, registry, repository, tag, arch, dest).await?; + let digest = + try_download_oci_image(&client, registry, repository, tag, arch, dest, progress).await?; // Persist the manifest digest so the next run can skip the download when // the tag hasn't changed (see fetch_oci_digest / digest_sidecar). @@ -206,6 +214,7 @@ async fn try_download_oci_image( tag: &str, arch: &str, dest: &Path, + progress: &dyn Progress, ) -> Result { // Get authentication token for the registry let token = get_auth_token(client, registry, repository).await?; @@ -370,7 +379,7 @@ async fn try_download_oci_image( let layer_path = temp_dir.path().join(&layer_name); let label = format!("Layer {}/{} ({})", i + 1, layers.len(), layer_name); - download_blob_with_auth(client, &blob_url, &layer_path, &token, &label).await?; + download_blob_with_auth(client, &blob_url, &layer_path, &token, &label, progress).await?; layer_names.push(layer_name); } @@ -553,6 +562,7 @@ async fn download_blob_with_auth( dest: &Path, token: &str, label: &str, + progress: &dyn Progress, ) -> Result<()> { let mut request = client.get(url); if !token.is_empty() { @@ -570,14 +580,7 @@ async fn download_blob_with_auth( let total_size = response.content_length().unwrap_or(0); - let pb = ProgressBar::new(total_size); - pb.set_style( - ProgressStyle::default_bar() - .template("{msg} [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})") - .unwrap() - .progress_chars("#>-"), - ); - pb.set_message(label.to_string()); + progress.begin(label, Unit::Bytes, (total_size > 0).then_some(total_size)); let mut file = tokio::fs::File::create(dest) .await @@ -592,11 +595,11 @@ async fn download_blob_with_auth( file.write_all(&chunk) .await .context("Failed to write chunk")?; - pb.inc(chunk.len() as u64); + progress.advance(chunk.len() as u64); } file.flush().await?; - pb.finish_and_clear(); + progress.finish(); Ok(()) } diff --git a/crates/ecr/src/rootfs.rs b/crates/ecr/src/rootfs.rs index fece38f..948d760 100644 --- a/crates/ecr/src/rootfs.rs +++ b/crates/ecr/src/rootfs.rs @@ -140,7 +140,7 @@ impl RootfsCache { std::fs::create_dir_all(&self.dir).with_context(|| { format!("Failed to create cache directory {}", self.dir.display()) })?; - download_image(&source, &cache_path, &arch)?; + download_image(&source, &cache_path, &arch, progress)?; } else { veprintln!("Using cached tarball: {}", cache_path.display()); }