Track connection DB role in DB load balancer

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/323164

app/assets/javascripts/performance_bar/components/performance_bar_app.vue

@@ -40,7 +40,7 @@ export default {
       metric: 'active-record',
       title: 'pg',
       header: s__('PerformanceBar|SQL queries'),
-      keys: ['sql', 'cached'],
+      keys: ['sql', 'cached', 'db_role'],
     },
     {
       metric: 'bullet',

doc/administration/monitoring/prometheus/gitlab_metrics.md

@@ -34,7 +34,7 @@ For enabling and viewing metrics from Sidekiq nodes, see [Sidekiq metrics](#side
 The following metrics are available:
 
 | Metric                                                           | Type        | Since   | Description                                                                                                           | Labels                                                    |
-|:---------------------------------------------------------------|:----------|------:|:------------------------------------------------------------------------------------------------------|:--------------------------------------------------------|
+| :--------------------------------------------------------------- | :---------- | ------: | :-------------------------------------------------------------------------------------------------------------------- | :-------------------------------------------------------- |
 | `gitlab_banzai_cached_render_real_duration_seconds`              | Histogram   | 9.4     | Duration of rendering Markdown into HTML when cached output exists                                                    | `controller`, `action`                                    |
 | `gitlab_banzai_cacheless_render_real_duration_seconds`           | Histogram   | 9.4     | Duration of rendering Markdown into HTML when cached output does not exist                                            | `controller`, `action`                                    |
 | `gitlab_cache_misses_total`                                      | Counter     | 10.2    | Cache read miss                                                                                                       | `controller`, `action`                                    |
@@ -49,6 +49,7 @@ The following metrics are available:
 | `gitlab_page_out_of_bounds`                                      | Counter     | 12.8    | Counter for the PageLimiter pagination limit being hit                                                                | `controller`, `action`, `bot`                             |
 | `gitlab_rails_queue_duration_seconds`                            | Histogram   | 9.4     | Measures latency between GitLab Workhorse forwarding a request to Rails                                               |                                                           |
 | `gitlab_sql_duration_seconds`                                    | Histogram   | 10.2    | SQL execution time, excluding `SCHEMA` operations and `BEGIN` / `COMMIT`                                              |                                                           |
+| `gitlab_sql_<role>_duration_seconds`                             | Histogram   | 13.10   | SQL execution time, excluding `SCHEMA` operations and `BEGIN` / `COMMIT`, grouped by database roles (primary/replica) |                                                           |
 | `gitlab_ruby_threads_max_expected_threads`                       | Gauge       | 13.3    | Maximum number of threads expected to be running and performing application work                                      |                                                           |
 | `gitlab_ruby_threads_running_threads`                            | Gauge       | 13.3    | Number of running Ruby threads by name                                                                                |                                                           |
 | `gitlab_transaction_cache_<key>_count_total`                     | Counter     | 10.2    | Counter for total Rails cache calls (per key)                                                                         |                                                           |
@@ -94,8 +95,10 @@ The following metrics are available:
 | `http_requests_total`                                            | Counter     | 9.4     | Rack request count                                                                                                    | `method`, `status`                                        |
 | `http_request_duration_seconds`                                  | Histogram   | 9.4     | HTTP response time from rack middleware                                                                               | `method`                                                  |
 | `gitlab_transaction_db_count_total`                              | Counter     | 13.1    | Counter for total number of SQL calls                                                                                 | `controller`, `action`                                    |
+| `gitlab_transaction_db_<role>_count_total`                       | Counter     | 13.10   | Counter for total number of SQL calls, grouped by database roles (primary/replica)                                    | `controller`, `action`                                    |
 | `gitlab_transaction_db_write_count_total`                        | Counter     | 13.1    | Counter for total number of write SQL calls                                                                           | `controller`, `action`                                    |
 | `gitlab_transaction_db_cached_count_total`                       | Counter     | 13.1    | Counter for total number of cached SQL calls                                                                          | `controller`, `action`                                    |
+| `gitlab_transaction_db_<role>_cached_count_total`                | Counter     | 13.1    | Counter for total number of cached SQL calls, grouped by database roles (primary/replica)                             | `controller`, `action`                                    |
 | `http_elasticsearch_requests_duration_seconds` **(PREMIUM)**     | Histogram   | 13.1    | Elasticsearch requests duration during web transactions                                                               | `controller`, `action`                                    |
 | `http_elasticsearch_requests_total` **(PREMIUM)**                | Counter     | 13.1    | Elasticsearch requests count during web transactions                                                                  | `controller`, `action`                                    |
 | `pipelines_created_total`                                        | Counter     | 9.4     | Counter of pipelines created                                                                                          |                                                           |

ee/changelogs/unreleased/873-split-structured-logs-db_duration-fields-between-primary-database-.yml

+---
+title: Differentiate metrics and logs from replica/primary databases
+merge_request: 54885
+author:
+type: changed

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

+# frozen_string_literal: true
+
+module EE
+  module Gitlab
+    module Metrics
+      module Subscribers
+        module ActiveRecord
+          extend ActiveSupport::Concern
+          extend ::Gitlab::Utils::Override
+
+          DB_LOAD_BALANCING_COUNTERS = %i{
+            db_replica_count db_replica_cached_count
+            db_primary_count db_primary_cached_count
+          }.freeze
+          DB_LOAD_BALANCING_DURATIONS = %i{db_primary_duration_s db_replica_duration_s}.freeze
+
+          class_methods do
+            extend ::Gitlab::Utils::Override
+
+            override :db_counter_payload
+            def db_counter_payload
+              super.tap do |payload|
+                if ::Gitlab::SafeRequestStore.active? && ::Gitlab::Database::LoadBalancing.enable?
+                  DB_LOAD_BALANCING_COUNTERS.each do |counter|
+                    payload[counter] = ::Gitlab::SafeRequestStore[counter].to_i
+                  end
+                  DB_LOAD_BALANCING_DURATIONS.each do |duration|
+                    payload[duration] = ::Gitlab::SafeRequestStore[duration].to_f.round(6)
+                  end
+                end
+              end
+            end
+          end
+
+          override :sql
+          def sql(event)
+            super
+
+            return unless ::Gitlab::Database::LoadBalancing.enable?
+
+            payload = event.payload
+            return if ignored_query?(payload)
+
+            db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
+            increment_db_role_counters(db_role, payload)
+            observe_db_role_duration(db_role, event.duration)
+          end
+
+          private
+
+          def increment_db_role_counters(db_role, payload)
+            increment("db_#{db_role}_count".to_sym)
+            increment("db_#{db_role}_cached_count".to_sym) if cached_query?(payload)
+          end
+
+          def observe_db_role_duration(db_role, duration)
+            duration /= 1000.0
+
+            current_transaction&.observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, duration) do
+              buckets [0.05, 0.1, 0.25]
+            end
+
+            duration_key = "db_#{db_role}_duration_s".to_sym
+            ::Gitlab::SafeRequestStore[duration_key] = (::Gitlab::SafeRequestStore[duration_key].presence || 0) + duration
+          end
+        end
+      end
+    end
+  end
+end

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

+# frozen_string_literal: true
+
+module EE
+  module Peek
+    module Views
+      module ActiveRecord
+        extend ::Gitlab::Utils::Override
+
+        override :generate_detail
+        def generate_detail(start, finish, data)
+          detail = super
+
+          if ::Gitlab::Database::LoadBalancing.enable?
+            detail[:db_role] = ::Gitlab::Database::LoadBalancing.db_role_for_connection(data[:connection])
+          end
+
+          detail
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/database/load_balancing.rb

@@ -96,11 +96,31 @@ module Gitlab
       # posting the write location of the database if load balancing is
       # configured.
       def self.configured?
-        return false unless ::License.feature_available?(:db_load_balancing)
+        return false unless feature_available?
 
         hosts.any? || service_discovery_enabled?
       end
 
+      def self.feature_available?
+        # If this method is called in any subscribers listening to
+        # sql.active_record, the SQL call below may cause infinite recursion.
+        # So, the memoization variable must have 3 states
+        # - First call: @feature_available is undefined
+        #   -> Set @feature_available to false
+        #   -> Trigger SQL
+        #   -> SQL subscriber triggers this method again
+        #     -> return false
+        #   -> Set @feature_available  to true
+        #   -> return true
+        # - Second call: return @feature_available right away
+
+        return @feature_available if defined?(@feature_available)
+
+        @feature_available = false
+        @feature_available = Gitlab::Database.cached_table_exists?('licenses') &&
+                             ::License.feature_available?(:db_load_balancing)
+      end
+
       def self.program_name
         @program_name ||= File.basename($0)
       end
@@ -121,9 +141,32 @@ module Gitlab
         ActiveRecord::Base.singleton_class.prepend(ActiveRecordProxy)
       end
 
+      # Clear configuration
+      def self.clear_configuration
+        @proxy = nil
+        remove_instance_variable(:@feature_available)
+      end
+
       def self.active_record_models
         ActiveRecord::Base.descendants
       end
+
+      DB_ROLES = [
+        ROLE_PRIMARY = :primary,
+        ROLE_REPLICA = :replica
+      ].freeze
+
+      # Returns the role (primary/replica) of the database the connection is
+      # connecting to. At the moment, the connection can only be retrieved by
+      # Gitlab::Database::LoadBalancer#read or #read_write or from the
+      # ActiveRecord directly. Therefore, if the load balancer doesn't
+      # recognize the connection, this method returns the primary role
+      # directly. In future, we may need to check for other sources.
+      def self.db_role_for_connection(connection)
+        return ROLE_PRIMARY if !enable? || @proxy.blank?
+
+        proxy.load_balancer.db_role_for_connection(connection) || ROLE_PRIMARY
+      end
     end
   end
 end

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

@@ -18,6 +18,7 @@ module Gitlab
         # hosts - The hostnames/addresses of the additional databases.
         def initialize(hosts = [])
           @host_list = HostList.new(hosts.map { |addr| Host.new(addr, self) })
+          @connection_db_roles = {}
         end
 
         # Yields a connection that can be used for reads.
@@ -25,14 +26,20 @@ module Gitlab
         # If no secondaries were available this method will use the primary
         # instead.
         def read(&block)
+          connection = nil
           conflict_retried = 0
 
           while host
             ensure_caching!
 
             begin
-              return yield host.connection
+              connection = host.connection
+              @connection_db_roles[connection.object_id] = ROLE_REPLICA
+
+              return yield connection
             rescue => error
+              @connection_db_roles.delete(connection.object_id) if connection.present?
+
               if serialization_failure?(error)
                 # This error can occur when a query conflicts. See
                 # https://www.postgresql.org/docs/current/static/hot-standby.html#HOT-STANDBY-CONFLICT
@@ -75,16 +82,31 @@ module Gitlab
           )
 
           read_write(&block)
