Fix db call logging numbers

- Use RequestStore for caching
- Fix Specs
- Add missing specs

lib/gitlab/lograge/custom_options.rb

@@ -20,7 +20,7 @@ module Gitlab
           username: event.payload[:username],
           ua: event.payload[:ua]
         }
-        add_db_counters!(payload)
+        payload.merge!(::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_payload)
 
         payload.merge!(event.payload[:metadata]) if event.payload[:metadata]
 
@@ -46,16 +46,6 @@ module Gitlab
 
         payload
       end
-
-      def self.add_db_counters!(payload)
-        current_transaction = Gitlab::Metrics::Transaction.current
-        if current_transaction
-          payload[:db_count] = current_transaction.get(:db_count, :counter).to_i
-          payload[:db_write_count] = current_transaction.get(:db_write_count, :counter).to_i
-          payload[:db_cached_count] = current_transaction.get(:db_cached_count, :counter).to_i
-        end
-      end
-      private_class_method :add_db_counters!
     end
   end
 end

lib/gitlab/metrics/sidekiq_middleware.rb

@@ -26,9 +26,7 @@ module Gitlab
       private
 
       def add_info_to_payload(payload, trans)
-        payload[:db_count] = trans.get(:db_count, :counter).to_i
-        payload[:db_write_count] = trans.get(:db_write_count, :counter).to_i
-        payload[:db_cached_count] = trans.get(:db_cached_count, :counter).to_i
+        payload.merge!(::Gitlab::Metrics::Subscribers::ActiveRecord.db_counter_payload)
       end
     end
   end

lib/gitlab/metrics/subscribers/active_record.rb

@@ -23,6 +23,18 @@ module Gitlab
           increment_db_counters(payload)
         end
 
+        def self.db_counter_payload
+          payload = {}
+
+          if Gitlab::SafeRequestStore.active?
+            DB_COUNTERS.each do |counter|
+              payload[counter] = Gitlab::SafeRequestStore[counter] if Gitlab::SafeRequestStore.exist?(counter)
+            end
+          end
+
+          payload
+        end
+
         private
 
         define_histogram :gitlab_sql_duration_seconds do
@@ -36,13 +48,30 @@ module Gitlab
         end
 
         def increment_db_counters(payload)
-          current_transaction.increment(:db_count, 1)
+          initialize_cache_keys
+          increment(:db_count)
 
           if payload.fetch(:cached, payload[:name] == 'CACHE')
-            current_transaction.increment(:db_cached_count, 1)
+            increment(:db_cached_count)
+          end
+
+          increment(:db_write_count) unless select_sql_command?(payload)
         end
 
-          current_transaction.increment(:db_write_count, 1) unless select_sql_command?(payload)
+        def initialize_cache_keys
+          return unless Gitlab::SafeRequestStore.active?
+
+          DB_COUNTERS.each do |counter|
+            Gitlab::SafeRequestStore[counter] = 0 unless Gitlab::SafeRequestStore.exist?(counter)
+          end
+        end
+
+        def increment(counter)
+          current_transaction.increment(counter, 1)
+
+          if Gitlab::SafeRequestStore.active?
+            Gitlab::SafeRequestStore[counter] += 1
+          end
         end
 
         def current_transaction

lib/gitlab/metrics/transaction.rb

@@ -92,12 +92,6 @@ module Gitlab
         self.class.transaction_metric(name, :gauge).set(labels, value) if use_prometheus
       end
 
-      def get(name, type, tags = {})
-        metric = self.class.transaction_metric(name, type)
-
-        metric.get(filter_tags(tags).merge(labels))
-      end
-
       def labels
         BASE_LABELS
       end

spec/lib/gitlab/lograge/custom_options_spec.rb

@@ -51,8 +51,10 @@ RSpec.describe Gitlab::Lograge::CustomOptions do
         allow(Gitlab::Metrics::Transaction).to receive(:current).and_return(transaction)
       end
 
-      it 'adds db counters' do
-        expect(subject).to include(:db_count, :db_write_count, :db_cached_count)
+      it 'adds db counters', :request_store do
+        ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+
+        expect(subject).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
       end
     end
 

spec/lib/gitlab/metrics/sidekiq_middleware_spec.rb

@@ -18,8 +18,20 @@ RSpec.describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, message, :test) do
         ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
       end
