feat: oci progress bars on launch

This commit is contained in:
Alex Zenla
2024-04-12 12:38:31 +00:00
parent 8d26cd8fb2
commit cb89ac02d1
9 changed files with 216 additions and 46 deletions

View File

@ -18,6 +18,7 @@ ctrlc = { workspace = true, features = ["termination"] }
env_logger = { workspace = true }
fancy-duration = { workspace = true }
human_bytes = { workspace = true }
indicatif = { workspace = true }
krata = { path = "../krata", version = "^0.0.8" }
log = { workspace = true }
prost-reflect = { workspace = true, features = ["serde"] }

View File

@ -2,6 +2,7 @@ use std::collections::HashMap;
use anyhow::Result;
use clap::Parser;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use krata::{
events::EventStream,
v1::{
@ -11,7 +12,7 @@ use krata::{
},
control::{
control_service_client::ControlServiceClient, watch_events_reply::Event,
CreateGuestRequest,
CreateGuestRequest, OciProgressEventLayerPhase, OciProgressEventPhase,
},
},
};
@ -125,9 +126,14 @@ impl LauchCommand {
async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
let mut stream = events.subscribe();
let mut multi_progress: Option<(MultiProgress, HashMap<String, ProgressBar>)> = None;
while let Ok(event) = stream.recv().await {
match event {
Event::GuestChanged(changed) => {
if let Some((multi_progress, _)) = multi_progress.as_mut() {
let _ = multi_progress.clear();
}
let Some(guest) = changed.guest else {
continue;
};
@ -159,7 +165,83 @@ async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
}
}
Event::OciProgress(_oci) => {}
Event::OciProgress(oci) => {
if multi_progress.is_none() {
multi_progress = Some((MultiProgress::new(), HashMap::new()));
}
let Some((multi_progress, progresses)) = multi_progress.as_mut() else {
continue;
};
match oci.phase() {
OciProgressEventPhase::Resolved
| OciProgressEventPhase::ConfigAcquire
| OciProgressEventPhase::LayerAcquire => {
if progresses.is_empty() && !oci.layers.is_empty() {
for layer in &oci.layers {
let bar = ProgressBar::new(layer.total);
bar.set_style(
ProgressStyle::with_template("{msg} {wide_bar} {pos}/{len}")
.unwrap(),
);
progresses.insert(layer.id.clone(), bar.clone());
multi_progress.add(bar);
}
}
for layer in oci.layers {
let Some(progress) = progresses.get_mut(&layer.id) else {
continue;
};
let phase = match layer.phase() {
OciProgressEventLayerPhase::Waiting => "waiting",
OciProgressEventLayerPhase::Downloading => "downloading",
OciProgressEventLayerPhase::Downloaded => "downloaded",
OciProgressEventLayerPhase::Extracting => "extracting",
OciProgressEventLayerPhase::Extracted => "extracted",
_ => "unknown",
};
progress.set_message(format!("{} {}", layer.id, phase));
progress.set_length(layer.total);
progress.set_position(layer.value);
}
}
OciProgressEventPhase::Packing => {
for (key, progress) in &mut *progresses {
if key == "packing" {
continue;
}
progress.finish_and_clear();
multi_progress.remove(progress);
}
progresses.retain(|k, _| k == "packing");
if progresses.is_empty() {
let progress = ProgressBar::new(100);
progress.set_style(
ProgressStyle::with_template("{msg} {wide_bar} {pos}/{len}")
.unwrap(),
);
progresses.insert("packing".to_string(), progress);
}
let Some(progress) = progresses.get("packing") else {
continue;
};
progress.set_message("packing image");
progress.set_length(oci.total);
progress.set_position(oci.value);
}
_ => {}
}
for progress in progresses {
progress.1.tick();
}
}
}
}
Ok(())