+        ensure
+          @connection_db_roles.delete(connection.object_id) if connection.present?
         end
 
         # Yields a connection that can be used for both reads and writes.
         def read_write
+          connection = nil
           # In the event of a failover the primary may be briefly unavailable.
           # Instead of immediately grinding to a halt we'll retry the operation
           # a few times.
           retry_with_backoff do
-            yield ActiveRecord::Base.retrieve_connection
+            connection = ActiveRecord::Base.retrieve_connection
+            @connection_db_roles[connection.object_id] = ROLE_PRIMARY
+
+            yield connection
           end
+        ensure
+          @connection_db_roles.delete(connection.object_id) if connection.present?
+        end
+
+        # Recognize the role (primary/replica) of the database this connection
+        # 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.object_id]
         end
 
         # Returns a host to use for queries.

ee/spec/finders/epics_finder_spec.rb

@@ -570,8 +570,9 @@ RSpec.describe EpicsFinder do
 
             # if user is not member of top-level group, we need to check
             # if he can read epics in each subgroup
-            it 'does not execute more than 10 SQL queries' do
-              expect { subject }.not_to exceed_all_query_limit(10)
+            it 'does not execute more than 15 SQL queries' do
+              # The limit here is fragile!
+              expect { subject }.not_to exceed_all_query_limit(15)
             end
 
             it 'checks permission for each subgroup' do

