Expose database metrics to sidekiq

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/323163

ee/lib/ee/gitlab/metrics/subscribers/active_record.rb

@@ -43,7 +43,7 @@ module EE
 
             db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
             increment_db_role_counters(db_role, payload)
-            observe_db_role_duration(db_role, event.duration)
+            observe_db_role_duration(db_role, event)
           end
 
           private
@@ -53,13 +53,10 @@ module EE
             increment("db_#{db_role}_cached_count".to_sym) if cached_query?(payload)
           end
 
-          def observe_db_role_duration(db_role, duration)
-            duration /= 1000.0
-
-            current_transaction&.observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, duration) do
-              buckets [0.05, 0.1, 0.25]
-            end
+          def observe_db_role_duration(db_role, event)
+            observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, event)
 
+            duration = event.duration / 1000.0
             duration_key = "db_#{db_role}_duration_s".to_sym
             ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration
           end

ee/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -6,7 +6,6 @@ RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do
   using RSpec::Parameterized::TableSyntax
 
   let(:env) { {} }
-  let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
   let(:subscriber) { described_class.new }
   let(:connection) { double(:connection) }
   let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } }
@@ -49,8 +48,6 @@ RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do
 
       before do
         allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(subscriber).to receive(:current_transaction)
-          .and_return(transaction)
       end
 
       context 'query using a connection to a replica' do

lib/gitlab/metrics/subscribers/active_record.rb

