Fix issues with tootctl's parallelization and progress reporting (#12093)

1 files changed, 25 insertions, 11 deletions

diff --git a/lib/mastodon/cli_helper.rb b/lib/mastodon/cli_helper.rb
index da7348349..c2950dffa 100644
--- a/lib/mastodon/cli_helper.rb
+++ b/lib/mastodon/cli_helper.rb
@@ -15,6 +15,11 @@ module Mastodon
     end
 
     def parallelize_with_progress(scope)
+      if options[:concurrency] < 1
+        say('Cannot run with this concurrency setting, must be at least 1', :red)
+        exit(1)
+      end
+
       ActiveRecord::Base.configurations[Rails.env]['pool'] = options[:concurrency]
 
       progress  = create_progress_bar(scope.count)
@@ -27,17 +32,26 @@ module Mastodon
 
         items.each do |item|
           futures << Concurrent::Future.execute(executor: pool) do
-            ActiveRecord::Base.connection_pool.with_connection do
-              begin
-                progress.log("Processing #{item.id}") if options[:verbose]
-
-                result = yield(item)
-                aggregate.increment(result) if result.is_a?(Integer)
-              rescue => e
-                progress.log pastel.red("Error processing #{item.id}: #{e}")
-              ensure
-                progress.increment
+            begin
+              if !progress.total.nil? && progress.progress + 1 > progress.total
+                # The number of items has changed between start and now,
+                # since there is no good way to predict the final count from
+                # here, just change the progress bar to an indeterminate one
+
+                progress.total = nil
               end
+
+              progress.log("Processing #{item.id}") if options[:verbose]
+
+              result = ActiveRecord::Base.connection_pool.with_connection do
+                yield(item)
+              end
+
+              aggregate.increment(result) if result.is_a?(Integer)
+            rescue => e
+              progress.log pastel.red("Error processing #{item.id}: #{e}")
+            ensure
+              progress.increment
             end
           end
         end
@@ -46,7 +60,7 @@ module Mastodon
         futures.map(&:value)
       end
 
-      progress.finish
+      progress.stop
 
       [total.value, aggregate.value]
     end