ee/spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -23,7 +23,7 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
 
     context 'with load balancing disabled', :request_store, :redis do
       before do
-        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(false)
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(false)
         expect(::Gitlab::Database::LoadBalancing::Sticking).not_to receive(:unstick_or_continue_sticking)
       end
 
@@ -43,7 +43,7 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
       let(:all_caught_up) { true }
 
       before do
-        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(true)
         expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:unstick_or_continue_sticking).and_call_original
         allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
       end

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

@@ -5,12 +5,17 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   let(:pool_spec) { ActiveRecord::Base.connection_pool.spec }
   let(:pool) { ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_spec) }
+  let(:conflict_error) { Class.new(RuntimeError) }
 
   let(:lb) { described_class.new(%w(localhost localhost)) }
 
   before do
     allow(Gitlab::Database).to receive(:create_connection_pool)
       .and_return(pool)
+    stub_const(
+      'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
+      conflict_error
+    )
   end
 
   def raise_and_wrap(wrapper, original)
@@ -36,15 +41,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   end
 
   describe '#read' do
-    let(:conflict_error) { Class.new(RuntimeError) }
-
-    before do
-      stub_const(
-        'Gitlab::Database::LoadBalancing::LoadBalancer::PG::TRSerializationFailure',
-        conflict_error
-      )
-    end
-
     it 'yields a connection for a read' do
       connection = double(:connection)
       host = double(:host)
@@ -139,6 +135,78 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
   end
 
+  describe '#db_role_for_connection' do
+    context 'when the load balancer creates the connection with #read' do
+      it 'returns :replica' do
+        role = nil
+        lb.read do |connection|
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:replica)
+      end
+    end
+
+    context 'when the load balancer creates the connection with #read_write' do
+      it 'returns :primary' do
+        role = nil
+        lb.read_write do |connection|
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:primary)
+      end
+    end
+
+    context 'when the load balancer falls back the connection creation to primary' do
+      it 'returns :primary' do
+        allow(lb).to receive(:serialization_failure?).and_return(true)
+
+        role = nil
+        raised = 7 # 2 hosts = 6 retries
+
+        lb.read do |connection|
+          if raised > 0
+            raised -= 1
+            raise
+          end
+
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:primary)
+      end
+    end
+
+    context 'when the load balancer uses replica after recovery from a failure' do
+      it 'returns :replica' do
+        allow(lb).to receive(:connection_error?).and_return(true)
+
+        role = nil
+        raised = false
+
+        lb.read do |connection|
+          unless raised
+            raised = true
+            raise
+          end
+
+          role = lb.db_role_for_connection(connection)
+        end
+
+        expect(role).to be(:replica)
+      end
+    end
+
+    context 'when the connection does not come from the load balancer' do
+      it 'returns nil' do
+        connection = double(:connection)
+
+        expect(lb.db_role_for_connection(connection)).to be(nil)
+      end
+    end
+  end
+
   describe '#host' do
     it 'returns the secondary host to use' do
       expect(lb.host).to be_an_instance_of(Gitlab::Database::LoadBalancing::Host)

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

@@ -11,7 +11,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       after do
-        subject.configure_proxy(nil)
+        subject.clear_configuration
       end
 
       it 'returns the connection proxy' do
@@ -132,6 +132,10 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   describe '.enable?' do
     let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
 
+    before do
+      subject.clear_configuration
+    end
+
     it 'returns false when no hosts are specified' do
       allow(described_class).to receive(:hosts).and_return([])
 
@@ -250,7 +254,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
   describe '.configure_proxy' do
     after do
-      described_class.configure_proxy(nil)
+      described_class.clear_configuration
     end
 
     it 'configures the connection proxy' do
@@ -343,4 +347,67 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       described_class.start_service_discovery
     end
   end
