Merge branch 'qmnguyen0711/inaccurate-database-role-classification' into 'master'

Fix inaccurate role classification in DB instrumentation

See merge request gitlab-org/gitlab!57998

ee/changelogs/unreleased/qmnguyen0711-inaccurate-database-role-classification.yml

+---
+title: Fix inaccurate role classification in DB instrumentation
+merge_request: 57998
+author:
+type: fixed

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

@@ -47,6 +47,8 @@ module EE
             return if ignored_query?(payload)
 
             db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
+            return if db_role.blank?
+
             increment_db_role_counters(db_role, payload)
             observe_db_role_duration(db_role, event)
           end

ee/lib/ee/peek/views/active_record.rb

@@ -11,7 +11,9 @@ module EE
           detail = super
 
           if ::Gitlab::Database::LoadBalancing.enable?
-            detail[:db_role] = ::Gitlab::Database::LoadBalancing.db_role_for_connection(data[:connection]).to_s.capitalize
+            role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(data[:connection]) ||
+              ::Gitlab::Database::LoadBalancing::ROLE_UNKNOWN
+            detail[:db_role] = role.to_s.capitalize
           end
 
           detail

ee/lib/gitlab/database/load_balancing.rb

@@ -143,7 +143,8 @@ module Gitlab
 
       DB_ROLES = [
         ROLE_PRIMARY = :primary,
-        ROLE_REPLICA = :replica
+        ROLE_REPLICA = :replica,
+        ROLE_UNKNOWN = :unknown
       ].freeze
 
       # Returns the role (primary/replica) of the database the connection is
@@ -155,7 +156,7 @@ module Gitlab
       def self.db_role_for_connection(connection)
         return ROLE_PRIMARY if !enable? || @proxy.blank?
 
-        proxy.load_balancer.db_role_for_connection(connection) || ROLE_PRIMARY
+        proxy.load_balancer.db_role_for_connection(connection)
       end
     end
   end

ee/lib/gitlab/database/load_balancing/host_list.rb

@@ -8,11 +8,13 @@ module Gitlab
         # hosts - The list of secondary hosts to add.
         def initialize(hosts = [])
           @hosts = hosts.shuffle
+          @pools = Set.new
           @index = 0
           @mutex = Mutex.new
           @hosts_gauge = Gitlab::Metrics.gauge(:db_load_balancing_hosts, 'Current number of load balancing hosts')
 
           set_metrics!
+          update_pools
         end
 
         def hosts
@@ -27,9 +29,14 @@ module Gitlab
           @mutex.synchronize { @hosts.map { |host| [host.host, host.port] } }
         end
 
+        def manage_pool?(pool)
+          @pools.include?(pool)
+        end
+
         def hosts=(hosts)
           @mutex.synchronize do
             @hosts = hosts.shuffle
+            update_pools
             @index = 0
           end
 
@@ -71,6 +78,10 @@ module Gitlab
         def set_metrics!
           @hosts_gauge.set({}, @hosts.length)
         end
+
+        def update_pools
+          @pools = Set.new(@hosts.map(&:pool))
+        end
       end
     end
   end

ee/lib/gitlab/database/load_balancing/load_balancer.rb

@@ -107,7 +107,9 @@ module Gitlab
         # is connecting to. If the connection is not issued by this load
         # balancer, return nil
         def db_role_for_connection(connection)
-          @connection_db_roles[connection]
+          return @connection_db_roles[connection] if @connection_db_roles[connection]
+          return ROLE_REPLICA if @host_list.manage_pool?(connection.pool)
+          return ROLE_PRIMARY if connection.pool == ActiveRecord::Base.connection_pool
         end
 
         # Returns a host to use for queries.

ee/spec/lib/gitlab/database/load_balancing/host_list_spec.rb

@@ -69,6 +69,48 @@ RSpec.describe Gitlab::Database::LoadBalancing::HostList do
     end
   end
 
