Fix the db-specific replica metrics are not stored in logs

lib/gitlab/metrics/subscribers/active_record.rb

@@ -5,6 +5,8 @@ module Gitlab
     module Subscribers
       # Class for tracking the total query duration of a transaction.
       class ActiveRecord < ActiveSupport::Subscriber
+        extend Gitlab::Utils::StrongMemoize
+
         attach_to :active_record
 
         IGNORABLE_SQL = %w{BEGIN COMMIT}.freeze
@@ -107,7 +109,7 @@ module Gitlab
 
           # Per database metrics
           db_config_name = db_config_name(event.payload)
-          duration_key = compose_metric_key(:duration_s, db_role, db_config_name)
+          duration_key = compose_metric_key(:duration_s, nil, db_config_name)
           ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration
         end
 
@@ -144,7 +146,7 @@ module Gitlab
           # when we are also logging the db_role. Otherwise it will be hard to
           # tell if the log key is referring to a db_role OR a db_config_name.
           if db_role.present? && db_config_name.present?
-            log_key = compose_metric_key(counter, db_role, db_config_name)
+            log_key = compose_metric_key(counter, nil, db_config_name)
             Gitlab::SafeRequestStore[log_key] = Gitlab::SafeRequestStore[log_key].to_i + 1
           end
         end
@@ -172,26 +174,34 @@ module Gitlab
         end
 
         def self.load_balancing_metric_counter_keys
-          load_balancing_metric_keys(DB_LOAD_BALANCING_COUNTERS)
+          strong_memoize(:load_balancing_metric_counter_keys) do
+            load_balancing_metric_keys(DB_LOAD_BALANCING_COUNTERS)
+          end
         end
 
         def self.load_balancing_metric_duration_keys
-          load_balancing_metric_keys(DB_LOAD_BALANCING_DURATIONS)
+          strong_memoize(:load_balancing_metric_duration_keys) do
+            load_balancing_metric_keys(DB_LOAD_BALANCING_DURATIONS)
+          end
         end
 
         def self.load_balancing_metric_keys(metrics)
-          [].tap do |counters|
+          counters = []
+
+          metrics.each do |metric|
             DB_LOAD_BALANCING_ROLES.each do |role|
-              metrics.each do |metric|
-                counters << compose_metric_key(metric, role)
-                next unless ENV['GITLAB_MULTIPLE_DATABASE_METRICS']
+              counters << compose_metric_key(metric, role)
+            end
 
-                ::Gitlab::Database.db_config_names.each do |config_name|
-                  counters << compose_metric_key(metric, role, config_name)
-                end
+            if ENV['GITLAB_MULTIPLE_DATABASE_METRICS']
+              ::Gitlab::Database.db_config_names.each do |config_name|
+                counters << compose_metric_key(metric, nil, config_name) # main
+                counters << compose_metric_key(metric, nil, config_name + ::Gitlab::Database::LoadBalancing::LoadBalancer::REPLICA_SUFFIX) # main_replica
               end
             end
           end
+
+          counters
         end
 
         def compose_metric_key(metric, db_role = nil, db_config_name = nil)

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

@@ -198,6 +198,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
         context 'query using a connection to a replica' do
           before do
             allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica)
+            allow(connection).to receive_message_chain(:pool, :db_config, :name).and_return(db_config_name)
           end
 
           it 'queries connection db role' do

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -272,12 +272,12 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         expected_end_payload.merge(
           'db_duration_s' => a_value >= 0.1,
           'db_count' => a_value >= 1,
-          "db_replica_#{db_config_name}_count" => 0,
+          "db_#{db_config_name}_replica_count" => 0,
           'db_replica_duration_s' => a_value >= 0,
           'db_primary_count' => a_value >= 1,
-          "db_primary_#{db_config_name}_count" => a_value >= 1,
+          "db_#{db_config_name}_count" => a_value >= 1,
           'db_primary_duration_s' => a_value > 0,
-          "db_primary_#{db_config_name}_duration_s" => a_value > 0
+          "db_#{db_config_name}_duration_s" => a_value > 0
         )
       end
 

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

 # frozen_string_literal: true
 
 RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
