Merge branch '351786-bbm-store-affected-rows-in-metrics' into 'master'

Store number of affected rows in metrics for batched background migrations

See merge request gitlab-org/gitlab!82960

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

spec/lib/gitlab/database/background_migration/batch_metrics_spec.rb

@@ -10,7 +10,6 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchMetrics do
       expect(batch_metrics.timings).to be_empty
 
       expect(Gitlab::Metrics::System).to receive(:monotonic_time)
-        .exactly(6).times
         .and_return(0.0, 111.0, 200.0, 290.0, 300.0, 410.0)
 
       batch_metrics.time_operation(:my_label) do
@@ -28,4 +27,33 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchMetrics do
       expect(batch_metrics.timings).to eq(my_label: [111.0, 110.0], my_other_label: [90.0])
     end
   end
+
+  describe '#instrument_operation' do
+    it 'tracks duration and affected rows' do
+      expect(batch_metrics.timings).to be_empty
+      expect(batch_metrics.affected_rows).to be_empty
+
+      expect(Gitlab::Metrics::System).to receive(:monotonic_time)
+        .and_return(0.0, 111.0, 200.0, 290.0, 300.0, 410.0, 420.0, 450.0)
+
+      batch_metrics.instrument_operation(:my_label) do
+        3
+      end
+
+      batch_metrics.instrument_operation(:my_other_label) do
+        42
+      end
+
+      batch_metrics.instrument_operation(:my_label) do
+        2
+      end
+
+      batch_metrics.instrument_operation(:my_other_label) do
+        :not_an_integer
+      end
+
+      expect(batch_metrics.timings).to eq(my_label: [111.0, 110.0], my_other_label: [90.0, 30.0])
+      expect(batch_metrics.affected_rows).to eq(my_label: [3, 2], my_other_label: [42])
+    end
+  end
 end