104 lines
3.3 KiB
Rust
104 lines
3.3 KiB
Rust
use miette::{IntoDiagnostic, Result};
|
|
use std::{collections::HashMap, fmt::Write, pin::pin};
|
|
|
|
use bollard::{errors::Error, secret::CreateImageInfo};
|
|
use futures::{Stream, TryStreamExt};
|
|
use indicatif::{MultiProgress, ProgressBar, ProgressState, ProgressStyle};
|
|
|
|
fn new_download_pb(multi: &MultiProgress, layer_id: &str) -> ProgressBar {
|
|
let pb = multi.add(ProgressBar::new(0));
|
|
pb.set_style(
|
|
ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] {msg} [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
|
|
.unwrap()
|
|
.with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
|
|
write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
|
|
})
|
|
.progress_chars("#>-"),
|
|
);
|
|
pb.set_message(format!("Layer {}", layer_id));
|
|
pb
|
|
}
|
|
|
|
fn download_drive_progress(pb: &mut ProgressBar, download_info: CreateImageInfo) {
|
|
match download_info.progress_detail {
|
|
Some(info) => match (info.current, info.total) {
|
|
(None, None) => {
|
|
pb.inc(1);
|
|
}
|
|
(current, total) => {
|
|
if let Some(total) = total {
|
|
pb.set_length(total as u64);
|
|
}
|
|
if let Some(current) = current {
|
|
pb.set_position(current as u64);
|
|
}
|
|
|
|
if let (Some(current), Some(total)) = (current, total)
|
|
&& (current == total)
|
|
{
|
|
pb.finish_with_message("Completed!");
|
|
}
|
|
}
|
|
},
|
|
None => {
|
|
// No progress detail, just show activity
|
|
pb.tick();
|
|
}
|
|
}
|
|
}
|
|
|
|
fn download_check_for_error(
|
|
layer_progress: &mut HashMap<String, ProgressBar>,
|
|
download_info: &CreateImageInfo,
|
|
) -> Result<()> {
|
|
if let Some(error_detail) = &download_info.error_detail {
|
|
for (_, pb) in layer_progress.drain() {
|
|
pb.finish_and_clear();
|
|
}
|
|
|
|
match (error_detail.code, &error_detail.message) {
|
|
(None, Some(msg)) => miette::bail!("docker image download error: {}", msg),
|
|
(Some(code), None) => miette::bail!("docker image download error: code {}", code),
|
|
(Some(code), Some(msg)) => {
|
|
miette::bail!(
|
|
"docker image download error: code {}, message: {}",
|
|
code,
|
|
msg
|
|
)
|
|
}
|
|
_ => (),
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// sadly type ... = impl ... is unstable
|
|
pub async fn perform_download(
|
|
multi: MultiProgress,
|
|
chunks: impl Stream<Item = Result<CreateImageInfo, Error>>,
|
|
) -> Result<()> {
|
|
let mut chunks = pin!(chunks);
|
|
let mut layer_progress: HashMap<String, ProgressBar> = HashMap::new();
|
|
|
|
while let Some(download_info) = chunks.try_next().await.into_diagnostic()? {
|
|
download_check_for_error(&mut layer_progress, &download_info)?;
|
|
|
|
let layer_id = download_info.id.as_deref().unwrap_or("unknown");
|
|
|
|
// Get or create progress bar for this layer
|
|
let pb = layer_progress
|
|
.entry(layer_id.to_string())
|
|
.or_insert_with(|| new_download_pb(&multi, layer_id));
|
|
|
|
download_drive_progress(pb, download_info);
|
|
}
|
|
|
|
// Clean up any remaining progress bars
|
|
for (_, pb) in layer_progress.drain() {
|
|
pb.finish_and_clear();
|
|
}
|
|
|
|
Ok(())
|
|
}
|