-  let(:db_config_name) { ::Gitlab::Database.db_config_names.first }
+  let(:db_config_name) do
+    db_config_name = ::Gitlab::Database.db_config_names.first
+    db_config_name += "_replica" if db_role == :secondary
+    db_config_name
+  end
 
   let(:expected_payload_defaults) do
     result = {}
@@ -39,15 +43,15 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
                        db_write_count: record_write_query ? 1 : 0,
                        db_cached_count: record_cached_query ? 1 : 0,
                        db_primary_cached_count: record_cached_query ? 1 : 0,
-                       "db_primary_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
+                       "db_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
                        db_primary_count: record_query ? 1 : 0,
-                       "db_primary_#{db_config_name}_count": record_query ? 1 : 0,
+                       "db_#{db_config_name}_count": record_query ? 1 : 0,
                        db_primary_duration_s: record_query ? 0.002 : 0.0,
-                       "db_primary_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
+                       "db_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
                        db_primary_wal_count: record_wal_query ? 1 : 0,
-                       "db_primary_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
+                       "db_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
                        db_primary_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
-                       "db_primary_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
+                       "db_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
                      })
                    elsif db_role == :replica
                      transform_hash(expected_payload_defaults, {
@@ -55,15 +59,15 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
                        db_write_count: record_write_query ? 1 : 0,
                        db_cached_count: record_cached_query ? 1 : 0,
                        db_replica_cached_count: record_cached_query ? 1 : 0,
-                       "db_replica_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
+                       "db_#{db_config_name}_cached_count": record_cached_query ? 1 : 0,
                        db_replica_count: record_query ? 1 : 0,
-                       "db_replica_#{db_config_name}_count": record_query ? 1 : 0,
+                       "db_#{db_config_name}_count": record_query ? 1 : 0,
                        db_replica_duration_s: record_query ? 0.002 : 0.0,
-                       "db_replica_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
+                       "db_#{db_config_name}_duration_s": record_query ? 0.002 : 0.0,
                        db_replica_wal_count: record_wal_query ? 1 : 0,
-                       "db_replica_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
+                       "db_#{db_config_name}_wal_count": record_wal_query ? 1 : 0,
                        db_replica_wal_cached_count: record_wal_query && record_cached_query ? 1 : 0,
-                       "db_replica_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
+                       "db_#{db_config_name}_wal_cached_count": record_wal_query && record_cached_query ? 1 : 0
                      })
                    else
                      transform_hash(expected_payload_defaults, {
@@ -71,15 +75,15 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
                        db_write_count: record_write_query ? 1 : 0,
                        db_cached_count: record_cached_query ? 1 : 0,
                        db_primary_cached_count: 0,
-                       "db_primary_#{db_config_name}_cached_count": 0,
+                       "db_#{db_config_name}_cached_count": 0,
                        db_primary_count: 0,
-                       "db_primary_#{db_config_name}_count": 0,
+                       "db_#{db_config_name}_count": 0,
                        db_primary_duration_s: 0.0,
-                       "db_primary_#{db_config_name}_duration_s": 0.0,
+                       "db_#{db_config_name}_duration_s": 0.0,
                        db_primary_wal_count: 0,
-                       "db_primary_#{db_config_name}_wal_count": 0,
+                       "db_#{db_config_name}_wal_count": 0,
                        db_primary_wal_cached_count: 0,
-                       "db_primary_#{db_config_name}_wal_cached_count": 0
+                       "db_#{db_config_name}_wal_cached_count": 0
                      })
                    end
 
@@ -105,7 +109,11 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
 end
 
 RSpec.shared_examples 'record ActiveRecord metrics in a metrics transaction' do |db_role|
-  let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.retrieve_connection) }
+  let(:db_config_name) do
+    db_config_name = ::Gitlab::Database.db_config_names.first
+    db_config_name += "_replica" if db_role == :secondary
+    db_config_name
+  end
 
   it 'increments only db counters' do
     if record_query