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.
This commit is contained in:
+28
-25
@@ -1,7 +1,7 @@
|
|||||||
|
use crate::progress::{Progress, Unit};
|
||||||
use crate::veprintln;
|
use crate::veprintln;
|
||||||
use anyhow::{anyhow, Context, Result};
|
use anyhow::{anyhow, Context, Result};
|
||||||
use futures_util::StreamExt;
|
use futures_util::StreamExt;
|
||||||
use indicatif::{ProgressBar, ProgressStyle};
|
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::path::{Path, PathBuf};
|
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"))
|
.ok_or_else(|| anyhow!("Registry did not return a Docker-Content-Digest header"))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Download a container image (either direct tarball or OCI image)
|
/// Download a container image (either direct tarball or OCI image),
|
||||||
pub fn download_image(source: &ImageSource, dest: &Path, arch: &str) -> Result<()> {
|
/// 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.
|
// Create a single Tokio runtime for all async I/O in this download.
|
||||||
// Previously each of the two sync wrappers (download_file_sync and
|
// Previously each of the two sync wrappers (download_file_sync and
|
||||||
// download_oci_image) created their own runtime; this consolidates them.
|
// 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://") {
|
if !url.starts_with("http://") && !url.starts_with("https://") {
|
||||||
return Err(anyhow!("Unsupported URL scheme: {}", url));
|
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 {
|
ImageSource::OciImage {
|
||||||
registry,
|
registry,
|
||||||
@@ -101,12 +107,13 @@ pub fn download_image(source: &ImageSource, dest: &Path, arch: &str) -> Result<(
|
|||||||
tag,
|
tag,
|
||||||
architecture,
|
architecture,
|
||||||
dest,
|
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);
|
veprintln!("Downloading: {}", url);
|
||||||
|
|
||||||
let client = Client::builder()
|
let client = Client::builder()
|
||||||
@@ -133,12 +140,11 @@ async fn download_file_async(url: &str, dest: &Path) -> Result<()> {
|
|||||||
.context("Failed to create cache directory")?;
|
.context("Failed to create cache directory")?;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Setup progress bar
|
progress.begin(
|
||||||
let pb = ProgressBar::new(total_size);
|
"Downloading",
|
||||||
pb.set_style(ProgressStyle::default_bar()
|
Unit::Bytes,
|
||||||
.template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({eta})")
|
(total_size > 0).then_some(total_size),
|
||||||
.unwrap()
|
);
|
||||||
.progress_chars("#>-"));
|
|
||||||
|
|
||||||
let temp_dest = dest.with_extension("partial");
|
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 {
|
while let Some(chunk) = stream.next().await {
|
||||||
let chunk = chunk.context("Download interrupted")?;
|
let chunk = chunk.context("Download interrupted")?;
|
||||||
tokio::io::copy(&mut &chunk[..], &mut file).await?;
|
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)
|
tokio::fs::rename(&temp_dest, dest)
|
||||||
.await
|
.await
|
||||||
@@ -182,13 +188,15 @@ async fn download_oci_image_async(
|
|||||||
tag: &str,
|
tag: &str,
|
||||||
arch: &str,
|
arch: &str,
|
||||||
dest: &Path,
|
dest: &Path,
|
||||||
|
progress: &dyn Progress,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let client = Client::builder()
|
let client = Client::builder()
|
||||||
.timeout(Duration::from_secs(300))
|
.timeout(Duration::from_secs(300))
|
||||||
.build()
|
.build()
|
||||||
.context("Failed to create HTTP client")?;
|
.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
|
// Persist the manifest digest so the next run can skip the download when
|
||||||
// the tag hasn't changed (see fetch_oci_digest / digest_sidecar).
|
// the tag hasn't changed (see fetch_oci_digest / digest_sidecar).
|
||||||
@@ -206,6 +214,7 @@ async fn try_download_oci_image(
|
|||||||
tag: &str,
|
tag: &str,
|
||||||
arch: &str,
|
arch: &str,
|
||||||
dest: &Path,
|
dest: &Path,
|
||||||
|
progress: &dyn Progress,
|
||||||
) -> Result<String> {
|
) -> Result<String> {
|
||||||
// Get authentication token for the registry
|
// Get authentication token for the registry
|
||||||
let token = get_auth_token(client, registry, repository).await?;
|
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 layer_path = temp_dir.path().join(&layer_name);
|
||||||
let label = format!("Layer {}/{} ({})", i + 1, layers.len(), 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);
|
layer_names.push(layer_name);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -553,6 +562,7 @@ async fn download_blob_with_auth(
|
|||||||
dest: &Path,
|
dest: &Path,
|
||||||
token: &str,
|
token: &str,
|
||||||
label: &str,
|
label: &str,
|
||||||
|
progress: &dyn Progress,
|
||||||
) -> Result<()> {
|
) -> Result<()> {
|
||||||
let mut request = client.get(url);
|
let mut request = client.get(url);
|
||||||
if !token.is_empty() {
|
if !token.is_empty() {
|
||||||
@@ -570,14 +580,7 @@ async fn download_blob_with_auth(
|
|||||||
|
|
||||||
let total_size = response.content_length().unwrap_or(0);
|
let total_size = response.content_length().unwrap_or(0);
|
||||||
|
|
||||||
let pb = ProgressBar::new(total_size);
|
progress.begin(label, Unit::Bytes, (total_size > 0).then_some(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());
|
|
||||||
|
|
||||||
let mut file = tokio::fs::File::create(dest)
|
let mut file = tokio::fs::File::create(dest)
|
||||||
.await
|
.await
|
||||||
@@ -592,11 +595,11 @@ async fn download_blob_with_auth(
|
|||||||
file.write_all(&chunk)
|
file.write_all(&chunk)
|
||||||
.await
|
.await
|
||||||
.context("Failed to write chunk")?;
|
.context("Failed to write chunk")?;
|
||||||
pb.inc(chunk.len() as u64);
|
progress.advance(chunk.len() as u64);
|
||||||
}
|
}
|
||||||
|
|
||||||
file.flush().await?;
|
file.flush().await?;
|
||||||
pb.finish_and_clear();
|
progress.finish();
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -140,7 +140,7 @@ impl RootfsCache {
|
|||||||
std::fs::create_dir_all(&self.dir).with_context(|| {
|
std::fs::create_dir_all(&self.dir).with_context(|| {
|
||||||
format!("Failed to create cache directory {}", self.dir.display())
|
format!("Failed to create cache directory {}", self.dir.display())
|
||||||
})?;
|
})?;
|
||||||
download_image(&source, &cache_path, &arch)?;
|
download_image(&source, &cache_path, &arch, progress)?;
|
||||||
} else {
|
} else {
|
||||||
veprintln!("Using cached tarball: {}", cache_path.display());
|
veprintln!("Using cached tarball: {}", cache_path.display());
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user