feat: implement oci image progress (#64)

* feat: oci progress events

* feat: oci progress bars on launch

crates/oci/examples/squashify.rs

@@ -2,8 +2,10 @@ use std::{env::args, path::PathBuf};
 
 use anyhow::Result;
 use env_logger::Env;
-use krataoci::{cache::ImageCache, compiler::ImageCompiler, name::ImageName};
-use tokio::fs;
+use krataoci::{
+    cache::ImageCache, compiler::ImageCompiler, name::ImageName, progress::OciProgressContext,
+};
+use tokio::{fs, sync::broadcast};
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -18,8 +20,18 @@ async fn main() -> Result<()> {
     }
 
     let cache = ImageCache::new(&cache_dir)?;
-    let compiler = ImageCompiler::new(&cache, seed)?;
-    let info = compiler.compile(&image).await?;
+
+    let (sender, mut receiver) = broadcast::channel(1000);
+    tokio::task::spawn(async move {
+        loop {
+            let Some(_) = receiver.recv().await.ok() else {
+                break;
+            };
+        }
+    });
+    let context = OciProgressContext::new(sender);
+    let compiler = ImageCompiler::new(&cache, seed, context)?;
+    let info = compiler.compile(&image.to_string(), &image).await?;
     println!(
         "generated squashfs of {} to {}",
         image,

crates/oci/src/compiler.rs

@@ -1,6 +1,7 @@
 use crate::cache::ImageCache;
 use crate::fetch::{OciImageDownloader, OciImageLayer};
 use crate::name::ImageName;
+use crate::progress::{OciProgress, OciProgressContext, OciProgressPhase};
 use crate::registry::OciRegistryPlatform;
 use anyhow::{anyhow, Result};
 use backhand::compression::Compressor;
@@ -8,6 +9,7 @@ use backhand::{FilesystemCompressor, FilesystemWriter, NodeHeader};
 use log::{debug, trace, warn};
 use oci_spec::image::{ImageConfiguration, ImageManifest};
 use std::borrow::Cow;
+use std::collections::BTreeMap;
 use std::fs::File;
 use std::io::{BufWriter, ErrorKind, Read};
 use std::os::unix::fs::{FileTypeExt, MetadataExt, PermissionsExt};
@@ -45,14 +47,23 @@ impl ImageInfo {
 pub struct ImageCompiler<'a> {
     cache: &'a ImageCache,
     seed: Option<PathBuf>,
+    progress: OciProgressContext,
 }
 
 impl ImageCompiler<'_> {
-    pub fn new(cache: &ImageCache, seed: Option<PathBuf>) -> Result<ImageCompiler> {
-        Ok(ImageCompiler { cache, seed })
+    pub fn new(
+        cache: &ImageCache,
+        seed: Option<PathBuf>,
+        progress: OciProgressContext,
+    ) -> Result<ImageCompiler> {
+        Ok(ImageCompiler {
+            cache,
+            seed,
+            progress,
+        })
     }
 
-    pub async fn compile(&self, image: &ImageName) -> Result<ImageInfo> {
+    pub async fn compile(&self, id: &str, image: &ImageName) -> Result<ImageInfo> {
         debug!("compile image={image}");
         let mut tmp_dir = std::env::temp_dir().clone();
         tmp_dir.push(format!("krata-compile-{}", Uuid::new_v4()));
@@ -68,7 +79,7 @@ impl ImageCompiler<'_> {
         let mut squash_file = tmp_dir.clone();
         squash_file.push("image.squashfs");
         let info = self
-            .download_and_compile(image, &layer_dir, &image_dir, &squash_file)
+            .download_and_compile(id, image, &layer_dir, &image_dir, &squash_file)
             .await?;
         fs::remove_dir_all(&tmp_dir).await?;
         Ok(info)
@@ -76,15 +87,25 @@ impl ImageCompiler<'_> {
 
     async fn download_and_compile(
         &self,
+        id: &str,
         image: &ImageName,
         layer_dir: &Path,
         image_dir: &Path,
         squash_file: &Path,
     ) -> Result<ImageInfo> {
+        let mut progress = OciProgress {
+            id: id.to_string(),
+            phase: OciProgressPhase::Resolving,
+            layers: BTreeMap::new(),
+            value: 0,
+            total: 0,
+        };
+        self.progress.update(&progress);
         let downloader = OciImageDownloader::new(
             self.seed.clone(),
             layer_dir.to_path_buf(),
             OciRegistryPlatform::current(),
+            self.progress.clone(),
         );
         let resolved = downloader.resolve(image.clone()).await?;
         let cache_key = format!(
@@ -93,28 +114,44 @@ impl ImageCompiler<'_> {
         );
         let cache_digest = sha256::digest(cache_key);
 
+        progress.phase = OciProgressPhase::Complete;
+        self.progress.update(&progress);
         if let Some(cached) = self.cache.recall(&cache_digest).await? {
             return Ok(cached);
         }
 
-        let local = downloader.download(resolved).await?;
+        progress.phase = OciProgressPhase::Resolved;
+        for layer in resolved.manifest.layers() {
+            progress.add_layer(layer.digest());
+        }
+        self.progress.update(&progress);
+
+        let local = downloader.download(resolved, &mut progress).await?;
         for layer in &local.layers {
             debug!(
                 "process layer digest={} compression={:?}",
                 &layer.digest, layer.compression,
             );
-            let whiteouts = self.process_layer_whiteout(layer, image_dir).await?;
+            progress.extracting_layer(&layer.digest, 0, 0);
+            self.progress.update(&progress);
+            let (whiteouts, count) = self.process_layer_whiteout(layer, image_dir).await?;
+            progress.extracting_layer(&layer.digest, 0, count);
+            self.progress.update(&progress);
             debug!(
                 "process layer digest={} whiteouts={:?}",
                 &layer.digest, whiteouts
             );
             let mut archive = layer.archive().await?;
             let mut entries = archive.entries()?;
+            let mut completed = 0;
             while let Some(entry) = entries.next().await {
                 let mut entry = entry?;
                 let path = entry.path()?;
                 let mut maybe_whiteout_path_str =
                     path.to_str().map(|x| x.to_string()).unwrap_or_default();
+                progress.extracting_layer(&layer.digest, completed, count);
+                completed += 1;
+                self.progress.update(&progress);
                 if whiteouts.contains(&maybe_whiteout_path_str) {
                     continue;
                 }
@@ -123,10 +160,10 @@ impl ImageCompiler<'_> {
                     continue;
                 }
                 let Some(name) = path.file_name() else {
-                    return Err(anyhow!("unable to get file name"));
+                    continue;
                 };
                 let Some(name) = name.to_str() else {
-                    return Err(anyhow!("unable to get file name as string"));
+                    continue;
                 };
 
                 if name.starts_with(".wh.") {
@@ -136,6 +173,8 @@ impl ImageCompiler<'_> {
                         .await?;
                 }
             }
+            progress.extracted_layer(&layer.digest);
+            self.progress.update(&progress);
         }
 
         for layer in &local.layers {
@@ -144,31 +183,51 @@ impl ImageCompiler<'_> {
             }
         }
 
-        self.squash(image_dir, squash_file)?;
+        let image_dir_squash = image_dir.to_path_buf();
+        let squash_file_squash = squash_file.to_path_buf();
+        let progress_squash = progress.clone();
+        let progress_context = self.progress.clone();
+        progress = tokio::task::spawn_blocking(move || {
+            ImageCompiler::squash(
+                &image_dir_squash,
+                &squash_file_squash,
+                progress_squash,
+                progress_context,
+            )
+        })
+        .await??;
+
         let info = ImageInfo::new(
             squash_file.to_path_buf(),
             local.image.manifest,
             local.config,
         )?;
-        self.cache.store(&cache_digest, &info).await
+        let info = self.cache.store(&cache_digest, &info).await?;
+        progress.phase = OciProgressPhase::Complete;
+        progress.value = 0;
+        progress.total = 0;
+        self.progress.update(&progress);
+        Ok(info)
     }
 
     async fn process_layer_whiteout(
         &self,
         layer: &OciImageLayer,
         image_dir: &Path,
-    ) -> Result<Vec<String>> {
+    ) -> Result<(Vec<String>, usize)> {
         let mut whiteouts = Vec::new();
         let mut archive = layer.archive().await?;
         let mut entries = archive.entries()?;
+        let mut count = 0usize;
         while let Some(entry) = entries.next().await {
             let entry = entry?;
+            count += 1;
             let path = entry.path()?;
             let Some(name) = path.file_name() else {
-                return Err(anyhow!("unable to get file name"));
+                continue;
             };
             let Some(name) = name.to_str() else {
-                return Err(anyhow!("unable to get file name as string"));
+                continue;
             };
 
             if name.starts_with(".wh.") {
@@ -180,7 +239,7 @@ impl ImageCompiler<'_> {
                 }
             }
         }
-        Ok(whiteouts)
+        Ok((whiteouts, count))
     }
 
     async fn process_whiteout_entry(
@@ -300,7 +359,16 @@ impl ImageCompiler<'_> {
         Ok(())
     }
 
-    fn squash(&self, image_dir: &Path, squash_file: &Path) -> Result<()> {
+    fn squash(
+        image_dir: &Path,
+        squash_file: &Path,
+        mut progress: OciProgress,
+        progress_context: OciProgressContext,
+    ) -> Result<OciProgress> {
+        progress.phase = OciProgressPhase::Packing;
+        progress.total = 2;
+        progress.value = 0;
+        progress_context.update(&progress);
         let mut writer = FilesystemWriter::default();
         writer.set_compressor(FilesystemCompressor::new(Compressor::Gzip, None)?);
         let walk = WalkDir::new(image_dir).follow_links(false);
@@ -358,6 +426,10 @@ impl ImageCompiler<'_> {
             }
         }
 
+        progress.phase = OciProgressPhase::Packing;
+        progress.value = 1;
+        progress_context.update(&progress);
+
         let squash_file_path = squash_file
             .to_str()
             .ok_or_else(|| anyhow!("failed to convert squashfs string"))?;
@@ -367,7 +439,10 @@ impl ImageCompiler<'_> {
         trace!("squash generate: {}", squash_file_path);
         writer.write(&mut bufwrite)?;
         std::fs::remove_dir_all(image_dir)?;
-        Ok(())
+        progress.phase = OciProgressPhase::Packing;
+        progress.value = 2;
+        progress_context.update(&progress);
+        Ok(progress)
     }
 }
 

