Merge pull request #278 from cgwalters/importer-progress

container: Add support for layer fetch notifications

Author: cgwalters

lib/src/cli.rs

@@ -15,10 +15,12 @@ use std::convert::TryFrom;
 use std::ffi::OsString;
 use std::path::PathBuf;
 use structopt::StructOpt;
+use tokio::sync::mpsc::Receiver;
 use tokio_stream::StreamExt;
 
 use crate::commit::container_commit;
 use crate::container as ostree_container;
+use crate::container::store::{ImportProgress, PreparedImport};
 use crate::container::{Config, ImageReference, OstreeImageReference, UnencapsulateOptions};
 use ostree_container::store::{ImageImporter, PrepareResult};
 use ostree_container::UnencapsulationProgress;
@@ -352,6 +354,46 @@ enum ProgressOrFinish {
     Finished(Result<ostree_container::Import>),
 }
 
+/// Render an import progress notification as a string.
+pub fn layer_progress_format(p: &ImportProgress) -> String {
+    let (starting, s, layer) = match p {
+        ImportProgress::OstreeChunkStarted(v) => (true, "ostree chunk", v),
+        ImportProgress::OstreeChunkCompleted(v) => (false, "ostree chunk", v),
+        ImportProgress::DerivedLayerStarted(v) => (true, "layer", v),
+        ImportProgress::DerivedLayerCompleted(v) => (false, "layer", v),
+    };
+    // podman outputs 12 characters of digest, let's add 7 for `sha256:`.
+    let short_digest = layer.digest().chars().take(12 + 7).collect::<String>();
+    if starting {
+        let size = glib::format_size(layer.size() as u64);
+        format!("Fetching {s} {short_digest} ({size})")
+    } else {
+        format!("Fetched {s} {short_digest}")
+    }
+}
+
+async fn handle_layer_progress_print(mut r: Receiver<ImportProgress>) {
+    while let Some(v) = r.recv().await {
+        println!("{}", layer_progress_format(&v));
+    }
+}
+
+fn print_layer_status(prep: &PreparedImport) {
+    let (stored, to_fetch, to_fetch_size) =
+        prep.all_layers()
+            .fold((0u32, 0u32, 0u64), |(stored, to_fetch, sz), v| {
+                if v.commit.is_some() {
+                    (stored + 1, to_fetch, sz)
+                } else {
+                    (stored, to_fetch + 1, sz + v.size())
+                }
+            });
+    if to_fetch > 0 {
+        let size = crate::glib::format_size(to_fetch_size);
+        println!("layers stored: {stored} needed: {to_fetch} ({size})");
+    }
+}
+
 /// Import a container image with an encapsulated ostree commit.
 async fn container_import(
     repo: &ostree::Repo,
@@ -455,22 +497,20 @@ async fn container_store(
     proxyopts: ContainerProxyOpts,
 ) -> Result<()> {
     let mut imp = ImageImporter::new(repo, imgref, proxyopts.into()).await?;
+    let layer_progress = imp.request_progress();
     let prep = match imp.prepare().await? {
         PrepareResult::AlreadyPresent(c) => {
             println!("No changes in {} => {}", imgref, c.merge_commit);
             return Ok(());
         }
         PrepareResult::Ready(r) => r,
     };
-    for layer in prep.all_layers() {
-        if layer.commit.is_some() {
-            println!("Using layer: {}", layer.digest());
-        } else {
-            let size = crate::glib::format_size(layer.size());
-            println!("Downloading layer: {} ({})", layer.digest(), size);
-        }
-    }
-    let import = imp.import(prep).await?;
+    print_layer_status(&prep);
+    let progress_printer =
+        tokio::task::spawn(async move { handle_layer_progress_print(layer_progress).await });
+    let import = imp.import(prep).await;
+    let _ = progress_printer.await;
+    let import = import?;
     let commit = &repo.load_commit(&import.merge_commit)?.0;
     let commit_meta = &glib::VariantDict::new(Some(&commit.child_value(0)));
     let filtered = commit_meta.lookup::<ostree_container::store::MetaFilteredData>(

lib/src/container/store.rs

@@ -16,6 +16,7 @@ use ostree::{gio, glib};
 use std::collections::HashMap;
 use std::iter::FromIterator;
 use std::sync::{Arc, Mutex};
+use tokio::sync::mpsc::{Receiver, Sender};
 
 /// Configuration for the proxy.
 ///
@@ -54,6 +55,19 @@ fn ref_for_image(l: &ImageReference) -> Result<String> {
     refescape::prefix_escape_for_ref(IMAGE_PREFIX, &l.to_string())
 }
 
+/// Sent across a channel to track start and end of a container fetch.
+#[derive(Debug)]
+pub enum ImportProgress {
+    /// Started fetching this layer.
+    OstreeChunkStarted(Descriptor),
+    /// Successfully completed the fetch of this layer.
+    OstreeChunkCompleted(Descriptor),
+    /// Started fetching this layer.
+    DerivedLayerStarted(Descriptor),
+    /// Successfully completed the fetch of this layer.
+    DerivedLayerCompleted(Descriptor),
+}
+
 /// State of an already pulled layered image.
 #[derive(Debug, PartialEq, Eq)]
 pub struct LayeredImageState {
@@ -95,6 +109,8 @@ pub struct ImageImporter {
     imgref: OstreeImageReference,
     target_imgref: Option<OstreeImageReference>,
     pub(crate) proxy_img: OpenedImage,
+
+    layer_progress: Option<Sender<ImportProgress>>,
 }
 
 /// Result of invoking [`LayeredImageImporter::prepare`].
@@ -274,6 +290,7 @@ impl ImageImporter {
             proxy_img,
             target_imgref: None,
             imgref: imgref.clone(),
+            layer_progress: None,
         })
     }
 
@@ -286,6 +303,14 @@ impl ImageImporter {
         self.prepare_internal(false).await
     }
 
+    /// Create a channel receiver that will get notifications for layer fetches.
+    pub fn request_progress(&mut self) -> Receiver<ImportProgress> {
+        assert!(self.layer_progress.is_none());
+        let (s, r) = tokio::sync::mpsc::channel(2);
+        self.layer_progress = Some(s);
+        r
+    }
+
     /// Determine if there is a new manifest, and if so return its digest.
     #[context("Fetching manifest")]
     pub(crate) async fn prepare_internal(&mut self, verify_layers: bool) -> Result<PrepareResult> {
@@ -397,6 +422,10 @@ impl ImageImporter {
             if layer.commit.is_some() {
                 continue;
             }
+            if let Some(p) = self.layer_progress.as_ref() {
+                p.send(ImportProgress::OstreeChunkStarted(layer.layer.clone()))
+                    .await?;
+            }
             let (blob, driver) =
                 fetch_layer_decompress(&mut self.proxy, &self.proxy_img, &layer.layer).await?;
             let blob = super::unencapsulate::ProgressReader {
@@ -425,8 +454,18 @@ impl ImageImporter {
                 });
             let commit = super::unencapsulate::join_fetch(import_task, driver).await?;
             layer.commit = commit;
+            if let Some(p) = self.layer_progress.as_ref() {
+                p.send(ImportProgress::OstreeChunkCompleted(layer.layer.clone()))
+                    .await?;
+            }
         }
         if import.ostree_commit_layer.commit.is_none() {
+            if let Some(p) = self.layer_progress.as_ref() {
+                p.send(ImportProgress::OstreeChunkStarted(
+                    import.ostree_commit_layer.layer.clone(),
+                ))
+                .await?;
+            }
             let (blob, driver) = fetch_layer_decompress(
                 &mut self.proxy,
                 &self.proxy_img,
@@ -457,6 +496,12 @@ impl ImageImporter {
                 });
             let commit = super::unencapsulate::join_fetch(import_task, driver).await?;
             import.ostree_commit_layer.commit = Some(commit);
+            if let Some(p) = self.layer_progress.as_ref() {
+                p.send(ImportProgress::OstreeChunkCompleted(
+                    import.ostree_commit_layer.layer.clone(),
+                ))
+                .await?;
+            }
         };
         Ok(())
     }
@@ -503,6 +548,10 @@ impl ImageImporter {
                 tracing::debug!("Reusing fetched commit {}", c);
                 layer_commits.push(c.to_string());
             } else {
+                if let Some(p) = self.layer_progress.as_ref() {
+                    p.send(ImportProgress::DerivedLayerStarted(layer.layer.clone()))
+                        .await?;
+                }
                 let (blob, driver) = super::unencapsulate::fetch_layer_decompress(
                     &mut proxy,
                     &self.proxy_img,
@@ -525,6 +574,10 @@ impl ImageImporter {
                     let filtered = HashMap::from_iter(r.filtered.into_iter());
                     layer_filtered_content.insert(layer.digest().to_string(), filtered);
                 }
+                if let Some(p) = self.layer_progress.as_ref() {
+                    p.send(ImportProgress::DerivedLayerCompleted(layer.layer.clone()))
+                        .await?;
+                }
             }
         }