1 files changed, 48 insertions, 9 deletions

diff --git a/lib/gitlab/metrics/subscribers/active_record.rb b/lib/gitlab/metrics/subscribers/active_record.rb
index 3db3317e833..9f7884e1364 100644
--- a/lib/gitlab/metrics/subscribers/active_record.rb
+++ b/lib/gitlab/metrics/subscribers/active_record.rb
@@ -14,6 +14,14 @@ module Gitlab
         SQL_DURATION_BUCKET = [0.05, 0.1, 0.25].freeze
         TRANSACTION_DURATION_BUCKET = [0.1, 0.25, 1].freeze
 
+        DB_LOAD_BALANCING_COUNTERS = %i{
+          db_replica_count db_replica_cached_count db_replica_wal_count
+          db_primary_count db_primary_cached_count db_primary_wal_count
+        }.freeze
+        DB_LOAD_BALANCING_DURATIONS = %i{db_primary_duration_s db_replica_duration_s}.freeze
+
+        SQL_WAL_LOCATION_REGEX = /(pg_current_wal_insert_lsn\(\)::text|pg_last_wal_replay_lsn\(\)::text)/.freeze
+
         # This event is published from ActiveRecordBaseTransactionMetrics and
         # used to record a database transaction duration when calling
         # ActiveRecord::Base.transaction {} block.
@@ -39,23 +47,56 @@ module Gitlab
           observe(:gitlab_sql_duration_seconds, event) do
             buckets SQL_DURATION_BUCKET
           end
+
+          if ::Gitlab::Database::LoadBalancing.enable?
+            db_role = ::Gitlab::Database::LoadBalancing.db_role_for_connection(payload[:connection])
+            return if db_role.blank?
+
+            increment_db_role_counters(db_role, payload)
+            observe_db_role_duration(db_role, event)
+          end
         end
 
         def self.db_counter_payload
           return {} unless Gitlab::SafeRequestStore.active?
 
-          payload = {}
-          DB_COUNTERS.each do |counter|
-            payload[counter] = Gitlab::SafeRequestStore[counter].to_i
+          {}.tap do |payload|
+            DB_COUNTERS.each do |counter|
+              payload[counter] = Gitlab::SafeRequestStore[counter].to_i
+            end
+
+            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(3)
+              end
+            end
           end
-          payload
         end
 
-        def self.known_payload_keys
-          DB_COUNTERS
+        private
+
+        def wal_command?(payload)
+          payload[:sql].match(SQL_WAL_LOCATION_REGEX)
+        end
+
+        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)
+          increment("db_#{db_role}_wal_count".to_sym) if !cached_query?(payload) && wal_command?(payload)
         end
 
-        private
+        def observe_db_role_duration(db_role, event)
+          observe("gitlab_sql_#{db_role}_duration_seconds".to_sym, event) do
+            buckets ::Gitlab::Metrics::Subscribers::ActiveRecord::SQL_DURATION_BUCKET
+          end
+
+          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
+        end
 
         def ignored_query?(payload)
           payload[:name] == 'SCHEMA' || IGNORABLE_SQL.include?(payload[:sql])
@@ -86,5 +127,3 @@ module Gitlab
     end
   end
 end
-
-Gitlab::Metrics::Subscribers::ActiveRecord.prepend_mod_with('Gitlab::Metrics::Subscribers::ActiveRecord')