Add latest changes from gitlab-org/gitlab@14-10-stable-eev14.10.0-rc42

1 files changed, 13 insertions, 2 deletions

diff --git a/lib/gitlab/database/background_migration/batch_metrics.rb b/lib/gitlab/database/background_migration/batch_metrics.rb
index 3e6d7ac3c9f..14fe0c14c24 100644
--- a/lib/gitlab/database/background_migration/batch_metrics.rb
+++ b/lib/gitlab/database/background_migration/batch_metrics.rb
@@ -5,17 +5,24 @@ module Gitlab
     module BackgroundMigration
       class BatchMetrics
         attr_reader :timings
+        attr_reader :affected_rows
 
         def initialize
           @timings = {}
+          @affected_rows = {}
         end
 
-        def time_operation(label)
+        def time_operation(label, &blk)
+          instrument_operation(label, instrument_affected_rows: false, &blk)
+        end
+
+        def instrument_operation(label, instrument_affected_rows: true)
           start_time = monotonic_time
 
-          yield
+          count = yield
 
           timings_for_label(label) << monotonic_time - start_time
+          affected_rows_for_label(label) << count if instrument_affected_rows && count.is_a?(Integer)
         end
 
         private
@@ -24,6 +31,10 @@ module Gitlab
           timings[label] ||= []
         end
 
+        def affected_rows_for_label(label)
+          affected_rows[label] ||= []
+        end
+
         def monotonic_time
           Gitlab::Metrics::System.monotonic_time
         end