Merge branch '346909-sample-host-metrics-for-db-replicas' into 'master'

Sample database metrics for replica hosts

See merge request gitlab-org/gitlab!76282

lib/gitlab/metrics/samplers/database_sampler.rb

@@ -38,6 +38,10 @@ module Gitlab
         end
 
         def host_stats
+          connection_class_stats + replica_host_stats
+        end
+
+        def connection_class_stats
           Gitlab::Database.database_base_models.each_value.with_object([]) do |base_model, stats|
             next unless base_model.connected?
 
@@ -45,6 +49,16 @@ module Gitlab
           end
         end
 
+        def replica_host_stats
+          Gitlab::Database::LoadBalancing.each_load_balancer.with_object([]) do |load_balancer, stats|
+            next if load_balancer.primary_only?
+
+            load_balancer.host_list.hosts.each do |host|
+              stats << { labels: labels_for_replica_host(load_balancer, host), stats: host.connection.pool.stat }
+            end
+          end
+        end
+
         def labels_for_class(klass)
           {
             host: klass.connection_db_config.host,
@@ -53,6 +67,15 @@ module Gitlab
             db_config_name: klass.connection_db_config.name
           }
         end
+
+        def labels_for_replica_host(load_balancer, host)
+          {
+            host: host.host,
+            port: host.port,
+            class: load_balancer.configuration.primary_connection_specification_name,
+            db_config_name: Gitlab::Database.db_config_name(host.connection)
+          }
+        end
       end
     end
   end

spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb

@@ -8,7 +8,7 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
   it_behaves_like 'metrics sampler', 'DATABASE_SAMPLER'
 
   describe '#sample' do
-    let(:active_record_labels) do
+    let(:main_labels) do
       {
         class: 'ActiveRecord::Base',
         host: ApplicationRecord.database.config['host'],
@@ -17,7 +17,7 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
       }
     end
 
-    let(:ci_application_record_labels) do
+    let(:ci_labels) do
       {
         class: 'Ci::ApplicationRecord',
         host: Ci::ApplicationRecord.database.config['host'],
@@ -26,6 +26,24 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
       }
     end
 
+    let(:main_replica_labels) do
+      {
+        class: 'ActiveRecord::Base',
+        host: 'main-replica-host',
+        port: 2345,
+        db_config_name: 'main_replica'
+      }
+    end
+
+    let(:ci_replica_labels) do
+      {
+        class: 'Ci::ApplicationRecord',
+        host: 'ci-replica-host',
+        port: 3456,
+        db_config_name: 'ci_replica'
+      }
+    end
+
     before do
       described_class::METRIC_DESCRIPTIONS.each_key do |metric|
         allow(subject.metrics[metric]).to receive(:set)
@@ -35,35 +53,124 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
         .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord })
     end
 
-    context 'when the database is connected', :add_ci_connection do
-      it 'samples connection pool statistics' do
-        expect(subject.metrics[:size]).to receive(:set).with(active_record_labels, a_value >= 1)
-        expect(subject.metrics[:connections]).to receive(:set).with(active_record_labels, a_value >= 1)
-        expect(subject.metrics[:busy]).to receive(:set).with(active_record_labels, a_value >= 1)
-        expect(subject.metrics[:dead]).to receive(:set).with(active_record_labels, a_value >= 0)
-        expect(subject.metrics[:waiting]).to receive(:set).with(active_record_labels, a_value >= 0)
+    context 'when all base models are connected', :add_ci_connection do
+      it 'samples connection pool statistics for all primaries' do
+        expect_metrics_with_labels(main_labels)
+        expect_metrics_with_labels(ci_labels)
+
+        subject.sample
+      end
+
+      context 'when replica hosts are configured' do
+        let(:main_load_balancer) { ActiveRecord::Base.load_balancer } # rubocop:disable Database/MultipleDatabases
+        let(:main_replica_host) { main_load_balancer.host }
+
+        let(:ci_load_balancer) { double(:load_balancer, host_list: ci_host_list, configuration: configuration) }
+        let(:configuration) { double(:configuration, primary_connection_specification_name: 'Ci::ApplicationRecord') }
+        let(:ci_host_list) { double(:host_list, hosts: [ci_replica_host]) }
+        let(:ci_replica_host) { double(:host, connection: ci_connection) }
+        let(:ci_connection) { double(:connection, pool: Ci::ApplicationRecord.connection_pool) }
+
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:each_load_balancer)
+            .and_return([main_load_balancer, ci_load_balancer].to_enum)
+
+          allow(main_load_balancer).to receive(:primary_only?).and_return(false)
+          allow(ci_load_balancer).to receive(:primary_only?).and_return(false)
+
+          allow(main_replica_host).to receive(:host).and_return('main-replica-host')
+          allow(ci_replica_host).to receive(:host).and_return('ci-replica-host')
+
+          allow(main_replica_host).to receive(:port).and_return(2345)
+          allow(ci_replica_host).to receive(:port).and_return(3456)
+
+          allow(Gitlab::Database).to receive(:db_config_name)
+            .with(main_replica_host.connection)
+            .and_return('main_replica')
+
+          allow(Gitlab::Database).to receive(:db_config_name)
+            .with(ci_replica_host.connection)
+            .and_return('ci_replica')
+        end
 