+
+  describe '.db_role_for_connection' do
+    let(:connection) { double(:conneciton) }
+
+    context 'when the load balancing is not configured' do
+      before do
+        allow(described_class).to receive(:enable?).and_return(false)
+      end
+
+      it 'returns primary' do
+        expect(described_class.db_role_for_connection(connection)).to be(:primary)
+      end
+    end
+
+    context 'when the load balancing is configured' do
+      let(:proxy) { described_class::ConnectionProxy.new(%w(foo)) }
+      let(:load_balancer) { described_class::LoadBalancer.new(%w(foo)) }
+
+      before do
+        allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
+
+        allow(described_class).to receive(:enable?).and_return(true)
+        allow(described_class).to receive(:proxy).and_return(proxy)
+        allow(proxy).to receive(:load_balancer).and_return(load_balancer)
+
+        subject.configure_proxy(proxy)
+      end
+
+      after do
+        subject.clear_configuration
+      end
+
+      context 'when the load balancer returns :replica' do
+        it 'returns :replica' do
+          allow(load_balancer).to receive(:db_role_for_connection).and_return(:replica)
+
+          expect(described_class.db_role_for_connection(connection)).to be(:replica)
+
+          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+
+      context 'when the load balancer returns :primary' do
+        it 'returns :primary' do
+          allow(load_balancer).to receive(:db_role_for_connection).and_return(:primary)
+
+          expect(described_class.db_role_for_connection(connection)).to be(:primary)
+
+          expect(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+
+      context 'when the load balancer returns nil' do
+        it 'returns :primary' 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(load_balancer).to have_received(:db_role_for_connection).with(connection)
+        end
+      end
+    end
+  end
 end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::Gitlab::Metrics::Subscribers::ActiveRecord do
+  using RSpec::Parameterized::TableSyntax
+
+  let(:env) { {} }
+  let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
+  let(:subscriber) { described_class.new }
+  let(:connection) { double(:connection) }
+  let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } }
+
+  let(:event) do
+    double(
+      :event,
+      name: 'sql.active_record',
+      duration: 2,
+      payload:  payload
+    )
+  end
+
+  # Emulate Marginalia pre-pending comments
+  def sql(query, comments: true)
+    if comments && !%w[BEGIN COMMIT].include?(query)
+      "/*application:web,controller:badges,action:pipeline,correlation_id:01EYN39K9VMJC56Z7808N7RSRH*/ #{query}"
+    else
+      query
+    end
+  end
+
+  shared_examples 'track sql events for each role' do
+    where(:name, :sql_query, :record_query, :record_write_query, :record_cached_query) do
+      'SQL' | 'SELECT * FROM users WHERE id = 10' | true | false | false
+      'SQL' | 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' | true | false | false
+      'SQL' | 'SELECT * FROM users WHERE id = 10 FOR UPDATE' | true | true | false
+      'SQL' | 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' | true | true | false
+      'SQL' | 'DELETE FROM users where id = 10' | true | true | false
+      'SQL' | 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' | true | true | false
+      'SQL' | 'UPDATE users SET admin = true WHERE id = 10' | true | true | false
+      'CACHE' | 'SELECT * FROM users WHERE id = 10' | true | false | true
+      'SCHEMA' | "SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass" | false | false | false
+      nil | 'BEGIN' | false | false | false
+      nil | 'COMMIT' | false | false | false
+    end
+
+    with_them do
+      let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
+
+      before do
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        allow(subscriber).to receive(:current_transaction)
+          .and_return(transaction)
+      end
+
+      context 'query using a connection to a replica' do
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:replica)
+        end
+
+        it 'queries connection db role' do
+          subscriber.sql(event)
+
+          if record_query
+            expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection)
+          end
+        end
+
+        it_behaves_like 'record ActiveRecord metrics', :replica
+        it_behaves_like 'store ActiveRecord info in RequestStore', :replica
+      end
+
+      context 'query using a connection to a primary' do
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:db_role_for_connection).and_return(:primary)
+        end
+
+        it 'queries connection db role' do
+          subscriber.sql(event)
+
+          if record_query
+            expect(Gitlab::Database::LoadBalancing).to have_received(:db_role_for_connection).with(connection)
+          end
+        end
+
+        it_behaves_like 'record ActiveRecord metrics', :primary
+        it_behaves_like 'store ActiveRecord info in RequestStore', :primary
+      end
+    end
+  end
+
+  context 'without Marginalia comments' do
+    let(:comments) { false }
+
+    it_behaves_like 'track sql events for each role'
+  end
+
+  context 'with Marginalia comments' do
+    let(:comments) { true }
+
+    it_behaves_like 'track sql events for each role'
+  end
+end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Peek::Views::ActiveRecord, :request_store do
+  subject { Peek.views.find { |v| v.class.name == Peek::Views::ActiveRecord.name } }
+
+  let(:connection_replica) { double(:connection_replica) }
+  let(:connection_primary) { double(:connection_primary) }
+
+  let(:event_1) do
+    {
+      name: 'SQL',
+      sql: 'SELECT * FROM users WHERE id = 10',
+      cached: false,
+      connection: connection_primary
+    }
+  end
+
+  let(:event_2) do
+    {
+      name: 'SQL',
+      sql: 'SELECT * FROM users WHERE id = 10',
+      cached: true,
+      connection: connection_replica
+    }
+  end
+
+  let(:event_3) do
+    {
+      name: 'SQL',
+      sql: 'UPDATE users SET admin = true WHERE id = 10',
+      cached: false,
+      connection: connection_primary
+    }
+  end
+
+  before do
+    allow(Gitlab::PerformanceBar).to receive(:enabled_for_request?).and_return(true)
+  end
+
+  context 'when database load balancing is not enabled' do
+    it 'subscribes and store data into peek views' do
+      Timecop.freeze(2021, 2, 23, 10, 0) 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)
+      end
+
+      expect(subject.results).to match(
+        calls: '3 (1 cached)',
+        duration: '6000.00ms',
+        warnings: ["active-record duration: 6000.0 over 3000"],
+        details: contain_exactly(
+          a_hash_including(
+            cached: '',
+            duration: 1000.0,
+            sql: 'SELECT * FROM users WHERE id = 10'
+          ),
+          a_hash_including(
+            cached: 'cached',
+            duration: 2000.0,
+            sql: 'SELECT * FROM users WHERE id = 10'
+          ),
+          a_hash_including(
+            cached: '',
+            duration: 3000.0,
+            sql: 'UPDATE users SET admin = true WHERE id = 10'
+          )
+        )
+      )
+    end
+  end
+
+  context 'when database load balancing is enabled' do
+    before do
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+      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).and_return(:primary)
+    end
+
+    it 'includes db role data' do
+      Timecop.freeze(2021, 2, 23, 10, 0) 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)
+      end
+
+      expect(subject.results).to match(
+        calls: '3 (1 cached)',
+        duration: '6000.00ms',
+        warnings: ["active-record duration: 6000.0 over 3000"],
+        details: contain_exactly(
+          a_hash_including(
+            cached: '',
+            duration: 1000.0,
+            sql: 'SELECT * FROM users WHERE id = 10',
+            db_role: :primary
+          ),
+          a_hash_including(
+            cached: 'cached',
+            duration: 2000.0,
+            sql: 'SELECT * FROM users WHERE id = 10',
+            db_role: :replica
+          ),
+          a_hash_including(
+            cached: '',
+            duration: 3000.0,
+            sql: 'UPDATE users SET admin = true WHERE id = 10',
+            db_role: :primary
+          )
+        )
+      )
+    end
+  end
+end

