feat: implement oci image progress (#64)

* feat: oci progress events

* feat: oci progress bars on launch

crates/daemon/Cargo.toml

@@ -18,6 +18,7 @@ clap = { workspace = true }
 env_logger = { workspace = true }
 futures = { workspace = true }
 krata = { path = "../krata", version = "^0.0.8" }
+krata-oci = { path = "../oci", version = "^0.0.8" }
 krata-runtime = { path = "../runtime", version = "^0.0.8" }
 log = { workspace = true }
 prost = { workspace = true }

crates/daemon/bin/daemon.rs

@@ -3,7 +3,6 @@ use clap::Parser;
 use env_logger::Env;
 use krata::dial::ControlDialAddress;
 use kratad::Daemon;
-use kratart::Runtime;
 use log::LevelFilter;
 use std::{
     str::FromStr,
@@ -27,8 +26,8 @@ async fn main() -> Result<()> {
 
     let args = DaemonCommand::parse();
     let addr = ControlDialAddress::from_str(&args.listen)?;
-    let runtime = Runtime::new(args.store.clone()).await?;
-    let mut daemon = Daemon::new(args.store.clone(), runtime).await?;
+
+    let mut daemon = Daemon::new(args.store.clone()).await?;
     daemon.listen(addr).await?;
     Ok(())
 }

crates/daemon/src/event.rs

@@ -9,6 +9,7 @@ use krata::{
     idm::protocol::{idm_event::Event, IdmEvent},
     v1::common::{GuestExitInfo, GuestState, GuestStatus},
 };
+use krataoci::progress::OciProgress;
 use log::{error, warn};
 use tokio::{
     select,
@@ -21,7 +22,7 @@ use tokio::{
 };
 use uuid::Uuid;
 
-use crate::{db::GuestStore, idm::DaemonIdmHandle};
+use crate::{db::GuestStore, idm::DaemonIdmHandle, oci::convert_oci_progress};
 
 pub type DaemonEvent = krata::v1::control::watch_events_reply::Event;
 
@@ -52,7 +53,8 @@ pub struct DaemonEventGenerator {
     idms: HashMap<u32, (Uuid, JoinHandle<()>)>,
     idm_sender: Sender<(u32, IdmEvent)>,
     idm_receiver: Receiver<(u32, IdmEvent)>,
-    _event_sender: broadcast::Sender<DaemonEvent>,
+    oci_progress_receiver: broadcast::Receiver<OciProgress>,
+    event_sender: broadcast::Sender<DaemonEvent>,
 }
 
 impl DaemonEventGenerator {
@@ -60,6 +62,7 @@ impl DaemonEventGenerator {
         guests: GuestStore,
         guest_reconciler_notify: Sender<Uuid>,
         idm: DaemonIdmHandle,
+        oci_progress_receiver: broadcast::Receiver<OciProgress>,
     ) -> Result<(DaemonEventContext, DaemonEventGenerator)> {
         let (sender, _) = broadcast::channel(EVENT_CHANNEL_QUEUE_LEN);
         let (idm_sender, idm_receiver) = channel(IDM_EVENT_CHANNEL_QUEUE_LEN);
@@ -71,55 +74,54 @@ impl DaemonEventGenerator {
             idms: HashMap::new(),
             idm_sender,
             idm_receiver,
-            _event_sender: sender.clone(),
+            oci_progress_receiver,
+            event_sender: sender.clone(),
         };
         let context = DaemonEventContext { sender };
         Ok((context, generator))
     }
 
     async fn handle_feed_event(&mut self, event: &DaemonEvent) -> Result<()> {
-        match event {
-            DaemonEvent::GuestChanged(changed) => {
-                let Some(ref guest) = changed.guest else {
-                    return Ok(());
-                };
+        if let DaemonEvent::GuestChanged(changed) = event {
+            let Some(ref guest) = changed.guest else {
+                return Ok(());
+            };
 
-                let Some(ref state) = guest.state else {
-                    return Ok(());
-                };
+            let Some(ref state) = guest.state else {
+                return Ok(());
+            };
 
-                let status = state.status();
-                let id = Uuid::from_str(&guest.id)?;
-                let domid = state.domid;
-                match status {
-                    GuestStatus::Started => {
-                        if let Entry::Vacant(e) = self.idms.entry(domid) {
-                            let client = self.idm.client(domid).await?;
-                            let mut receiver = client.subscribe().await?;
-                            let sender = self.idm_sender.clone();
-                            let task = tokio::task::spawn(async move {
-                                loop {
-                                    let Ok(event) = receiver.recv().await else {
-                                        break;
-                                    };
+            let status = state.status();
+            let id = Uuid::from_str(&guest.id)?;
+            let domid = state.domid;
+            match status {
+                GuestStatus::Started => {
+                    if let Entry::Vacant(e) = self.idms.entry(domid) {
+                        let client = self.idm.client(domid).await?;
+                        let mut receiver = client.subscribe().await?;
+                        let sender = self.idm_sender.clone();
+                        let task = tokio::task::spawn(async move {
+                            loop {
+                                let Ok(event) = receiver.recv().await else {
+                                    break;
+                                };
 
-                                    if let Err(error) = sender.send((domid, event)).await {
-                                        warn!("unable to deliver idm event: {}", error);
-                                    }
+                                if let Err(error) = sender.send((domid, event)).await {
+                                    warn!("unable to deliver idm event: {}", error);
                                 }
-                            });
-                            e.insert((id, task));
-                        }
+                            }
+                        });
+                        e.insert((id, task));
                     }
-
-                    GuestStatus::Destroyed => {
-                        if let Some((_, handle)) = self.idms.remove(&domid) {
-                            handle.abort();
-                        }
-                    }
-
-                    _ => {}
                 }
+
+                GuestStatus::Destroyed => {
+                    if let Some((_, handle)) = self.idms.remove(&domid) {
+                        handle.abort();
+                    }
+                }
+
+                _ => {}
             }
         }
         Ok(())
@@ -148,6 +150,17 @@ impl DaemonEventGenerator {
         Ok(())
     }
 
+    async fn handle_oci_progress_event(&mut self, progress: OciProgress) -> Result<()> {
+        let Some(_) = Uuid::from_str(&progress.id).ok() else {
+            return Ok(());
+        };
+
+        let event = convert_oci_progress(progress);
+        self.event_sender.send(DaemonEvent::OciProgress(event))?;
+
+        Ok(())
+    }
+
     async fn evaluate(&mut self) -> Result<()> {
         select! {
             x = self.idm_receiver.recv() => match x {
@@ -168,6 +181,14 @@ impl DaemonEventGenerator {
                 Err(error) => {
                     Err(error.into())
                 }
+            },
+            x = self.oci_progress_receiver.recv() => match x {
+                Ok(event) => {
+                    self.handle_oci_progress_event(event).await
+                },
+                Err(error) => {
+                    Err(error.into())
+                }
             }
         }
     }

crates/daemon/src/lib.rs

@@ -7,12 +7,16 @@ use db::GuestStore;
 use event::{DaemonEventContext, DaemonEventGenerator};
 use idm::{DaemonIdm, DaemonIdmHandle};
 use krata::{dial::ControlDialAddress, v1::control::control_service_server::ControlServiceServer};
+use krataoci::progress::OciProgressContext;
 use kratart::Runtime;
 use log::info;
 use reconcile::guest::GuestReconciler;
 use tokio::{
     net::UnixListener,
-    sync::mpsc::{channel, Sender},
+    sync::{
+        broadcast,
+        mpsc::{channel, Sender},
+    },
     task::JoinHandle,
 };
 use tokio_stream::wrappers::UnixListenerStream;
@@ -25,6 +29,7 @@ pub mod db;
 pub mod event;
 pub mod idm;
 pub mod metrics;
+pub mod oci;
 pub mod reconcile;
 
 pub struct Daemon {
@@ -39,9 +44,14 @@ pub struct Daemon {
 }
 
 const GUEST_RECONCILER_QUEUE_LEN: usize = 1000;
+const OCI_PROGRESS_QUEUE_LEN: usize = 1000;
 
 impl Daemon {
-    pub async fn new(store: String, runtime: Runtime) -> Result<Self> {
+    pub async fn new(store: String) -> Result<Self> {
+        let (oci_progress_sender, oci_progress_receiver) =
+            broadcast::channel(OCI_PROGRESS_QUEUE_LEN);
+        let runtime =
+            Runtime::new(OciProgressContext::new(oci_progress_sender), store.clone()).await?;
         let guests_db_path = format!("{}/guests.db", store);
         let guests = GuestStore::open(&PathBuf::from(guests_db_path))?;
         let (guest_reconciler_notify, guest_reconciler_receiver) =
@@ -50,9 +60,13 @@ impl Daemon {
         let idm = idm.launch().await?;
         let console = DaemonConsole::new().await?;
         let console = console.launch().await?;
-        let (events, generator) =
-            DaemonEventGenerator::new(guests.clone(), guest_reconciler_notify.clone(), idm.clone())
-                .await?;
+        let (events, generator) = DaemonEventGenerator::new(
+            guests.clone(),
+            guest_reconciler_notify.clone(),
+            idm.clone(),
+            oci_progress_receiver,
+        )
+        .await?;
         let runtime_for_reconciler = runtime.dupe().await?;
         let guest_reconciler = GuestReconciler::new(
             guests.clone(),

crates/daemon/src/oci.rs

@@ -0,0 +1,42 @@
+use krata::v1::control::{
+    OciProgressEvent, OciProgressEventLayer, OciProgressEventLayerPhase, OciProgressEventPhase,
+};
+use krataoci::progress::{OciProgress, OciProgressLayer, OciProgressLayerPhase, OciProgressPhase};
+
+fn convert_oci_layer_progress(layer: OciProgressLayer) -> OciProgressEventLayer {
+    OciProgressEventLayer {
+        id: layer.id,
+        phase: match layer.phase {
+            OciProgressLayerPhase::Waiting => OciProgressEventLayerPhase::Waiting,
+            OciProgressLayerPhase::Downloading => OciProgressEventLayerPhase::Downloading,
+            OciProgressLayerPhase::Downloaded => OciProgressEventLayerPhase::Downloaded,
+            OciProgressLayerPhase::Extracting => OciProgressEventLayerPhase::Extracting,
+            OciProgressLayerPhase::Extracted => OciProgressEventLayerPhase::Extracted,
+        }
+        .into(),
+        value: layer.value,
+        total: layer.total,
+    }
+}
+
+pub fn convert_oci_progress(oci: OciProgress) -> OciProgressEvent {
+    OciProgressEvent {
+        guest_id: oci.id,
+        phase: match oci.phase {
+            OciProgressPhase::Resolving => OciProgressEventPhase::Resolving,
+            OciProgressPhase::Resolved => OciProgressEventPhase::Resolved,
+            OciProgressPhase::ConfigAcquire => OciProgressEventPhase::ConfigAcquire,
+            OciProgressPhase::LayerAcquire => OciProgressEventPhase::LayerAcquire,
+            OciProgressPhase::Packing => OciProgressEventPhase::Packing,
+            OciProgressPhase::Complete => OciProgressEventPhase::Complete,
+        }
+        .into(),
+        layers: oci
+            .layers
+            .into_values()
+            .map(convert_oci_layer_progress)
+            .collect::<Vec<_>>(),
+        value: oci.value,
+        total: oci.total,
+    }
+}