+    end
+
+    it 'prevents database counters from leaking to the next transaction' do
+      worker = double(:worker, class: double(:class, name: 'TestWorker'))
+
+      2.times do
+        Gitlab::WithRequestStore.with_request_store do
+          middleware.call(worker, message, :test) do
+            ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+          end
+        end
+      end
 
-      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
+      expect(message).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
     end
 
     it 'tracks the transaction (for messages without `enqueued_at`)', :aggregate_failures do
@@ -35,19 +47,20 @@ RSpec.describe Gitlab::Metrics::SidekiqMiddleware do
       middleware.call(worker, {}, :test) { nil }
     end
 
-    it 'tracks any raised exceptions', :aggregate_failures do
+    it 'tracks any raised exceptions', :aggregate_failures, :request_store do
       worker = double(:worker, class: double(:class, name: 'TestWorker'))
 
-      expect_any_instance_of(Gitlab::Metrics::Transaction)
-        .to receive(:run).and_raise(RuntimeError)
-
       expect_any_instance_of(Gitlab::Metrics::Transaction)
         .to receive(:add_event).with(:sidekiq_exception)
 
-      expect { middleware.call(worker, message, :test) }
-        .to raise_error(RuntimeError)
+      expect do
+        middleware.call(worker, message, :test) do
+          ActiveRecord::Base.connection.execute('SELECT pg_sleep(0.1);')
+          raise RuntimeError
+        end
+      end.to raise_error(RuntimeError)
 
-      expect(message).to include(:db_count, :db_write_count, :db_cached_count)
+      expect(message).to include(db_count: 1, db_write_count: 0, db_cached_count: 0)
     end
   end
 end

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

@@ -28,109 +28,123 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
     end
 
     describe 'with a current transaction' do
-      shared_examples 'read only query' do
-        it 'increments only db count value' do
+      shared_examples 'track executed query' do
+        before do
           allow(subscriber).to receive(:current_transaction)
                                  .at_least(:once)
                                  .and_return(transaction)
+        end
 
-          expect(transaction).to receive(:increment)
-                                   .with(:db_count, 1)
-
-          expect(transaction).not_to receive(:increment)
-                                       .with(:db_cached_count, 1)
-
-          expect(transaction).not_to receive(:increment)
-                                       .with(:db_write_count, 1)
-
-          subscriber.sql(event)
+        it 'increments only db count value' do
+          described_class::DB_COUNTERS.each do |counter|
+            if expected_counters[counter] > 0
+              expect(transaction).to receive(:increment).with(counter, 1)
+            else
+              expect(transaction).not_to receive(:increment).with(counter, 1)
             end
           end
 
-      shared_examples 'write query' do
-        it 'increments db_write_count and db_count value' do
-          expect(subscriber).to receive(:current_transaction)
-                                  .at_least(:once)
-                                  .and_return(transaction)
-
-          expect(transaction).to receive(:increment)
-                                   .with(:db_count, 1)
+          subscriber.sql(event)
+        end
 
-          expect(transaction).not_to receive(:increment)
-                                       .with(:db_cached_count, 1)
+        context 'when RequestStore is enabled' do
+          it 'caches db count value', :request_store, :aggregate_failures do
+            subscriber.sql(event)
 
-          expect(transaction).to receive(:increment)
-                                       .with(:db_write_count, 1)
+            described_class::DB_COUNTERS.each do |counter|
+              expect(Gitlab::SafeRequestStore[counter]).to eq expected_counters[counter]
+            end
+          end
 
+          it 'prevents db counters from leaking to the next transaction' do
+            2.times do
+              Gitlab::WithRequestStore.with_request_store do
                 subscriber.sql(event)
+
+                described_class::DB_COUNTERS.each do |counter|
+                  expect(Gitlab::SafeRequestStore[counter]).to eq expected_counters[counter]
+                end
+              end
+            end
+          end
         end
       end
 
-      shared_examples 'cached query' do
-        it 'increments db_cached_count and db_count value' do
+      it 'observes sql_duration metric' do
         expect(subscriber).to receive(:current_transaction)
                                 .at_least(:once)
                                 .and_return(transaction)
+        expect(described_class.send(:gitlab_sql_duration_seconds)).to receive(:observe).with({}, 0.002)
+        subscriber.sql(event)
+      end
 
-          expect(transaction).to receive(:increment)
-                                   .with(:db_count, 1)
+      context 'with read query' do
+        let(:expected_counters) do
+          {
+            db_count: 1,
+            db_write_count: 0,
+            db_cached_count: 0
+          }
+        end
 
