Record connection pool metrics

Add a sampler for the `#stat` method of a connection pool, and record
those stats in Prometheus. This applies to all stats except
`checkout_timeout`, which is really more of a configuration item. It
also applies to all database connections we use: the main one, the Geo
tracking database, and any connection pools for database load balancing
hosts.

This requires a fix to Geo::TrackingBase.connected?, which wasn't
correctly reporting the connection's status.

config/initializers/7_prometheus_metrics.rb

@@ -44,6 +44,10 @@ if !Rails.env.test? && Gitlab::Metrics.prometheus_metrics_enabled?
 
     Gitlab::Metrics::Samplers::RubySampler.initialize_instance(Settings.monitoring.ruby_sampler_interval).start
 
+    if Gitlab::Utils.to_boolean(ENV['ENABLE_DATABASE_CONNECTION_POOL_METRICS'])
+      Gitlab::Metrics::Samplers::DatabaseSampler.initialize_instance(Gitlab::Metrics::Samplers::DatabaseSampler::SAMPLING_INTERVAL_SECONDS).start
+    end
+
     if Gitlab.ee? && Gitlab::Runtime.sidekiq?
       Gitlab::Metrics::Samplers::GlobalSearchSampler.instance(Settings.monitoring.global_search_sampler_interval).start
     end

ee/app/models/geo/tracking_base.rb

@@ -14,9 +14,9 @@ module Geo
     end
 
     def self.connected?
-      super if ::Gitlab::Geo.geo_database_configured?
+      return false unless ::Gitlab::Geo.geo_database_configured?
 
-      false
+      connection_handler.connected?(connection_specification_name)
     end
 
     def self.connection