+  describe '#manage_pool?' do
+    before do
+      allow(Gitlab::Database).to receive(:create_connection_pool) { double(:connection) }
+    end
+
+    context 'when the testing pool belongs to one host of the host list' do
+      it 'returns true' do
+        pool = host_list.hosts.first.pool
+
+        expect(host_list.manage_pool?(pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool belongs to a former host of the host list' do
+      it 'returns false' do
+        pool = host_list.hosts.first.pool
+        host_list.hosts = [
+          Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        ]
+
+        expect(host_list.manage_pool?(pool)).to be(false)
+      end
+    end
+
+    context 'when the testing pool belongs to a new host of the host list' do
+      it 'returns true' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+        host_list.hosts = [host]
+
+        expect(host_list.manage_pool?(host.pool)).to be(true)
+      end
+    end
+
+    context 'when the testing pool does not have any relation with the host list' do
+      it 'returns false' do
+        host = Gitlab::Database::LoadBalancing::Host.new('foo', load_balancer)
+
+        expect(host_list.manage_pool?(host.pool)).to be(false)
+      end
+    end
+  end
+
   describe '#hosts=' do
     it 'updates the list of hosts to use' do
       host_list.hosts = [

ee/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb

@@ -226,9 +226,29 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
       end
     end
 
-    context 'when the connection does not come from the load balancer' do
+    context 'when the connection comes from a pool managed by the host list' do
+      it 'returns :replica' do
+        connection = double(:connection)
+        allow(connection).to receive(:pool).and_return(lb.host_list.hosts.first.pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(:replica)
+      end
+    end
+
+    context 'when the connection comes from the primary pool' do
+      it 'returns :primary' do
+        connection = double(:connection)
+        allow(connection).to receive(:pool).and_return(ActiveRecord::Base.connection_pool)
+
+        expect(lb.db_role_for_connection(connection)).to be(:primary)
+      end
+    end
+
+    context 'when the connection does not come from any known pool' do
       it 'returns nil' do
         connection = double(:connection)
+        pool = double(:connection_pool)
+        allow(connection).to receive(:pool).and_return(pool)
 
         expect(lb.db_role_for_connection(connection)).to be(nil)
       end

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -399,10 +399,10 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       context 'when the load balancer returns nil' do
-        it 'returns :primary' do
+        it 'returns nil' do
           allow(load_balancer).to receive(:db_role_for_connection).and_return(nil)
 
-          expect(described_class.db_role_for_connection(connection)).to be(:primary)
+          expect(described_class.db_role_for_connection(connection)).to be(nil)
 
           expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
         end
@@ -702,6 +702,62 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
     end
 
+    context 'custom connection handling' do
+      where(:queries, :expected_role) do
+        [
+          # Reload cache. The schema loading queries should be handled by
+          # primary.
+          [
+            -> {
+              model.connection.clear_cache!
+              model.connection.schema_cache.add('users')
+              model.connection.pool.release_connection
+            },
+            :primary
+          ],
+
+          # Call model's connection method
+          [
+            -> {
+              connection = model.connection
+              connection.select_one('SELECT 1')
+              connection.pool.release_connection
+            },
+            :replica
+          ],
+
+          # Retrieve connection via #retrieve_connection
+          [
+            -> {
+              connection = model.retrieve_connection
+              connection.select_one('SELECT 1')
+              connection.pool.release_connection
+            },
+            :primary
+          ]
+        ]
+      end
+
+      with_them do
+        include_context 'LoadBalancing setup'
+
+        it 'redirects queries to the right roles' do
+          roles = []
+
+          subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |event|
+            role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(event.payload[:connection])
+            roles << role if role.present?
+          end
+
+          self.instance_exec(&queries)
+
+          expect(roles).to all(eql(expected_role))
+        ensure
+          ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+        end
+      end
+    end
+
     context 'a write inside a transaction inside use_replica_if_possible block' do
       include_context 'LoadBalancing setup'
 

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

@@ -83,6 +83,48 @@ RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do
         it_behaves_like 'record ActiveRecord metrics', :primary
         it_behaves_like 'store ActiveRecord info in RequestStore', :primary
       end
+
+      context 'query using a connection to an unknown source' do
+        let(:transaction) { double('Gitlab::Metrics::WebTransaction') }
+
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(nil)
+
+          allow(::Gitlab::Metrics::WebTransaction).to receive(:current).and_return(transaction)
+          allow(::Gitlab::Metrics::BackgroundTransaction).to receive(:current).and_return(nil)
+
+          allow(transaction).to receive(:increment)
+          allow(transaction).to receive(:observe)
+        end
+
+        it 'does not record DB role metrics' do
+          expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_count_total".to_sym, any_args)
+          expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_count_total".to_sym, any_args)
+
+          expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_primary_cached_count_total".to_sym, any_args)
+          expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_replica_cached_count_total".to_sym, any_args)
+
+          expect(transaction).not_to receive(:observe).with("gitlab_sql_primary_duration_seconds".to_sym, any_args)
+          expect(transaction).not_to receive(:observe).with("gitlab_sql_replica_duration_seconds".to_sym, any_args)
+
+          subscriber.sql(event)
+        end
+
+        it 'does not store DB roles into into RequestStore' do
+          Gitlab::WithRequestStore.with_request_store do
+            subscriber.sql(event)
+
+            expect(described_class.db_counter_payload).to include(
+              db_primary_cached_count: 0,
+              db_primary_count: 0,
+              db_primary_duration_s: 0,
+              db_replica_cached_count: 0,
+              db_replica_count: 0,
+              db_replica_duration_s: 0
+            )
+          end
+        end
+      end
     end
   end
 

ee/spec/lib/peek/views/active_record_spec.rb

@@ -8,6 +8,7 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
   let(:connection_replica) { double(:connection_replica) }
   let(:connection_primary_1) { double(:connection_primary) }
   let(:connection_primary_2) { double(:connection_primary) }
+  let(:connection_unknown) { double(:connection_unknown) }
 
   let(:event_1) do
     {
@@ -36,11 +37,21 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
     }
   end
 
+  let(:event_4) do
+    {
+      name: 'SCHEMA',
+      sql: 'SELECT VERSION()',
+      cached: false,
+      connection: connection_unknown
+    }
+  end
+
   before do
     allow(Gitlab::PerformanceBar).to receive(:enabled_for_request?).and_return(true)
     allow(connection_replica).to receive(:transaction_open?).and_return(false)
     allow(connection_primary_1).to receive(:transaction_open?).and_return(false)
     allow(connection_primary_2).to receive(:transaction_open?).and_return(true)
+    allow(connection_unknown).to receive(:transaction_open?).and_return(false)
   end
 
   context 'when database load balancing is not enabled' do
@@ -49,16 +60,17 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
+        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
       end
 
       expect(subject.results).to match(
-        calls: 3,
+        calls: 4,
         summary: {
           "Cached" => 1,
           "In a transaction" => 1
         },
-        duration: '6000.00ms',
-        warnings: ["active-record duration: 6000.0 over 3000"],
+        duration: '10000.00ms',
+        warnings: ["active-record duration: 10000.0 over 3000"],
         details: contain_exactly(
           a_hash_including(
             start: be_a(Time),
@@ -80,6 +92,13 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
             transaction: 'In a transaction',
             duration: 3000.0,
             sql: 'UPDATE users SET admin = true WHERE id = 10'
+          ),
+          a_hash_including(
+            start: be_a(Time),
+            cached: '',
+            transaction: '',
+            duration: 4000.0,
+            sql: 'SELECT VERSION()'
           )
         )
       )
@@ -92,6 +111,7 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
       allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_replica).and_return(:replica)
       allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_1).and_return(:primary)
       allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_primary_2).and_return(:primary)
+      allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).with(connection_unknown).and_return(nil)
     end
 
     it 'includes db role data' do
