Log database duration per database

As part of https://gitlab.com/groups/gitlab-org/-/epics/5759 we're
preparing the GitLab application to handle connections to multiple
databases. As such we want to improve our observability to understand
these multiple databases.

This MR is a first step which will break down the total database
duration of a request by the database name of the database connection.
At present this will only be 1 extra field (`gitlabhq_production` on
production) but in future it will automatically start logging
information about the other databases we add.

We assume that database name should be sufficient for distinguishing
these connections in our logs for the time being. Technically we could
be conneting to multiple database servers with the same database name
but we don't plan on doing that first up and we can come up with a
different way to identify hosts with the same database name later if
that becomes necessary.

config/feature_flags/development/multiple_database_metrics.yml

+---
+name: multiple_database_metrics
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/63490
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/333227
+milestone: '14.1'
+type: development
+group: group::sharding
+default_enabled: false

lib/gitlab/database.rb

@@ -333,6 +333,16 @@ module Gitlab
       end
     end
 
+    def self.dbname(ar_connection)
+      if ar_connection.respond_to?(:pool) &&
+          ar_connection.pool.respond_to?(:db_config) &&
+          ar_connection.pool.db_config.respond_to?(:database)
+        return ar_connection.pool.db_config.database
+      end
+
+      'unknown'
+    end
+
     # inside_transaction? will return true if the caller is running within a transaction. Handles special cases
     # when running inside a test environment, where tests may be wrapped in transactions
     def self.inside_transaction?

lib/gitlab/metrics/subscribers/active_record.rb

@@ -72,6 +72,14 @@ module Gitlab
               DB_LOAD_BALANCING_DURATIONS.each do |duration|
                 payload[duration] = ::Gitlab::SafeRequestStore[duration].to_f.round(3)
               end
+
+              if Feature.enabled?(:multiple_database_metrics, default_enabled: :yaml)
+                ::Gitlab::SafeRequestStore[:duration_by_database]&.each do |dbname, duration_by_role|
+                  duration_by_role.each do |db_role, duration|
+                    payload[:"db_#{db_role}_#{dbname}_duration_s"] = duration.to_f.round(3)
+                  end
+                end
+              end
             end
           end
         end
@@ -93,9 +101,18 @@ module Gitlab
             buckets ::Gitlab::Metrics::Subscribers::ActiveRecord::SQL_DURATION_BUCKET
           end
 
+          return unless ::Gitlab::SafeRequestStore.active?
+
           duration = event.duration / 1000.0
           duration_key = "db_#{db_role}_duration_s".to_sym
           ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration
+
+          # Per database metrics
+          dbname = ::Gitlab::Database.dbname(event.payload[:connection])
+          ::Gitlab::SafeRequestStore[:duration_by_database] ||= {}
+          ::Gitlab::SafeRequestStore[:duration_by_database][dbname] ||= {}
+          ::Gitlab::SafeRequestStore[:duration_by_database][dbname][db_role] ||= 0
+          ::Gitlab::SafeRequestStore[:duration_by_database][dbname][db_role] += duration
         end
 
         def ignored_query?(payload)

spec/lib/gitlab/database_spec.rb

@@ -414,6 +414,23 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.dbname' do
+    it 'returns the dbname for the connection' do
+      connection = ActiveRecord::Base.connection
+
+      expect(described_class.dbname(connection)).to be_a(String)
+      expect(described_class.dbname(connection)).to eq(connection.pool.db_config.database)
+    end
+
+    context 'when the pool is a NullPool' do
+      it 'returns unknown' do
+        connection = double(:active_record_connection, pool: ActiveRecord::ConnectionAdapters::NullPool.new)
+
+        expect(described_class.dbname(connection)).to eq('unknown')
+      end
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value' do
       expect(described_class.true_value).to eq "'t'"

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

@@ -7,7 +7,7 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
   let(:env) { {} }
   let(:subscriber) { described_class.new }
-  let(:connection) { double(:connection) }
+  let(:connection) { ActiveRecord::Base.connection }
 
   describe '#transaction' do
     let(:web_transaction) { double('Gitlab::Metrics::WebTransaction') }

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -298,6 +298,8 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
           allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
         end
 
+        let(:dbname) { ::Gitlab::Database.dbname(ActiveRecord::Base.connection) }
+
         let(:expected_end_payload_with_db) do
           expected_end_payload.merge(
             'db_duration_s' => a_value >= 0.1,
@@ -311,7 +313,8 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
             'db_primary_count' => a_value >= 1,
             'db_primary_cached_count' => 0,
             'db_primary_wal_count' => 0,
-            'db_primary_duration_s' => a_value > 0
+            'db_primary_duration_s' => a_value > 0,
+            "db_primary_#{dbname}_duration_s" => a_value > 0
           )
         end
 

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

@@ -5,9 +5,10 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
     2.times do
       Gitlab::WithRequestStore.with_request_store do
         subscriber.sql(event)
+        connection = event.payload[:connection]
 
         if db_role == :primary
-          expect(described_class.db_counter_payload).to eq(
+          expected = {
             db_count: record_query ? 1 : 0,
             db_write_count: record_write_query ? 1 : 0,
             db_cached_count: record_cached_query ? 1 : 0,
@@ -19,9 +20,10 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
             db_replica_duration_s:  0.0,
             db_primary_wal_count: record_wal_query ? 1 : 0,
             db_replica_wal_count: 0
-          )
+          }
+          expected[:"db_primary_#{::Gitlab::Database.dbname(connection)}_duration_s"] = 0.002 if record_query
         elsif db_role == :replica
-          expect(described_class.db_counter_payload).to eq(
+          expected = {
             db_count: record_query ? 1 : 0,
             db_write_count: record_write_query ? 1 : 0,
             db_cached_count: record_cached_query ? 1 : 0,
@@ -33,14 +35,32 @@ RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
             db_replica_duration_s:  record_query ? 0.002 : 0,
             db_replica_wal_count: record_wal_query ? 1 : 0,
             db_primary_wal_count: 0
-          )
+          }
+          expected[:"db_replica_#{::Gitlab::Database.dbname(connection)}_duration_s"] = 0.002 if record_query
         else
-          expect(described_class.db_counter_payload).to eq(
+          expected = {
             db_count: record_query ? 1 : 0,
             db_write_count: record_write_query ? 1 : 0,
             db_cached_count: record_cached_query ? 1 : 0
-          )
+          }
         end
+
+        expect(described_class.db_counter_payload).to eq(expected)
+      end
+    end
+  end
+
+  context 'when multiple_database_metrics is disabled' do
+    before do
+      stub_feature_flags(multiple_database_metrics: false)
+    end
+
+    it 'does not include per database metrics' do
+      Gitlab::WithRequestStore.with_request_store do
+        subscriber.sql(event)
+        connection = event.payload[:connection]
+
+        expect(described_class.db_counter_payload).not_to include(:"db_replica_#{::Gitlab::Database.dbname(connection)}_duration_s")
       end
     end
   end