@@ -11,6 +11,14 @@ module Gitlab
         DB_COUNTERS = %i{db_count db_write_count db_cached_count}.freeze
         SQL_COMMANDS_WITH_COMMENTS_REGEX = /\A(\/\*.*\*\/\s)?((?!(.*[^\w'"](DELETE|UPDATE|INSERT INTO)[^\w'"])))(WITH.*)?(SELECT)((?!(FOR UPDATE|FOR SHARE)).)*$/i.freeze
 
+        DURATION_BUCKET = [0.05, 0.1, 0.25].freeze
+
+        # observe_transaction_duration is called from ActiveRecordBaseTransactionMetrics.transaction and used to
+        # record transaction durations.
+        def transaction(event)
+          observe(:gitlab_database_transaction_seconds, event)
+        end
+
         def sql(event)
           # Mark this thread as requiring a database connection. This is used
           # by the Gitlab::Metrics::Samplers::ThreadsSampler to count threads
@@ -20,10 +28,11 @@ module Gitlab
           payload = event.payload
           return if ignored_query?(payload)
 
-          increment_db_counters(payload)
-          current_transaction&.observe(:gitlab_sql_duration_seconds, event.duration / 1000.0) do
-            buckets [0.05, 0.1, 0.25]
-          end
+          increment(:db_count)
+          increment(:db_cached_count) if cached_query?(payload)
+          increment(:db_write_count) unless select_sql_command?(payload)
+
+          observe(:gitlab_sql_duration_seconds, event)
         end
 
         def self.db_counter_payload
@@ -50,20 +59,30 @@ module Gitlab
           payload[:sql].match(SQL_COMMANDS_WITH_COMMENTS_REGEX)
         end
 
-        def increment_db_counters(payload)
-          increment(:db_count)
-          increment(:db_cached_count) if cached_query?(payload)
-          increment(:db_write_count) unless select_sql_command?(payload)
-        end
-
         def increment(counter)
-          current_transaction&.increment("gitlab_transaction_#{counter}_total".to_sym, 1)
+          web_transaction&.increment("gitlab_transaction_#{counter}_total".to_sym, 1)
+          background_transaction&.increment("gitlab_transaction_#{counter}_total".to_sym, 1)
 
           Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
         end
 
-        def current_transaction
-          ::Gitlab::Metrics::Transaction.current
+        def observe(histogram, event)
+          # puts "web_transaction = #{web_transaction}"
+          web_transaction&.observe(histogram, event.duration / 1000.0) do
+            buckets DURATION_BUCKET
+          end
+          # puts "background_transaction = #{background_transaction}"
+          background_transaction&.observe(histogram, event.duration / 1000.0) do
+            buckets DURATION_BUCKET
+          end
+        end
+
+        def web_transaction
+          ::Gitlab::Metrics::WebTransaction.current
+        end
+
+        def background_transaction
+          ::Gitlab::Metrics::BackgroundTransaction.current
         end
       end
     end

spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

@@ -6,7 +6,6 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
   using RSpec::Parameterized::TableSyntax
 
   let(:env) { {} }
-  let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
   let(:subscriber)  { described_class.new }
   let(:connection) { double(:connection) }
   let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } }
@@ -47,33 +46,15 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
     with_them do
       let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
 
-      describe 'with a current transaction' do
-        before do
-          allow(subscriber).to receive(:current_transaction)
-            .at_least(:once)
-            .and_return(transaction)
-        end
+      it 'marks the current thread as using the database' do
+        # since it would already have been toggled by other specs
+        Thread.current[:uses_db_connection] = nil
 
-        it 'marks the current thread as using the database' do
-          # since it would already have been toggled by other specs
-          Thread.current[:uses_db_connection] = nil
-
-          expect { subscriber.sql(event) }.to change { Thread.current[:uses_db_connection] }.from(nil).to(true)
-        end
-
-        it_behaves_like 'record ActiveRecord metrics'
-        it_behaves_like 'store ActiveRecord info in RequestStore'
+        expect { subscriber.sql(event) }.to change { Thread.current[:uses_db_connection] }.from(nil).to(true)
       end
 
-      describe 'without a current transaction' do
-        it 'does not track any metrics' do
-          expect_any_instance_of(Gitlab::Metrics::Transaction)
-            .not_to receive(:increment)
-          subscriber.sql(event)
-        end
-
-        it_behaves_like 'store ActiveRecord info in RequestStore'
-      end
+      it_behaves_like 'record ActiveRecord metrics'
+      it_behaves_like 'store ActiveRecord info in RequestStore'
     end
   end
 

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

@@ -42,7 +42,7 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
   end
 end
 
-RSpec.shared_examples 'record ActiveRecord metrics' do |db_role|
+RSpec.shared_examples 'record ActiveRecord metrics in a metrics transaction' do |db_role|
   it 'increments only db counters' do
     if record_query
       expect(transaction).to receive(:increment).with(:gitlab_transaction_db_count_total, 1)
@@ -80,3 +80,47 @@ RSpec.shared_examples 'record ActiveRecord metrics' do |db_role|
     subscriber.sql(event)
   end
 end
+
+RSpec.shared_examples 'record ActiveRecord metrics' do |db_role|
+  context 'when both web and background transaction are available' do
+    before do
+      allow(::Gitlab::Metrics::WebTransaction).to receive(:current)
+        .and_return(nil)
+      allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current)
+        .and_return(nil)
+    end
+
+    it 'does not capture the metrics' do
+      expect_any_instance_of(Gitlab::Metrics::Transaction).not_to receive(:increment) # rubocop:disable RSpec/AnyInstanceOf
+      expect_any_instance_of(Gitlab::Metrics::Transaction).not_to receive(:observe) # rubocop:disable RSpec/AnyInstanceOf
+
+      subscriber.sql(event)
+    end
+  end
+
+  context 'when web transaction is available' do
+    let(:transaction) { double('Gitlab::Metrics::WebTransaction', increment: nil, observe: nil) }
+
+    before do
+      allow(::Gitlab::Metrics::WebTransaction).to receive(:current)
+        .and_return(transaction)
+      allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current)
+        .and_return(nil)
+    end
+
+    it_behaves_like 'record ActiveRecord metrics in a metrics transaction', db_role
+  end
+
+  context 'when background transaction is available' do
+    let(:transaction) { double('Gitlab::Metrics::BackgroundTransaction', increment: nil, observe: nil) }
+
+    before do
+      allow(::Gitlab::Metrics::WebTransaction).to receive(:current)
+        .and_return(nil)
+      allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current)
+        .and_return(transaction)
+    end
+
+    it_behaves_like 'record ActiveRecord metrics in a metrics transaction', db_role
+  end
+end