Add instrumentation for recording subtransaction depth

config/feature_flags/ops/active_record_transactions_tracking.yml

+---
+name: active_record_transactions_tracking
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/67918
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/338306
+milestone: '14.2'
+type: ops
+group: group::pipeline execution
+default_enabled: false

config/initializers/active_record_transaction_observer.rb

+# frozen_string_literal: true
+
+return unless Gitlab.com? || Gitlab.dev_or_test_env?
+
+def feature_flags_available?
+  # When the DBMS is not available, an exception (e.g. PG::ConnectionBad) is raised
+  active_db_connection = ActiveRecord::Base.connection.active? rescue false
+
+  active_db_connection && Feature::FlipperFeature.table_exists?
+rescue ActiveRecord::NoDatabaseError
+  false
+end
+
+Gitlab::Application.configure do
+  if feature_flags_available? && ::Feature.enabled?(:active_record_transactions_tracking, type: :ops, default_enabled: :yaml)
+    Gitlab::Database::Transaction::Observer.register!
+  end
+end

config/initializers/active_record_transaction_patches.rb

+# frozen_string_literal: true
+
+if ENV['ACTIVE_RECORD_DISABLE_TRANSACTION_METRICS_PATCHES'].blank?
+  Gitlab::Database.install_transaction_metrics_patches!
+end
+
+return unless Gitlab.com? || Gitlab.dev_or_test_env?
+
+if ENV['ACTIVE_RECORD_DISABLE_TRANSACTION_CONTEXT_PATCHES'].blank?
+  Gitlab::Database.install_transaction_context_patches!
+end

config/initializers/transaction_metrics.rb

-# frozen_string_literal: true
-
-Gitlab::Database.install_monkey_patches

lib/gitlab/database.rb

@@ -177,11 +177,6 @@ module Gitlab
       'unknown'
     end
 
-    # Monkeypatch rails with upgraded database observability
-    def self.install_monkey_patches
-      ActiveRecord::Base.prepend(ActiveRecordBaseTransactionMetrics)
-    end
-
     def self.read_only?
       false
     end
@@ -190,6 +185,18 @@ module Gitlab
       !read_only?
     end
 
+    # Monkeypatch rails with upgraded database observability
+    def self.install_transaction_metrics_patches!
+      ActiveRecord::Base.prepend(ActiveRecordBaseTransactionMetrics)
+    end
+
+    def self.install_transaction_context_patches!
+      ActiveRecord::ConnectionAdapters::TransactionManager
+        .prepend(TransactionManagerContext)
+      ActiveRecord::ConnectionAdapters::RealTransaction
+        .prepend(RealTransactionContext)
+    end
+
     # MonkeyPatch for ActiveRecord::Base for adding observability
     module ActiveRecordBaseTransactionMetrics
       extend ActiveSupport::Concern
@@ -204,6 +211,32 @@ module Gitlab
         end
       end
     end
+
+    # rubocop:disable Gitlab/ModuleWithInstanceVariables
+    module TransactionManagerContext
+      def transaction_context
+        @stack.first.try(:gitlab_transaction_context)
+      end
+    end
+
+    module RealTransactionContext
+      def gitlab_transaction_context
+        @gitlab_transaction_context ||= ::Gitlab::Database::Transaction::Context.new
+      end
+
+      def commit
+        gitlab_transaction_context.commit
+
+        super
+      end
+
+      def rollback
+        gitlab_transaction_context.rollback
+
+        super
+      end
+    end
+    # rubocop:enable Gitlab/ModuleWithInstanceVariables
   end
 end
 