-          expect(transaction).to receive(:increment)
-                                       .with(:db_cached_count, 1)
+        it_behaves_like 'track executed query'
 
-          expect(transaction).not_to receive(:increment)
-                                   .with(:db_write_count, 1)
+        context 'with only select' do
+          let(:payload) { { sql: 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' } }
 
-          subscriber.sql(event)
+          it_behaves_like 'track executed query'
         end
       end
 
-      it 'observes sql_duration metric' do
-        expect(subscriber).to receive(:current_transaction)
-                                .at_least(:once)
-                                .and_return(transaction)
-        expect(described_class.send(:gitlab_sql_duration_seconds)).to receive(:observe).with({}, 0.002)
-        subscriber.sql(event)
+      context 'write query' do
+        let(:expected_counters) do
+          {
+            db_count: 1,
+            db_write_count: 1,
+            db_cached_count: 0
+          }
         end
 
-      it_behaves_like 'read only query'
-
         context 'with select for update sql event' do
           let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10 FOR UPDATE' } }
 
-        it_behaves_like 'write query'
+          it_behaves_like 'track executed query'
         end
 
         context 'with common table expression' do
           context 'with insert' do
             let(:payload) { { sql: 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' } }
 
-          it_behaves_like 'write query'
-        end
-
-        context 'with only select' do
-          let(:payload) { { sql: 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' } }
-
-          it_behaves_like 'read only query'
+            it_behaves_like 'track executed query'
           end
         end
 
         context 'with delete sql event' do
           let(:payload) { { sql: 'DELETE FROM users where id = 10' } }
 
-        it_behaves_like 'write query'
+          it_behaves_like 'track executed query'
         end
 
         context 'with insert sql event' do
           let(:payload) { { sql: 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' } }
 
-        it_behaves_like 'write query'
+          it_behaves_like 'track executed query'
         end
 
         context 'with update sql event' do
           let(:payload) { { sql: 'UPDATE users SET admin = true WHERE id = 10' } }
 
-        it_behaves_like 'write query'
+          it_behaves_like 'track executed query'
+        end
+      end
+
+      context 'with cached query' do
+        let(:expected_counters) do
+          {
+            db_count: 1,
+            db_write_count: 0,
+            db_cached_count: 1
+          }
         end
 
         context 'with cached payload ' do
@@ -141,7 +155,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
             }
           end
 
-        it_behaves_like 'cached query'
+          it_behaves_like 'track executed query'
         end
 
         context 'with cached payload name' do
@@ -152,7 +166,8 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
             }
           end
 
-        it_behaves_like 'cached query'
+          it_behaves_like 'track executed query'
+        end
       end
 
       context 'events are internal to Rails or irrelevant' do
@@ -215,4 +230,48 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
       end
     end
   end
+
+  describe 'self.db_counter_payload' do
+    before do
+      allow(subscriber).to receive(:current_transaction)
+                             .at_least(:once)
+                             .and_return(transaction)
+    end
+
+    context 'when RequestStore is enabled', :request_store do
+      context 'when query is executed' do
+        let(:expected_payload) do
+          {
+            db_count: 1,
+            db_cached_count: 0,
+            db_write_count: 0
+          }
+        end
+
+        it 'returns correct payload' do
+          subscriber.sql(event)
+
+          expect(described_class.db_counter_payload).to eq(expected_payload)
+        end
+      end
+
+      context 'when query is not executed' do
+        let(:expected_payload) { {} }
+
+        it 'returns empty payload' do
+          expect(described_class.db_counter_payload).to eq(expected_payload)
+        end
+      end
+    end
+
+    context 'when RequestStore is disabled' do
+      let(:expected_payload) { {} }
+
+      it 'returns correct payload' do
+        subscriber.sql(event)
+
+        expect(described_class.db_counter_payload).to eq(expected_payload)
+      end
+    end
+  end
 end

spec/lib/gitlab/metrics/transaction_spec.rb

@@ -114,15 +114,4 @@ RSpec.describe Gitlab::Metrics::Transaction do
       transaction.set(:meow, 1)
     end
   end
-
-  describe '#get' do
-    let(:prometheus_metric) { instance_double(Prometheus::Client::Counter, get: nil) }
-
-    it 'gets a metric' do
-      expect(described_class).to receive(:fetch_metric).with(:counter, :gitlab_transaction_meow_total).and_return(prometheus_metric)
-      expect(prometheus_metric).to receive(:get)
-
-      transaction.get(:meow, :counter)
-    end
-  end
 end