@@ -99,18 +119,20 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 1.second, '1', event_1)
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 2.seconds, '2', event_2)
         ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 3.seconds, '3', event_3)
+        ActiveSupport::Notifications.publish('sql.active_record', Time.current, Time.current + 4.seconds, '4', event_4)
       end
 
       expect(subject.results).to match(
-        calls: 3,
+        calls: 4,
         summary: {
           "Cached" => 1,
           "In a transaction" => 1,
           "Primary" => 2,
-          "Replica" => 1
+          "Replica" => 1,
+          "Unknown" => 1
         },
-        duration: '6000.00ms',
-        warnings: ["active-record duration: 6000.0 over 3000"],
+        duration: '10000.00ms',
+        warnings: ["active-record duration: 10000.0 over 3000"],
         details: contain_exactly(
           a_hash_including(
             start: be_a(Time),
@@ -135,6 +157,14 @@ RSpec.describe Peek::Views::ActiveRecord, :request_store do
             duration: 3000.0,
             sql: 'UPDATE users SET admin = true WHERE id = 10',
             db_role: 'Primary'
+          ),
+          a_hash_including(
+            start: be_a(Time),
+            cached: '',
+            transaction: '',
+            duration: 4000.0,
+            sql: 'SELECT VERSION()',
+            db_role: 'Unknown'
           )
         )
       )