From d6312117c2857d3424a6f6888411f33642f0caa7 Mon Sep 17 00:00:00 2001
From: Colin Walters <walters@verbum.org>
Date: Wed, 21 Aug 2024 15:18:18 -0400
Subject: [PATCH] sysroot: Add total layer count to output

It's useful to see in the progress output how many layers
there are to fetch.

This is similar to
https://github.com/containers/bootc/pull/524/commits/6eb57180a3c82393b70edb702fb9c4db53e54c32
which ended up being totally reworked in a nicer way in
https://github.com/containers/bootc/pull/655/commits/d8b5df2f378b54c3a60ccaa673a60a7bbca5d653

But doing the latter would require nontrivial changes to our
DBus API around status and progress reporting...and I'd
like to think about how we tackle that more generally
in e.g. https://github.com/containers/skopeo/issues/658

Closes: https://github.com/coreos/rpm-ostree/issues/5024
---
 rust/src/sysroot_upgrade.rs | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/rust/src/sysroot_upgrade.rs b/rust/src/sysroot_upgrade.rs
index 89dcd76cff..8e6ea5a25c 100644
--- a/rust/src/sysroot_upgrade.rs
+++ b/rust/src/sysroot_upgrade.rs
@@ -59,12 +59,14 @@ fn layer_counts<'a>(layers: impl Iterator<Item = &'a ManifestLayerState>) -> (u3
 // Output when we begin/end fetching a layer.  Ideally, we'd handle
 // byte-level progress here too, but that requires some more sophisticated
 // binding with the rpmostree-output.h via cxx.rs.
-async fn layer_progress_print(mut r: Receiver<ImportProgress>) {
+async fn layer_progress_print(mut r: Receiver<ImportProgress>, total_to_fetch: u32) {
     // This is just used to hold a reference to the task.
     #[allow(unused_variables, unused_assignments)]
     let mut task = None;
+    let mut n_fetched = 0u64;
     while let Some(v) = r.recv().await {
-        let msg = ostree_ext::cli::layer_progress_format(&v);
+        let mut msg = ostree_ext::cli::layer_progress_format(&v);
+        msg.insert_str(0, &format!("[{n_fetched}/{total_to_fetch}] "));
         tracing::debug!("layer progress: {msg}");
         match v {
             ImportProgress::OstreeChunkStarted(_) => {
@@ -73,6 +75,7 @@ async fn layer_progress_print(mut r: Receiver<ImportProgress>) {
             }
             ImportProgress::OstreeChunkCompleted(_) => {
                 assert!(task.take().is_some());
+                n_fetched += 1;
             }
             ImportProgress::DerivedLayerStarted(_) => {
                 assert!(task.is_none());
@@ -80,6 +83,7 @@ async fn layer_progress_print(mut r: Receiver<ImportProgress>) {
             }
             ImportProgress::DerivedLayerCompleted(_) => {
                 assert!(task.take().is_some());
+                n_fetched += 1;
             }
         }
     }
@@ -127,6 +131,7 @@ async fn pull_container_async(
         .ostree_layers
         .iter()
         .chain(std::iter::once(&prep.ostree_commit_layer));
+    let mut total_to_fetch = 0;
     let (stored, (n_to_fetch, size_to_fetch)) = layer_counts(ostree_layers);
     if stored > 0 {
         output_message(&format!("ostree chunk layers already present: {stored}"));
@@ -136,6 +141,7 @@ async fn pull_container_async(
         output_message(&format!(
             "ostree chunk layers needed: {n_to_fetch} ({size})"
         ));
+        total_to_fetch += n_to_fetch;
     }
     let (stored, (n_to_fetch, size_to_fetch)) = layer_counts(prep.layers.iter());
     if stored > 0 {
@@ -144,12 +150,13 @@ async fn pull_container_async(
     if n_to_fetch > 0 {
         let size = glib::format_size(size_to_fetch);
         output_message(&format!("custom layers needed: {n_to_fetch} ({size})"));
+        total_to_fetch += n_to_fetch;
     }
     let local = tokio::task::LocalSet::new();
     let import = local
         .run_until(async move {
             let _progress_printer =
-                tokio::task::spawn_local(async move { layer_progress_print(layer_progress).await });
+                tokio::task::spawn_local(layer_progress_print(layer_progress, total_to_fetch));
             imp.import(prep).await
         })
         .await;