fix(yt/update): Correct the progress display in `--grouped` mode

4 files changed, 51 insertions, 8 deletions

diff --git a/crates/yt/src/cli.rs b/crates/yt/src/cli.rs
index c13693b..634e422 100644
--- a/crates/yt/src/cli.rs
+++ b/crates/yt/src/cli.rs
@@ -127,13 +127,30 @@ pub enum Command {
         #[arg(short, long)]
         max_backlog: Option<usize>,
 
+        /// How many subs were already checked.
+        ///
+        /// Only used in the progress display in combination with `--grouped`.
+        #[arg(short, long, hide = true)]
+        current_progress: Option<usize>,
+
+        /// How many subs are to be checked.
+        ///
+        /// Only used in the progress display in combination with `--grouped`.
+        #[arg(short, long, hide = true)]
+        total_number: Option<usize>,
+
         /// The subscriptions to update
         subscriptions: Vec<String>,
 
         /// Perform the updates in blocks.
         ///
         /// This works around the memory leaks in the default update invocation.
-        #[arg(short, long)]
+        #[arg(
+            short,
+            long,
+            conflicts_with = "total_number",
+            conflicts_with = "current_progress"
+        )]
         grouped: bool,
     },
 
diff --git a/crates/yt/src/main.rs b/crates/yt/src/main.rs
index 1affed9..930d269 100644
--- a/crates/yt/src/main.rs
+++ b/crates/yt/src/main.rs
@@ -162,6 +162,8 @@ async fn main() -> Result<()> {
             max_backlog,
             subscriptions,
             grouped,
+            current_progress,
+            total_number,
         } => {
             let all_subs = subscriptions::get(&app).await?;
 
@@ -177,6 +179,10 @@ async fn main() -> Result<()> {
             let max_backlog = max_backlog.unwrap_or(app.config.update.max_backlog);
 
             if grouped {
+                const CHUNK_SIZE: usize = 50;
+
+                assert!(current_progress.is_none() && total_number.is_none());
+
                 let subs = {
                     if subscriptions.is_empty() {
                         all_subs.0.into_iter().map(|sub| sub.0).collect()
@@ -185,7 +191,9 @@ async fn main() -> Result<()> {
                     }
                 };
 
-                for chunk in subs.chunks(50) {
+                let total_number = subs.len();
+                let mut current_progress = 0;
+                for chunk in subs.chunks(CHUNK_SIZE) {
                     info!(
                         "$ yt update {}",
                         chunk
@@ -199,15 +207,19 @@ async fn main() -> Result<()> {
                         current_exe().context("Failed to get the current exe to re-execute")?,
                     )
                     .arg("update")
+                    .args(["--current-progress", current_progress.to_string().as_str()])
+                    .args(["--total-number", total_number.to_string().as_str()])
                     .args(chunk)
                     .status()?;
 
                     if !status.success() {
                         bail!("grouped yt update: Child process failed.");
                     }
+
+                    current_progress += CHUNK_SIZE;
                 }
             } else {
-                update::update(&app, max_backlog, subscriptions).await?;
+                update::update(&app, max_backlog, subscriptions, total_number, current_progress).await?;
             }
         }
         Command::Subscriptions { cmd } => match cmd {
diff --git a/crates/yt/src/update/mod.rs b/crates/yt/src/update/mod.rs
index 36c8c81..d866882 100644
--- a/crates/yt/src/update/mod.rs
+++ b/crates/yt/src/update/mod.rs
@@ -36,6 +36,8 @@ pub async fn update(
     app: &App,
     max_backlog: usize,
     subscription_names_to_update: Vec<String>,
+    total_number: Option<usize>,
+    current_progress: Option<usize>,
 ) -> Result<()> {
     let subscriptions = subscriptions::get(app).await?;
 
@@ -54,7 +56,9 @@ pub async fn update(
     let hashes = get_all_hashes(app).await?;
 
     let updater = Updater::new(max_backlog, hashes);
-    updater.update(app, subs).await?;
+    updater
+        .update(app, subs, total_number, current_progress)
+        .await?;
 
     Ok(())
 }
diff --git a/crates/yt/src/update/updater.rs b/crates/yt/src/update/updater.rs
index 21e30ed..04bcaa1 100644
--- a/crates/yt/src/update/updater.rs
+++ b/crates/yt/src/update/updater.rs
@@ -10,7 +10,7 @@
 
 use std::{
     io::{Write, stderr},
-    sync::atomic::AtomicUsize,
+    sync::atomic::{AtomicUsize, Ordering},
 };
 
 use anyhow::{Context, Result};
@@ -49,8 +49,18 @@ impl Updater {
         }
     }
 
-    pub(super) async fn update(self, app: &App, subscriptions: Vec<Subscription>) -> Result<()> {
-        let total_number = subscriptions.len();
+    pub(super) async fn update(
+        self,
+        app: &App,
+        subscriptions: Vec<Subscription>,
+        total_number: Option<usize>,
+        current_progress: Option<usize>,
+    ) -> Result<()> {
+        let total_number = total_number.unwrap_or(subscriptions.len());
+
+        if let Some(current_progress) = current_progress {
+            REACHED_NUMBER.store(current_progress, Ordering::Relaxed);
+        }
 
         let mut stream = stream::iter(subscriptions)
             .map(|sub| self.get_new_entries(sub, total_number))
@@ -102,7 +112,7 @@ impl Updater {
                         move_to_col(1);
                         eprint!(
                             "({}/{total_number}) Checking playlist {}...",
-                            REACHED_NUMBER.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
+                            REACHED_NUMBER.fetch_add(1, Ordering::Relaxed),
                             sub.name
                         );
                         move_to_col(1);