1 files changed, 143 insertions, 18 deletions

diff --git a/spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb b/spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb
index e97a4fdddcb..e8f8947c9e8 100644
--- a/spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb
+++ b/spec/lib/gitlab/metrics/samplers/database_sampler_spec.rb
@@ -8,44 +8,169 @@ RSpec.describe Gitlab::Metrics::Samplers::DatabaseSampler do
   it_behaves_like 'metrics sampler', 'DATABASE_SAMPLER'
 
   describe '#sample' do
+    let(:main_labels) do
+      {
+        class: 'ActiveRecord::Base',
+        host: ApplicationRecord.database.config['host'],
+        port: ApplicationRecord.database.config['port'],
+        db_config_name: 'main'
+      }
+    end
+
+    let(:ci_labels) do
+      {
+        class: 'Ci::ApplicationRecord',
+        host: Ci::ApplicationRecord.database.config['host'],
+        port: Ci::ApplicationRecord.database.config['port'],
+        db_config_name: 'ci'
+      }
+    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)
       end
+
+      allow(Gitlab::Database).to receive(:database_base_models)
+        .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord })
     end
 
-    context 'for ActiveRecord::Base' do
-      let(:labels) do
-        {
-          class: 'ActiveRecord::Base',
-          host: ApplicationRecord.database.config['host'],
-          port: ApplicationRecord.database.config['port']
-        }
+    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 the database is connected' do
-        it 'samples connection pool statistics' do
-          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)
+      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
+
+        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 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 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) }
 
-      context 'when the database is not connected' do
         before do
-          allow(ActiveRecord::Base).to receive(:connected?).and_return(false)
+          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 'records no samples' do
-          expect(subject.metrics[:size]).not_to receive(:set).with(labels, anything)
+        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