Record transaction duration during migrations

- Records, during migrations, the duration of each transaction

Related to https://gitlab.com/gitlab-org/database-team/gitlab-com-database-testing/-/issues/33

Changelog: added

lib/gitlab/database.rb

@@ -263,14 +263,28 @@ module Gitlab
         # A patch over ActiveRecord::Base.transaction that provides
         # observability into transactional methods.
         def transaction(**options, &block)
-          if options[:requires_new] && connection.transaction_open?
-            ::Gitlab::Database::Metrics.subtransactions_increment(self.name)
-          end
+          transaction_type = get_transaction_type(connection.transaction_open?, options[:requires_new])
+
+          ::Gitlab::Database::Metrics.subtransactions_increment(self.name) if transaction_type == :sub_transaction
+
+          payload = { connection: connection, transaction_type: transaction_type }
 
-          ActiveSupport::Notifications.instrument('transaction.active_record', { connection: connection }) do
+          ActiveSupport::Notifications.instrument('transaction.active_record', payload) do
             super(**options, &block)
           end
         end
+
+        private
+
+        def get_transaction_type(transaction_open, requires_new_flag)
+          if transaction_open
+            return :sub_transaction if requires_new_flag
+
+            return :fake_transaction
+          end
+
+          :real_transaction
+        end
       end
     end
 

lib/gitlab/database/migrations/observers.rb

@@ -9,7 +9,8 @@ module Gitlab
             TotalDatabaseSizeChange,
             QueryStatistics,
             QueryLog,
-            QueryDetails
+            QueryDetails,
+            TransactionDuration
           ]
         end
       end

lib/gitlab/database/migrations/observers/transaction_duration.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        class TransactionDuration < MigrationObserver
+          def before
+            file_path = File.join(output_dir, "#{observation.version}_#{observation.name}-transaction-duration.json")
+            @file = File.open(file_path, 'wb')
+            @writer = Oj::StreamWriter.new(@file, {})
+            @writer.push_array
+            @subscriber = ActiveSupport::Notifications.subscribe('transaction.active_record') do |*args|
+              record_sql_event(*args)
+            end
+          end
+
+          def after
+            ActiveSupport::Notifications.unsubscribe(@subscriber)
+            @writer.pop_all
+            @writer.flush
+            @file.close
+          end
+
+          def record
+            # no-op
+          end
+
+          def record_sql_event(_name, started, finished, _unique_id, payload)
+            return if payload[:transaction_type] == :fake_transaction
+
+            @writer.push_value({
+                                 start_time: started.iso8601(6),
+                                 end_time: finished.iso8601(6),
+                                 transaction_type: payload[:transaction_type]
+                               })
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migrations/observers/transaction_duration_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::Observers::TransactionDuration do
+  subject(:transaction_duration_observer) { described_class.new(observation, directory_path) }
+
+  let(:observation) { Gitlab::Database::Migrations::Observation.new(migration_version, migration_name) }
+  let(:directory_path) { Dir.mktmpdir }
+  let(:log_file) { "#{directory_path}/#{migration_version}_#{migration_name}-transaction-duration.json" }
+  let(:transaction_duration) { Gitlab::Json.parse(File.read(log_file)) }
+  let(:migration_version) { 20210422152437 }
+  let(:migration_name) { 'test' }
+
+  after do
+    FileUtils.remove_entry(directory_path)
+  end
+
+  it 'records real and sub transactions duration', :delete do
+    observe
+
+    entry = transaction_duration[0]
+    start_time, end_time, transaction_type = entry.values_at('start_time', 'end_time', 'transaction_type')
+    start_time = DateTime.parse(start_time)
+    end_time = DateTime.parse(end_time)
+
+    aggregate_failures do
+      expect(transaction_duration.size).to eq(3)
+      expect(start_time).to be_before(end_time)
+      expect(transaction_type).not_to be_nil
+    end
+  end
+
+  context 'when there are sub-transactions' do
+    it 'records transaction duration' do
+      observe_sub_transaction
+
+      expect(transaction_duration.size).to eq(1)
+
+      entry = transaction_duration[0]['transaction_type']
+
+      expect(entry).to eql 'sub_transaction'
+    end
+  end
+
+  context 'when there are real-transactions' do
+    it 'records transaction duration', :delete do
+      observe_real_transaction
+
+      expect(transaction_duration.size).to eq(1)
+
+      entry = transaction_duration[0]['transaction_type']
+
+      expect(entry).to eql 'real_transaction'
+    end
+  end
+
+  private
+
+  def observe
+    transaction_duration_observer.before
+    run_transaction
+    transaction_duration_observer.after
+    transaction_duration_observer.record
+  end
+
+  def observe_sub_transaction
+    transaction_duration_observer.before
+    run_sub_transactions
+    transaction_duration_observer.after
+    transaction_duration_observer.record
+  end
+
+  def observe_real_transaction
+    transaction_duration_observer.before
+    run_real_transactions
+    transaction_duration_observer.after
+    transaction_duration_observer.record
+  end
+
+  def run_real_transactions
+    ActiveRecord::Base.transaction do
+    end
+  end
+
+  def run_sub_transactions
+    ActiveRecord::Base.transaction(requires_new: true) do
+    end
+  end
+
+  def run_transaction
+    ActiveRecord::Base.connection_pool.with_connection do |connection|
+      Gitlab::Database::SharedModel.using_connection(connection) do
+        Gitlab::Database::SharedModel.transaction do
+          Gitlab::Database::SharedModel.transaction(requires_new: true) do
+            Gitlab::Database::SharedModel.transaction do
+              Gitlab::Database::SharedModel.transaction do
+                Gitlab::Database::SharedModel.transaction(requires_new: true) do
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end