lib/gitlab/database/transaction/context.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Transaction
+      class Context
+        attr_reader :context
+
+        LOG_DEPTH_THRESHOLD = 8
+        LOG_SAVEPOINTS_THRESHOLD = 32
+        LOG_DURATION_S_THRESHOLD = 300
+        LOG_THROTTLE_DURATION = 1
+
+        def initialize
+          @context = {}
+        end
+
+        def set_start_time
+          @context[:start_time] = current_timestamp
+        end
+
+        def increment_savepoints
+          @context[:savepoints] = @context[:savepoints].to_i + 1
+        end
+
+        def increment_rollbacks
+          @context[:rollbacks] = @context[:rollbacks].to_i + 1
+        end
+
+        def increment_releases
+          @context[:releases] = @context[:releases].to_i + 1
+        end
+
+        def set_depth(depth)
+          @context[:depth] = [@context[:depth].to_i, depth].max
+        end
+
+        def track_sql(sql)
+          (@context[:queries] ||= []).push(sql)
+        end
+
+        def duration
+          return unless @context[:start_time].present?
+
+          current_timestamp - @context[:start_time]
+        end
+
+        def depth_threshold_exceeded?
+          @context[:depth].to_i > LOG_DEPTH_THRESHOLD
+        end
+
+        def savepoints_threshold_exceeded?
+          @context[:savepoints].to_i > LOG_SAVEPOINTS_THRESHOLD
+        end
+
+        def duration_threshold_exceeded?
+          duration.to_i > LOG_DURATION_S_THRESHOLD
+        end
+
+        def log_savepoints?
+          depth_threshold_exceeded? || savepoints_threshold_exceeded?
+        end
+
+        def log_duration?
+          duration_threshold_exceeded?
+        end
+
+        def should_log?
+          !logged_already? && (log_savepoints? || log_duration?)
+        end
+
+        def commit
+          log(:commit)
+        end
+
+        def rollback
+          log(:rollback)
+        end
+
+        private
+
+        def queries
+          @context[:queries].to_a.join("\n")
+        end
+
+        def current_timestamp
+          ::Gitlab::Metrics::System.monotonic_time
+        end
+
+        def logged_already?
+          return false if @context[:last_log_timestamp].nil?
+
+          (current_timestamp - @context[:last_log_timestamp].to_i) < LOG_THROTTLE_DURATION
+        end
+
+        def set_last_log_timestamp
+          @context[:last_log_timestamp] = current_timestamp
+        end
+
+        def log(operation)
+          return unless should_log?
+
+          set_last_log_timestamp
+
+          attributes = {
+            class: self.class.name,
+            result: operation,
+            duration_s: duration,
+            depth: @context[:depth].to_i,
+            savepoints_count: @context[:savepoints].to_i,
+            rollbacks_count: @context[:rollbacks].to_i,
+            releases_count: @context[:releases].to_i,
+            sql: queries
+          }
+
+          application_info(attributes)
+        end
+
+        def application_info(attributes)
+          Gitlab::AppJsonLogger.info(attributes)
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/transaction/observer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Transaction
+      class Observer
+        INSTRUMENTED_STATEMENTS = %w[BEGIN SAVEPOINT ROLLBACK RELEASE].freeze
+        LONGEST_COMMAND_LENGTH = 'ROLLBACK TO SAVEPOINT'.length
+        START_COMMENT = '/*'
+        END_COMMENT = '*/'
+
+        def self.instrument_transactions(cmd, event)
+          connection = event.payload[:connection]
+          manager = connection&.transaction_manager
+          return unless manager.respond_to?(:transaction_context)
+
+          context = manager.transaction_context
+          return if context.nil?
+
+          if cmd.start_with?('BEGIN')
+            context.set_start_time
+            context.set_depth(0)
+            context.track_sql(event.payload[:sql])
+          elsif cmd.start_with?('SAVEPOINT ')
+            context.set_depth(manager.open_transactions)
+            context.increment_savepoints
+          elsif cmd.start_with?('ROLLBACK TO SAVEPOINT')
+            context.increment_rollbacks
+          elsif cmd.start_with?('RELEASE SAVEPOINT ')
+            context.increment_releases
+          end
+        end
+
+        def self.register!
+          ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+            sql = event.payload.dig(:sql).to_s
+            cmd = extract_sql_command(sql)
+
+            if cmd.start_with?(*INSTRUMENTED_STATEMENTS)
+              self.instrument_transactions(cmd, event)
+            end
+          end
+        end
+
+        def self.extract_sql_command(sql)
+          return sql unless sql.start_with?(START_COMMENT)
+
+          index = sql.index(END_COMMENT)
+
+          return sql unless index
+
+          # /* comment */ SELECT
+          #
+          # We offset using a position of the end comment + 1 character to
+          # accomodate a space between Marginalia comment and a SQL statement.
+          offset = index + END_COMMENT.length + 1
+
+          # Avoid duplicating the entire string. This isn't optimized to
+          # strip extra spaces, but we assume that this doesn't happen
+          # for performance reasons.
+          sql[offset..offset + LONGEST_COMMAND_LENGTH]
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/transaction/context_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Transaction::Context do
+  subject { described_class.new }
+
+  let(:data) { subject.context }
+
+  before do
+    stub_const("#{described_class}::LOG_THROTTLE", 100)
+  end
+
+  describe '#set_start_time' do
+    before do
+      subject.set_start_time
+    end
+
+    it 'sets start_time' do
+      expect(data).to have_key(:start_time)
+    end
+  end
+
+  describe '#increment_savepoints' do
+    before do
+      2.times { subject.increment_savepoints }
+    end
+
+    it { expect(data[:savepoints]).to eq(2) }
+  end
+
+  describe '#increment_rollbacks' do
+    before do
+      3.times { subject.increment_rollbacks }
+    end
+
+    it { expect(data[:rollbacks]).to eq(3) }
+  end
+
+  describe '#increment_releases' do
+    before do
+      4.times { subject.increment_releases }
+    end
+
+    it { expect(data[:releases]).to eq(4) }
+  end
+
+  describe '#set_depth' do
+    before do
+      subject.set_depth(2)
+    end
+
+    it { expect(data[:depth]).to eq(2) }
+  end
+
+  describe '#track_sql' do
+    before do
+      subject.track_sql('SELECT 1')
+      subject.track_sql('SELECT * FROM users')
+    end
+
+    it { expect(data[:queries]).to eq(['SELECT 1', 'SELECT * FROM users']) }
+  end
+
+  describe '#duration' do
+    before do
+      subject.set_start_time
+    end
+
+    it { expect(subject.duration).to be >= 0 }
+  end
+
+  context 'when depth is low' do
+    it 'does not log data upon COMMIT' do
+      expect(subject).not_to receive(:application_info)
+
+      subject.commit
+    end
+
+    it 'does not log data upon ROLLBACK' do
+      expect(subject).not_to receive(:application_info)
+
+      subject.rollback
+    end
+
+    it '#should_log? returns false' do
+      expect(subject.should_log?).to be false
+    end
+  end
+
+  shared_examples 'logs transaction data' do
+    it 'logs once upon COMMIT' do
+      expect(subject).to receive(:application_info).and_call_original
+
+      2.times { subject.commit }
+    end
+
+    it 'logs once upon ROLLBACK' do
+      expect(subject).to receive(:application_info).once
+
+      2.times { subject.rollback }
+    end
+
+    it 'logs again when log throttle duration passes' do
+      expect(subject).to receive(:application_info).twice.and_call_original
+
+      2.times { subject.commit }
+
+      data[:last_log_timestamp] -= (described_class::LOG_THROTTLE_DURATION + 1)
+
+      subject.commit
+    end
+
+    it '#should_log? returns true' do
+      expect(subject.should_log?).to be true
+    end
+  end
+
+  context 'when depth exceeds threshold' do
+    before do
+      subject.set_depth(described_class::LOG_DEPTH_THRESHOLD + 1)
+    end
+
+    it_behaves_like 'logs transaction data'
+  end
+
+  context 'when savepoints count exceeds threshold' do
+    before do
+      data[:savepoints] = described_class::LOG_SAVEPOINTS_THRESHOLD + 1
+    end
+
+    it_behaves_like 'logs transaction data'
+  end
+
+  context 'when duration exceeds threshold' do
+    before do
+      subject.set_start_time
+
+      data[:start_time] -= (described_class::LOG_DURATION_S_THRESHOLD + 1)
+    end
+
+    it_behaves_like 'logs transaction data'
+  end
+end