ee/lib/ee/gitlab/metrics/samplers/database_sampler.rb

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module Metrics
+      module Samplers
+        module DatabaseSampler
+          private
+
+          def host_stats
+            super
+              .concat(geo_connection_stats)
+              .concat(load_balancing_connection_stats)
+          end
+
+          def geo_connection_stats
+            return [] unless Geo::TrackingBase.connected?
+
+            [{ labels: labels_for_class(Geo::TrackingBase), stats: Geo::TrackingBase.connection_pool.stat }]
+          end
+
+          def load_balancing_connection_stats
+            return [] unless ::Gitlab::Database::LoadBalancing.enable?
+
+            ActiveRecord::Base.connection.load_balancer.host_list.hosts.map do |host|
+              {
+                labels: { host: host.host, port: host.port, class: 'Gitlab::Database::LoadBalancing::Host' },
+                stats: host.pool.stat
+              }
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Metrics::Samplers::DatabaseSampler do
+  subject { described_class.new(described_class::SAMPLING_INTERVAL_SECONDS) }
+
+  describe '#sample' do
+    before do
+      described_class::METRIC_DESCRIPTIONS.each_key do |metric|
+        allow(subject.metrics[metric]).to receive(:set)
+      end
+    end
+
+    context 'for Geo::TrackingBase', :geo do
+      let(:labels) { a_hash_including(class: 'Geo::TrackingBase', host: anything, port: anything) }
+
+      context 'when Geo is enabled' 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 >= 0)
+          expect(subject.metrics[:busy]).to receive(:set).with(labels, a_value >= 0)
+          expect(subject.metrics[:dead]).to receive(:set).with(labels, a_value >= 0)
+          expect(subject.metrics[:waiting]).to receive(:set).with(labels, a_value >= 0)
+
+          subject.sample
+        end
+      end
+
+      context 'when Geo is not enabled' do
+        before do
+          allow(Geo::TrackingBase).to receive(:connected?).and_return(false)
+        end
+
+        it 'records no samples' do
+          expect(subject.metrics[:size]).not_to receive(:set).with(labels, anything)
+
+          subject.sample
+        end
+
+        it 'still records samples for other connections' do
+          expect(subject.metrics[:size]).to receive(:set)
+
+          subject.sample
+        end
+      end
+    end
+
+    context 'for Gitlab::Database::LoadBalancing::Host' do
+      let(:labels) { { class: 'Gitlab::Database::LoadBalancing::Host' } }
+
+      context 'when database load balancing is enabled' do
+        let(:hosts) { %w[secondary-1 secondary-2] }
+        let(:proxy) { ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts) }
+
+        before do
+          allow(ActiveRecord::Base).to receive(:connection).and_return(proxy)
+          allow(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        end
+
+        it 'samples connection pool statistics for all hosts' do
+          hosts.each do |host|
+            expected_labels = a_hash_including(host: host, **labels)
+
+            expect(subject.metrics[:size]).to receive(:set).with(expected_labels, a_value >= 1)
+            expect(subject.metrics[:connections]).to receive(:set).with(expected_labels, a_value >= 0)
+            expect(subject.metrics[:busy]).to receive(:set).with(expected_labels, a_value >= 0)
+            expect(subject.metrics[:dead]).to receive(:set).with(expected_labels, a_value >= 0)
+            expect(subject.metrics[:waiting]).to receive(:set).with(expected_labels, a_value >= 0)
+          end
+
+          subject.sample
+        end
+      end
+
+      context 'when database load balancing is not enabled' do
+        it 'records no samples' do
+          expect(subject.metrics[:size]).not_to receive(:set).with(a_hash_including(labels), anything)
+
+          subject.sample
+        end
+
+        it 'still records samples for other connections' do
+          expect(subject.metrics[:size]).to receive(:set)
+
+          subject.sample
+        end
+      end
+    end
+  end
+end

lib/gitlab/metrics/samplers/database_sampler.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Metrics
+    module Samplers
+      class DatabaseSampler < BaseSampler
+        SAMPLING_INTERVAL_SECONDS = 5
+
+        METRIC_PREFIX = 'gitlab_database_connection_pool_'
+
+        METRIC_DESCRIPTIONS = {
+          size: 'Total connection pool capacity',
+          connections: 'Current connections in the pool',
+          busy: 'Connections in use where the owner is still alive',
+          dead: 'Connections in use where the owner is not alive',
+          idle: 'Connections not in use',
+          waiting: 'Threads currently waiting on this queue'
+        }.freeze
+
+        def metrics
+          @metrics ||= init_metrics
+        end
+
+        def sample
+          host_stats.each do |host_stat|
+            METRIC_DESCRIPTIONS.each_key do |metric|
+              metrics[metric].set(host_stat[:labels], host_stat[:stats][metric])
+            end
+          end
+        end
+
+        private
+
+        def init_metrics
+          METRIC_DESCRIPTIONS.map do |name, description|
+            [name, ::Gitlab::Metrics.gauge(:"#{METRIC_PREFIX}#{name}", description)]
+          end.to_h
+        end
+
+        def host_stats
+          return [] unless ActiveRecord::Base.connected?
+
+          [{ labels: labels_for_class(ActiveRecord::Base), stats: ActiveRecord::Base.connection_pool.stat }]
+        end
+
+        def labels_for_class(klass)
+          {
+            host: klass.connection_config[:host],
+            port: klass.connection_config[:port],
+            class: klass.to_s
+          }
+        end
+      end
+    end
+  end
+end
+
+Gitlab::Metrics::Samplers::DatabaseSampler.prepend_if_ee('EE::Gitlab::Metrics::Samplers::DatabaseSampler')

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Metrics::Samplers::DatabaseSampler do
+  subject { described_class.new(described_class::SAMPLING_INTERVAL_SECONDS) }
+
+  describe '#sample' do
+    before do
+      described_class::METRIC_DESCRIPTIONS.each_key do |metric|
+        allow(subject.metrics[metric]).to receive(:set)
+      end
+    end
+
+    context 'for ActiveRecord::Base' do
+      let(:labels) do
+        {
+          class: 'ActiveRecord::Base',
+          host: Gitlab::Database.config['host'],
+          port: Gitlab::Database.config['port']
+        }
+      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)
+
+          subject.sample
+        end
+      end
+
+      context 'when the database is not connected' do
+        before do
+          allow(ActiveRecord::Base).to receive(:connected?).and_return(false)
+        end
+
+        it 'records no samples' do
+          expect(subject.metrics[:size]).not_to receive(:set).with(labels, anything)
+
+          subject.sample
+        end
+      end
+    end
+  end
+end