lib/gitlab/metrics/subscribers/active_record.rb

@@ -18,10 +18,9 @@ module Gitlab
           Thread.current[:uses_db_connection] = true
 
           payload = event.payload
-          return if payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql])
+          return if ignored_query?(payload)
 
           increment_db_counters(payload)
-
           current_transaction&.observe(:gitlab_sql_duration_seconds, event.duration / 1000.0) do
             buckets [0.05, 0.1, 0.25]
           end
@@ -30,34 +29,38 @@ module Gitlab
         def self.db_counter_payload
           return {} unless Gitlab::SafeRequestStore.active?
 
-          DB_COUNTERS.map do |counter|
-            [counter, Gitlab::SafeRequestStore[counter].to_i]
-          end.to_h
+          payload = {}
+          DB_COUNTERS.each do |counter|
+            payload[counter] = Gitlab::SafeRequestStore[counter].to_i
+          end
+          payload
         end
 
         private
 
+        def ignored_query?(payload)
+          payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql])
+        end
+
+        def cached_query?(payload)
+          payload.fetch(:cached, payload[:name] == 'CACHE')
+        end
+
         def select_sql_command?(payload)
           payload[:sql].match(SQL_COMMANDS_WITH_COMMENTS_REGEX)
         end
 
         def increment_db_counters(payload)
           increment(:db_count)
-
-          if payload.fetch(:cached, payload[:name] == 'CACHE')
-            increment(:db_cached_count)
-          end
-
+          increment(:db_cached_count) if cached_query?(payload)
           increment(:db_write_count) unless select_sql_command?(payload)
         end
 
         def increment(counter)
           current_transaction&.increment("gitlab_transaction_#{counter}_total".to_sym, 1)
 
-          if Gitlab::SafeRequestStore.active?
           Gitlab::SafeRequestStore[counter] = Gitlab::SafeRequestStore[counter].to_i + 1
         end
-        end
 
         def current_transaction
           ::Gitlab::Metrics::Transaction.current
@@ -66,3 +69,5 @@ module Gitlab
     end
   end
 end
+
+Gitlab::Metrics::Subscribers::ActiveRecord.prepend_if_ee('EE::Gitlab::Metrics::Subscribers::ActiveRecord')

lib/peek/views/active_record.rb

@@ -39,8 +39,12 @@ module Peek
         super
 
         subscribe('sql.active_record') do |_, start, finish, _, data|