spec/lib/gitlab/database/transaction/observer_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Transaction::Observer do
+  # Use the delete DB strategy so that the test won't be wrapped in a transaction
+  describe '.instrument_transactions', :delete do
+    let(:transaction_context) { ActiveRecord::Base.connection.transaction_manager.transaction_context }
+    let(:context) { transaction_context.context }
+
+    around do |example|
+      # Emulate production environment when SQL comments come first to avoid truncation
+      Marginalia::Comment.prepend_comment = true
+      subscriber = described_class.register!
+
+      example.run
+
+      ActiveSupport::Notifications.unsubscribe(subscriber)
+      Marginalia::Comment.prepend_comment = false
+    end
+
+    it 'tracks transaction data', :aggregate_failures do
+      ActiveRecord::Base.transaction do
+        ActiveRecord::Base.transaction(requires_new: true) do
+          User.first
+
+          expect(transaction_context).to be_a(::Gitlab::Database::Transaction::Context)
+          expect(context.keys).to match_array(%i(start_time depth savepoints queries))
+          expect(context[:depth]).to eq(2)
+          expect(context[:savepoints]).to eq(1)
+          expect(context[:queries].length).to eq(1)
+        end
+      end
+
+      expect(context[:depth]).to eq(2)
+      expect(context[:savepoints]).to eq(1)
+      expect(context[:releases]).to eq(1)
+    end
+
+    describe '.extract_sql_command' do
+      using RSpec::Parameterized::TableSyntax
+
+      where(:sql, :expected) do
+        'SELECT 1' | 'SELECT 1'
+        '/* test comment */ SELECT 1' | 'SELECT 1'
+        '/* test comment */ ROLLBACK TO SAVEPOINT point1' | 'ROLLBACK TO SAVEPOINT '
+        'SELECT 1 /* trailing comment */' | 'SELECT 1 /* trailing comment */'
+      end
+
+      with_them do
+        it do
+          expect(described_class.extract_sql_command(sql)).to eq(expected)
+        end
+      end
+    end
+  end
+end