crates/oci/src/fetch.rs

@@ -1,3 +1,5 @@
+use crate::progress::{OciProgress, OciProgressContext, OciProgressPhase};
+
 use super::{
     name::ImageName,
     registry::{OciRegistryClient, OciRegistryPlatform},
@@ -26,6 +28,7 @@ pub struct OciImageDownloader {
     seed: Option<PathBuf>,
     storage: PathBuf,
     platform: OciRegistryPlatform,
+    progress: OciProgressContext,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -79,11 +82,13 @@ impl OciImageDownloader {
         seed: Option<PathBuf>,
         storage: PathBuf,
         platform: OciRegistryPlatform,
+        progress: OciProgressContext,
     ) -> OciImageDownloader {
         OciImageDownloader {
             seed,
             storage,
             platform,
+            progress,
         }
     }
 
@@ -208,9 +213,15 @@ impl OciImageDownloader {
         })
     }
 
-    pub async fn download(&self, image: OciResolvedImage) -> Result<OciLocalImage> {
+    pub async fn download(
+        &self,
+        image: OciResolvedImage,
+        progress: &mut OciProgress,
+    ) -> Result<OciLocalImage> {
         let config: ImageConfiguration;
 
+        progress.phase = OciProgressPhase::ConfigAcquire;
+        self.progress.update(progress);
         let mut client = OciRegistryClient::new(image.name.registry_url()?, self.platform.clone())?;
         if let Some(seeded) = self
             .load_seed_json_blob::<ImageConfiguration>(image.manifest.config())
@@ -223,9 +234,18 @@ impl OciImageDownloader {
                 .await?;
             config = serde_json::from_slice(&config_bytes)?;
         }
+        progress.phase = OciProgressPhase::LayerAcquire;
+        self.progress.update(progress);
         let mut layers = Vec::new();
         for layer in image.manifest.layers() {
-            layers.push(self.acquire_layer(&image.name, layer, &mut client).await?);
+            progress.downloading_layer(layer.digest(), 0, layer.size() as usize);
+            self.progress.update(progress);
+            layers.push(
+                self.acquire_layer(&image.name, layer, &mut client, progress)
+                    .await?,
+            );
+            progress.downloaded_layer(layer.digest());
+            self.progress.update(progress);
         }
         Ok(OciLocalImage {
             image,
@@ -239,6 +259,7 @@ impl OciImageDownloader {
         image: &ImageName,
         layer: &Descriptor,
         client: &mut OciRegistryClient,
+        progress: &mut OciProgress,
     ) -> Result<OciImageLayer> {
         debug!(
             "acquire layer digest={} size={}",
@@ -251,7 +272,15 @@ impl OciImageDownloader {
         let seeded = self.extract_seed_blob(layer, &layer_path).await?;
         if !seeded {
             let file = File::create(&layer_path).await?;
-            let size = client.write_blob_to_file(&image.name, layer, file).await?;
+            let size = client
+                .write_blob_to_file(
+                    &image.name,
+                    layer,
+                    file,
+                    Some(progress),
+                    Some(&self.progress),
+                )
+                .await?;
             if layer.size() as u64 != size {
                 return Err(anyhow!(
                     "downloaded layer size differs from size in manifest",

crates/oci/src/lib.rs

@@ -2,4 +2,5 @@ pub mod cache;
 pub mod compiler;
 pub mod fetch;
 pub mod name;
+pub mod progress;
 pub mod registry;

crates/oci/src/progress.rs

@@ -0,0 +1,98 @@
+use std::collections::BTreeMap;
+
+use tokio::sync::broadcast::Sender;
+
+#[derive(Clone, Debug)]
+pub struct OciProgress {
+    pub id: String,
+    pub phase: OciProgressPhase,
+    pub layers: BTreeMap<String, OciProgressLayer>,
+    pub value: u64,
+    pub total: u64,
+}
+
+impl OciProgress {
+    pub fn add_layer(&mut self, id: &str) {
+        self.layers.insert(
+            id.to_string(),
+            OciProgressLayer {
+                id: id.to_string(),
+                phase: OciProgressLayerPhase::Waiting,
+                value: 0,
+                total: 0,
+            },
+        );
+    }
+
+    pub fn downloading_layer(&mut self, id: &str, downloaded: usize, total: usize) {
+        if let Some(entry) = self.layers.get_mut(id) {
+            entry.phase = OciProgressLayerPhase::Downloading;
+            entry.value = downloaded as u64;
+            entry.total = total as u64;
+        }
+    }
+
+    pub fn downloaded_layer(&mut self, id: &str) {
+        if let Some(entry) = self.layers.get_mut(id) {
+            entry.phase = OciProgressLayerPhase::Downloaded;
+            entry.value = entry.total;
+        }
+    }
+
+    pub fn extracting_layer(&mut self, id: &str, extracted: usize, total: usize) {
+        if let Some(entry) = self.layers.get_mut(id) {
+            entry.phase = OciProgressLayerPhase::Extracting;
+            entry.value = extracted as u64;
+            entry.total = total as u64;
+        }
+    }
+
+    pub fn extracted_layer(&mut self, id: &str) {
+        if let Some(entry) = self.layers.get_mut(id) {
+            entry.phase = OciProgressLayerPhase::Extracted;
+            entry.value = entry.total;
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum OciProgressPhase {
+    Resolving,
+    Resolved,
+    ConfigAcquire,
+    LayerAcquire,
+    Packing,
+    Complete,
+}
+
+#[derive(Clone, Debug)]
+pub struct OciProgressLayer {
+    pub id: String,
+    pub phase: OciProgressLayerPhase,
+    pub value: u64,
+    pub total: u64,
+}
+
+#[derive(Clone, Debug)]
+pub enum OciProgressLayerPhase {
+    Waiting,
+    Downloading,
+    Downloaded,
+    Extracting,
+    Extracted,
+}
+
+#[derive(Clone)]
+pub struct OciProgressContext {
+    sender: Sender<OciProgress>,
+}
+
+impl OciProgressContext {
+    pub fn new(sender: Sender<OciProgress>) -> OciProgressContext {
+        OciProgressContext { sender }
+    }
+
+    pub fn update(&self, progress: &OciProgress) {
+        let _ = self.sender.send(progress.clone());
+    }
+}

crates/oci/src/registry.rs

@@ -7,6 +7,8 @@ use reqwest::{Client, RequestBuilder, Response, StatusCode};
 use tokio::{fs::File, io::AsyncWriteExt};
 use url::Url;
 
+use crate::progress::{OciProgress, OciProgressContext};
+
 #[derive(Clone, Debug)]
 pub struct OciRegistryPlatform {
     pub os: Os,
@@ -138,6 +140,8 @@ impl OciRegistryClient {
         name: N,
         descriptor: &Descriptor,
         mut dest: File,
+        mut progress_handle: Option<&mut OciProgress>,
+        progress_context: Option<&OciProgressContext>,
     ) -> Result<u64> {
         let url = self.url.join(&format!(
             "/v2/{}/blobs/{}",
@@ -146,9 +150,24 @@ impl OciRegistryClient {
         ))?;
         let mut response = self.call(self.agent.get(url.as_str())).await?;
         let mut size: u64 = 0;
+        let mut last_progress_size: u64 = 0;
         while let Some(chunk) = response.chunk().await? {
             dest.write_all(&chunk).await?;
             size += chunk.len() as u64;
+
+            if (size - last_progress_size) > (5 * 1024 * 1024) {
+                last_progress_size = size;
+                if let Some(progress_handle) = progress_handle.as_mut() {
+                    progress_handle.downloading_layer(
+                        descriptor.digest(),
+                        size as usize,
+                        descriptor.size() as usize,
+                    );
+                    if let Some(progress_context) = progress_context.as_ref() {
+                        progress_context.update(progress_handle);
+                    }
+                }
+            }
         }
         Ok(size)
     }