-          if Gitlab::PerformanceBar.enabled_for_request?
-            detail_store << {
+          detail_store << generate_detail(start, finish, data) if Gitlab::PerformanceBar.enabled_for_request?
+        end
+      end
+
+      def generate_detail(start, finish, data)
+        {
           duration: finish - start,
           sql: data[:sql].strip,
           backtrace: Gitlab::BacktraceCleaner.clean_backtrace(caller),
@@ -49,6 +53,6 @@ module Peek
       end
     end
   end
-    end
-  end
 end
+
+Peek::Views::ActiveRecord.prepend_if_ee('EE::Peek::Views::ActiveRecord')

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

@@ -3,10 +3,13 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
+  using RSpec::Parameterized::TableSyntax
+
   let(:env) { {} }
   let(:transaction) { Gitlab::Metrics::WebTransaction.new(env) }
   let(:subscriber)  { described_class.new }
-  let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10' } }
+  let(:connection) { double(:connection) }
+  let(:payload) { { sql: 'SELECT * FROM users WHERE id = 10', connection: connection } }
 
   let(:event) do
     double(
@@ -19,88 +22,36 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
 
   # Emulate Marginalia pre-pending comments
   def sql(query, comments: true)
-    if comments
+    if comments && !%w[BEGIN COMMIT].include?(query)
       "/*application:web,controller:badges,action:pipeline,correlation_id:01EYN39K9VMJC56Z7808N7RSRH*/ #{query}"
     else
       query
     end
   end
 
-  describe '#sql' do
-    shared_examples 'track query in metrics' do
-      before do
-        allow(subscriber).to receive(:current_transaction)
-                               .at_least(:once)
-                               .and_return(transaction)
-      end
-
-      it 'increments only db count value' do
-        described_class::DB_COUNTERS.each do |counter|
-          prometheus_counter = "gitlab_transaction_#{counter}_total".to_sym
-          if expected_counters[counter] > 0
-            expect(transaction).to receive(:increment).with(prometheus_counter, 1)
-          else
-            expect(transaction).not_to receive(:increment).with(prometheus_counter, 1)
-          end
-        end
-
-        subscriber.sql(event)
-      end
-    end
-
-    shared_examples 'track query in RequestStore' do
-      context 'when RequestStore is enabled' do
-        it 'caches db count value', :request_store, :aggregate_failures do
-          subscriber.sql(event)
-
-          described_class::DB_COUNTERS.each do |counter|
-            expect(Gitlab::SafeRequestStore[counter].to_i).to eq expected_counters[counter]
-          end
-        end
-
-        it 'prevents db counters from leaking to the next transaction' do
-          2.times do
-            Gitlab::WithRequestStore.with_request_store do
-              subscriber.sql(event)
-
-              described_class::DB_COUNTERS.each do |counter|
-                expect(Gitlab::SafeRequestStore[counter].to_i).to eq expected_counters[counter]
-              end
-            end
-          end
-        end
-      end
-    end
-
-    describe 'without a current transaction' do
-      it 'does not track any metrics' do
-        expect_any_instance_of(Gitlab::Metrics::Transaction)
-          .not_to receive(:increment)
-
-        subscriber.sql(event)
-      end
-
-      context 'with read query' do
-        let(:expected_counters) do
-          {
-            db_count: 1,
-            db_write_count: 0,
-            db_cached_count: 0
-          }
+  shared_examples 'track generic sql events' do
+    where(:name, :sql_query, :record_query, :record_write_query, :record_cached_query) do
+      'SQL' | 'SELECT * FROM users WHERE id = 10' | true | false | false
+      'SQL' | 'WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones' | true | false | false
+      'SQL' | 'SELECT * FROM users WHERE id = 10 FOR UPDATE' | true | true | false
+      'SQL' | 'WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows' | true | true | false
+      'SQL' | 'DELETE FROM users where id = 10' | true | true | false
+      'SQL' | 'INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects' | true | true | false
+      'SQL' | 'UPDATE users SET admin = true WHERE id = 10' | true | true | false
+      'CACHE' | 'SELECT * FROM users WHERE id = 10' | true | false | true
+      'SCHEMA' | "SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass" | false | false | false
+      nil | 'BEGIN' | false | false | false
+      nil | 'COMMIT' | false | false | false
     end
 
-        it_behaves_like 'track query in RequestStore'
-      end
-    end
+    with_them do
+      let(:payload) { { name: name, sql: sql(sql_query, comments: comments), connection: connection } }
 
       describe 'with a current transaction' do
-      it 'observes sql_duration metric' do
-        expect(subscriber).to receive(:current_transaction)
+        before do
+          allow(subscriber).to receive(:current_transaction)
             .at_least(:once)
             .and_return(transaction)
-        expect(transaction).to receive(:observe).with(:gitlab_sql_duration_seconds, 0.002)
-
-        subscriber.sql(event)
         end
 
         it 'marks the current thread as using the database' do
@@ -110,229 +61,31 @@ RSpec.describe Gitlab::Metrics::Subscribers::ActiveRecord do
           expect { subscriber.sql(event) }.to change { Thread.current[:uses_db_connection] }.from(nil).to(true)
         end
 
-      shared_examples 'read queries' do
-        let(:expected_counters) do
-          {
-            db_count: 1,
-            db_write_count: 0,
-            db_cached_count: 0
-          }
-        end
-
-        it_behaves_like 'track query in metrics'
-        it_behaves_like 'track query in RequestStore'
-
-        context 'with only select' do
-          let(:payload) { { sql: sql('WITH active_milestones AS (SELECT COUNT(*), state FROM milestones GROUP BY state) SELECT * FROM active_milestones', comments: comments) } }
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
-        end
-      end
-
-      shared_examples 'write queries' do
-        let(:expected_counters) do
-          {
-            db_count: 1,
-            db_write_count: 1,
-            db_cached_count: 0
-          }
+        it_behaves_like 'record ActiveRecord metrics'
+        it_behaves_like 'store ActiveRecord info in RequestStore'
       end
 
-        context 'with select for update sql event' do
-          let(:payload) { { sql: sql('SELECT * FROM users WHERE id = 10 FOR UPDATE', comments: comments) } }
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
-        end
-
-        context 'with common table expression' do
-          context 'with insert' do
-            let(:payload) { { sql: sql('WITH archived_rows AS (SELECT * FROM users WHERE archived = true) INSERT INTO products_log SELECT * FROM archived_rows', comments: comments) } }
-
-            it_behaves_like 'track query in metrics'
-            it_behaves_like 'track query in RequestStore'
-          end
-        end
-
-        context 'with delete sql event' do
-          let(:payload) { { sql: sql('DELETE FROM users where id = 10', comments: comments) } }
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
+      describe 'without a current transaction' do
+        it 'does not track any metrics' do
+          expect_any_instance_of(Gitlab::Metrics::Transaction)
+            .not_to receive(:increment)
+          subscriber.sql(event)
         end
 
-        context 'with insert sql event' do
-          let(:payload) { { sql: sql('INSERT INTO project_ci_cd_settings (project_id) SELECT id FROM projects', comments: comments) } }
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
+        it_behaves_like 'store ActiveRecord info in RequestStore'
       end
-
-        context 'with update sql event' do
-          let(:payload) { { sql: sql('UPDATE users SET admin = true WHERE id = 10', comments: comments) } }
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
     end
   end
 
   context 'without Marginalia comments' do
     let(:comments) { false }
 
-        it_behaves_like 'write queries'
-        it_behaves_like 'read queries'
+    it_behaves_like 'track generic sql events'
   end
 
   context 'with Marginalia comments' do
     let(:comments) { true }
 
-        it_behaves_like 'write queries'
-        it_behaves_like 'read queries'
-      end
-
-      context 'with cached query' do
-        let(:expected_counters) do
-          {
-            db_count: 1,
-            db_write_count: 0,
-            db_cached_count: 1
-          }
-        end
-
-        context 'with cached payload ' do
-          let(:payload) do
-            {
-              sql: sql('SELECT * FROM users WHERE id = 10'),
-              cached: true
-            }
-          end
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
-        end
-
-        context 'with cached payload name' do
-          let(:payload) do
-            {
-              sql: sql('SELECT * FROM users WHERE id = 10'),
-              name: 'CACHE'
-            }
-          end
-
-          it_behaves_like 'track query in metrics'
-          it_behaves_like 'track query in RequestStore'
-        end
-      end
-
-      context 'events are internal to Rails or irrelevant' do
-        let(:schema_event) do
-          double(
-            :event,
-            name: 'sql.active_record',
-            payload: {
-              sql: sql("SELECT attr.attname FROM pg_attribute attr INNER JOIN pg_constraint cons ON attr.attrelid = cons.conrelid AND attr.attnum = any(cons.conkey) WHERE cons.contype = 'p' AND cons.conrelid = '\"projects\"'::regclass"),
-              name: 'SCHEMA',
-              connection_id: 135,
-              statement_name: nil,
-              binds: []
-            },
-            duration: 0.7
-          )
-        end
-
-        let(:begin_event) do
-          double(
-            :event,
-            name: 'sql.active_record',
-            payload: {
-              sql: "BEGIN",
-              name: nil,
-              connection_id: 231,
-              statement_name: nil,
-              binds: []
-            },
-            duration: 1.1
-          )
-        end
-
-        let(:commit_event) do
-          double(
-            :event,
-            name: 'sql.active_record',
-            payload: {
-              sql: "COMMIT",
-              name: nil,
-              connection_id: 212,
-              statement_name: nil,
-              binds: []
-            },
-            duration: 1.6
-          )
-        end
-
-        it 'skips schema/begin/commit sql commands' do
-          allow(subscriber).to receive(:current_transaction)
-                                 .at_least(:once)
-                                 .and_return(transaction)
-
-          expect(transaction).not_to receive(:increment)
-
-          subscriber.sql(schema_event)
-          subscriber.sql(begin_event)
-          subscriber.sql(commit_event)
-        end
-      end
-    end
-  end
-
-  describe 'self.db_counter_payload' do
-    before do
-      allow(subscriber).to receive(:current_transaction)
-                             .at_least(:once)
-                             .and_return(transaction)
-    end
-
-    context 'when RequestStore is enabled', :request_store do
-      context 'when query is executed' do
-        let(:expected_payload) do
-          {
-            db_count: 1,
-            db_cached_count: 0,
-            db_write_count: 0
-          }
-        end
-
-        it 'returns correct payload' do
-          subscriber.sql(event)
-
-          expect(described_class.db_counter_payload).to eq(expected_payload)
-        end
-      end
-
-      context 'when query is not executed' do
-        let(:expected_payload) do
-          {
-            db_count: 0,
-            db_cached_count: 0,
-            db_write_count: 0
-          }
-        end
-
-        it 'returns correct payload' do
-          expect(described_class.db_counter_payload).to eq(expected_payload)
-        end
-      end
-    end
-
-    context 'when RequestStore is disabled' do
-      let(:expected_payload) { {} }
-
-      it 'returns empty payload' do
-        subscriber.sql(event)
-
-        expect(described_class.db_counter_payload).to eq(expected_payload)
-      end
-    end
+    it_behaves_like 'track generic sql events'
   end
 end

spec/lib/peek/views/active_record_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Peek::Views::ActiveRecord, :request_store do
+  subject { Peek.views.find { |v| v.class.name == Peek::Views::ActiveRecord.name } }
+
+  let(:connection) { double(:connection) }
+
+  let(:event_1) do
+    {
+      name: 'SQL',
+      sql: 'SELECT * FROM users WHERE id = 10',
+      cached: false,
+      connection: connection
+    }
+  end
+
+  let(:event_2) do
+    {
+      name: 'SQL',
+      sql: 'SELECT * FROM users WHERE id = 10',
+      cached: true,
+      connection: connection
+    }
+  end
+
+  let(:event_3) do
+    {
+      name: 'SQL',
+      sql: 'UPDATE users SET admin = true WHERE id = 10',
+      cached: false,
+      connection: connection
+    }
+  end
+
+  before do
+    allow(Gitlab::PerformanceBar).to receive(:enabled_for_request?).and_return(true)
+  end
+
+  it 'subscribes and store data into peek views' do
+    Timecop.freeze(2021, 2, 23, 10, 0) 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)
+    end
+
+    expect(subject.results).to match(
+      calls: '3 (1 cached)',
+      duration: '6000.00ms',
+      warnings: ["active-record duration: 6000.0 over 3000"],
+      details: contain_exactly(
+        a_hash_including(
+          cached: '',
+          duration: 1000.0,
+          sql: 'SELECT * FROM users WHERE id = 10'
+        ),
+        a_hash_including(
+          cached: 'cached',
+          duration: 2000.0,
+          sql: 'SELECT * FROM users WHERE id = 10'
+        ),
+        a_hash_including(
+          cached: '',
+          duration: 3000.0,
+          sql: 'UPDATE users SET admin = true WHERE id = 10'
+        )
+      )
+    )
+  end
+end

spec/support/shared_examples/metrics/active_record_subscriber_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'store ActiveRecord info in RequestStore' do |db_role|
+  it 'prevents db counters from leaking to the next transaction' do
+    2.times do
+      Gitlab::WithRequestStore.with_request_store do
+        subscriber.sql(event)
+
+        if db_role == :primary
+          expect(described_class.db_counter_payload).to eq(
+            db_count: record_query ? 1 : 0,
+            db_write_count: record_write_query ? 1 : 0,
+            db_cached_count: record_cached_query ? 1 : 0,
+            db_primary_cached_count:  record_cached_query ? 1 : 0,
+            db_primary_count:  record_query ? 1 : 0,
+            db_primary_duration_s:  record_query ? 0.002 : 0,
+            db_replica_cached_count:  0,
+            db_replica_count:  0,
+            db_replica_duration_s:  0.0
+          )
+        elsif db_role == :replica
+          expect(described_class.db_counter_payload).to eq(
+            db_count: record_query ? 1 : 0,
+            db_write_count: record_write_query ? 1 : 0,
+            db_cached_count: record_cached_query ? 1 : 0,
+            db_primary_cached_count:  0,
+            db_primary_count:  0,
+            db_primary_duration_s:  0.0,
+            db_replica_cached_count:  record_cached_query ? 1 : 0,
+            db_replica_count:  record_query ? 1 : 0,
+            db_replica_duration_s:  record_query ? 0.002 : 0
+          )
+        else
+          expect(described_class.db_counter_payload).to eq(
+            db_count: record_query ? 1 : 0,
+            db_write_count: record_write_query ? 1 : 0,
+            db_cached_count: record_cached_query ? 1 : 0
+          )
+        end
+      end
+    end
+  end
+end
+
+RSpec.shared_examples 'record ActiveRecord metrics' do |db_role|
+  it 'increments only db counters' do
+    if record_query
+      expect(transaction).to receive(:increment).with(:gitlab_transaction_db_count_total, 1)
+      expect(transaction).to receive(:increment).with("gitlab_transaction_db_#{db_role}_count_total".to_sym, 1) if db_role
+    else
+      expect(transaction).not_to receive(:increment).with(:gitlab_transaction_db_count_total, 1)
+      expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_#{db_role}_count_total".to_sym, 1) if db_role
+    end
+
+    if record_write_query
+      expect(transaction).to receive(:increment).with(:gitlab_transaction_db_write_count_total, 1)
+    else
+      expect(transaction).not_to receive(:increment).with(:gitlab_transaction_db_write_count_total, 1)
+    end
+
+    if record_cached_query
+      expect(transaction).to receive(:increment).with(:gitlab_transaction_db_cached_count_total, 1)
+      expect(transaction).to receive(:increment).with("gitlab_transaction_db_#{db_role}_cached_count_total".to_sym, 1) if db_role
+    else
+      expect(transaction).not_to receive(:increment).with(:gitlab_transaction_db_cached_count_total, 1)
+      expect(transaction).not_to receive(:increment).with("gitlab_transaction_db_#{db_role}_cached_count_total".to_sym, 1) if db_role
+    end
+
+    subscriber.sql(event)
+  end
+
+  it 'observes sql_duration metric' do
+    if record_query
+      expect(transaction).to receive(:observe).with(:gitlab_sql_duration_seconds, 0.002)
+      expect(transaction).to receive(:observe).with("gitlab_sql_#{db_role}_duration_seconds".to_sym, 0.002) if db_role
+    else
+      expect(transaction).not_to receive(:observe)
+    end
+
+    subscriber.sql(event)
+  end
+end