-        expect(subject.metrics[:size]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
-        expect(subject.metrics[:connections]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
-        expect(subject.metrics[:busy]).to receive(:set).with(ci_application_record_labels, a_value >= 1)
-        expect(subject.metrics[:dead]).to receive(:set).with(ci_application_record_labels, a_value >= 0)
-        expect(subject.metrics[:waiting]).to receive(:set).with(ci_application_record_labels, a_value >= 0)
+        it 'samples connection pool statistics for primaries and replicas' do
+          expect_metrics_with_labels(main_labels)
+          expect_metrics_with_labels(ci_labels)
+          expect_metrics_with_labels(main_replica_labels)
+          expect_metrics_with_labels(ci_replica_labels)
 
           subject.sample
         end
       end
+    end
 
-    context 'when a database is not connected', :add_ci_connection do
+    context 'when a base model is not connected', :add_ci_connection do
       before do
         allow(Ci::ApplicationRecord).to receive(:connected?).and_return(false)
       end
 
-      it 'records no samples for that database' do
-        expect(subject.metrics[:size]).to receive(:set).with(active_record_labels, anything)
-        expect(subject.metrics[:size]).not_to receive(:set).with(ci_application_record_labels, anything)
+      it 'records no samples for that primary' do
+        expect_metrics_with_labels(main_labels)
+        expect_no_metrics_with_labels(ci_labels)
+
+        subject.sample
+      end
+
+      context 'when the base model has replica connections' do
+        let(:main_load_balancer) { ActiveRecord::Base.load_balancer } # rubocop:disable Database/MultipleDatabases
+        let(:main_replica_host) { main_load_balancer.host }
+
+        let(:ci_load_balancer) { double(:load_balancer, host_list: ci_host_list, configuration: configuration) }
+        let(:configuration) { double(:configuration, primary_connection_specification_name: 'Ci::ApplicationRecord') }
+        let(:ci_host_list) { double(:host_list, hosts: [ci_replica_host]) }
+        let(:ci_replica_host) { double(:host, connection: ci_connection) }
+        let(:ci_connection) { double(:connection, pool: Ci::ApplicationRecord.connection_pool) }
+
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:each_load_balancer)
+            .and_return([main_load_balancer, ci_load_balancer].to_enum)
+
+          allow(main_load_balancer).to receive(:primary_only?).and_return(false)
+          allow(ci_load_balancer).to receive(:primary_only?).and_return(false)
+
+          allow(main_replica_host).to receive(:host).and_return('main-replica-host')
+          allow(ci_replica_host).to receive(:host).and_return('ci-replica-host')
+
+          allow(main_replica_host).to receive(:port).and_return(2345)
+          allow(ci_replica_host).to receive(:port).and_return(3456)
+
+          allow(Gitlab::Database).to receive(:db_config_name)
+            .with(main_replica_host.connection)
+            .and_return('main_replica')
+
+          allow(Gitlab::Database).to receive(:db_config_name)
+            .with(ci_replica_host.connection)
+            .and_return('ci_replica')
+        end
+
+        it 'still records the replica metrics' do
+          expect_metrics_with_labels(main_labels)
+          expect_metrics_with_labels(main_replica_labels)
+          expect_no_metrics_with_labels(ci_labels)
+          expect_metrics_with_labels(ci_replica_labels)
 
           subject.sample
         end
       end
     end
+
+    def expect_metrics_with_labels(labels)
+      expect(subject.metrics[:size]).to receive(:set).with(labels, a_value >= 1)
+      expect(subject.metrics[:connections]).to receive(:set).with(labels, a_value >= 1)
+      expect(subject.metrics[:busy]).to receive(:set).with(labels, a_value >= 1)
+      expect(subject.metrics[:dead]).to receive(:set).with(labels, a_value >= 0)
+      expect(subject.metrics[:waiting]).to receive(:set).with(labels, a_value >= 0)
+    end
+
+    def expect_no_metrics_with_labels(labels)
+      described_class::METRIC_DESCRIPTIONS.each_key do |metric|
+        expect(subject.metrics[metric]).not_to receive(:set).with(labels, anything)
+      end
+    end
+  end
 end