feat: implement oci image progress (#64)

* feat: oci progress events

* feat: oci progress bars on launch

crates/ctl/Cargo.toml

@@ -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"] }

crates/ctl/src/cli/destroy.rs

@@ -52,32 +52,30 @@ impl DestroyCommand {
 async fn wait_guest_destroyed(id: &str, events: EventStream) -> Result<()> {
     let mut stream = events.subscribe();
     while let Ok(event) = stream.recv().await {
-        match event {
-            Event::GuestChanged(changed) => {
-                let Some(guest) = changed.guest else {
-                    continue;
-                };
+        if let Event::GuestChanged(changed) = event {
+            let Some(guest) = changed.guest else {
+                continue;
+            };
 
-                if guest.id != id {
-                    continue;
+            if guest.id != id {
+                continue;
+            }
+
+            let Some(state) = guest.state else {
+                continue;
+            };
+
+            if let Some(ref error) = state.error_info {
+                if state.status() == GuestStatus::Failed {
+                    error!("destroy failed: {}", error.message);
+                    std::process::exit(1);
+                } else {
+                    error!("guest error: {}", error.message);
                 }
+            }
 
-                let Some(state) = guest.state else {
-                    continue;
-                };
-
-                if let Some(ref error) = state.error_info {
-                    if state.status() == GuestStatus::Failed {
-                        error!("destroy failed: {}", error.message);
-                        std::process::exit(1);
-                    } else {
-                        error!("guest error: {}", error.message);
-                    }
-                }
-
-                if state.status() == GuestStatus::Destroyed {
-                    std::process::exit(0);
-                }
+            if state.status() == GuestStatus::Destroyed {
+                std::process::exit(0);
             }
         }
     }

crates/ctl/src/cli/launch.rs

@@ -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;
                 };
@@ -158,6 +164,84 @@ async fn wait_guest_started(id: &str, events: EventStream) -> Result<()> {
                     break;
                 }
             }
+
+            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(())

crates/ctl/src/cli/watch.rs

@@ -28,11 +28,10 @@ impl WatchCommand {
         let mut stream = events.subscribe();
         loop {
             let event = stream.recv().await?;
-            match event {
-                Event::GuestChanged(changed) => {
-                    let guest = changed.guest.clone();
-                    self.print_event("guest.changed", changed, guest)?;
-                }
+
+            if let Event::GuestChanged(changed) = event {
+                let guest = changed.guest.clone();
+                self.print_event("guest.changed", changed, guest)?;
             }
         }
     }

crates/ctl/src/console.rs

@@ -69,28 +69,26 @@ impl StdioConsoleStream {
         Ok(tokio::task::spawn(async move {
             let mut stream = events.subscribe();
             while let Ok(event) = stream.recv().await {
-                match event {
-                    Event::GuestChanged(changed) => {
-                        let Some(guest) = changed.guest else {
-                            continue;
-                        };
+                if let Event::GuestChanged(changed) = event {
+                    let Some(guest) = changed.guest else {
+                        continue;
+                    };
 
-                        let Some(state) = guest.state else {
-                            continue;
-                        };
+                    let Some(state) = guest.state else {
+                        continue;
+                    };
 
-                        if guest.id != id {
-                            continue;
-                        }
+                    if guest.id != id {
+                        continue;
+                    }
 
-                        if let Some(exit_info) = state.exit_info {
-                            return Some(exit_info.code);
-                        }
+                    if let Some(exit_info) = state.exit_info {
+                        return Some(exit_info.code);
+                    }
 
-                        let status = state.status();
-                        if status == GuestStatus::Destroying || status == GuestStatus::Destroyed {
-                            return Some(10);
-                        }
+                    let status = state.status();
+                    if status == GuestStatus::Destroying || status == GuestStatus::Destroyed {
+                        return Some(10);
                     }
                 }
             }