Add latest changes from gitlab-org/gitlab@master

5 files changed, 76 insertions, 93 deletions

diff --git a/lib/gitlab/metrics/subscribers/active_record.rb b/lib/gitlab/metrics/subscribers/active_record.rb
index 24006d75c4d..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,20 +47,57 @@ 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
 
         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
+
+        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])
         end
@@ -82,5 +127,3 @@ module Gitlab
     end
   end
 end
-
-Gitlab::Metrics::Subscribers::ActiveRecord.prepend_mod_with('Gitlab::Metrics::Subscribers::ActiveRecord')
diff --git a/lib/gitlab/redis/cache.rb b/lib/gitlab/redis/cache.rb
index a634f12345a..cf4dbe48643 100644
--- a/lib/gitlab/redis/cache.rb
+++ b/lib/gitlab/redis/cache.rb
@@ -1,35 +1,13 @@
 # frozen_string_literal: true
 
-# please require all dependencies below:
-require_relative 'wrapper' unless defined?(::Rails) && ::Rails.root.present?
-
 module Gitlab
   module Redis
     class Cache < ::Gitlab::Redis::Wrapper
       CACHE_NAMESPACE = 'cache:gitlab'
-      DEFAULT_REDIS_CACHE_URL = 'redis://localhost:6380'
-      REDIS_CACHE_CONFIG_ENV_VAR_NAME = 'GITLAB_REDIS_CACHE_CONFIG_FILE'
 
       class << self
         def default_url
-          DEFAULT_REDIS_CACHE_URL
-        end
-
-        def config_file_name
-          # if ENV set for this class, use it even if it points to a file does not exist
-          file_name = ENV[REDIS_CACHE_CONFIG_ENV_VAR_NAME]
-          return file_name unless file_name.nil?
-
-          # otherwise, if config files exists for this class, use it
-          file_name = config_file_path('redis.cache.yml')
-          return file_name if File.file?(file_name)
-
-          # this will force use of DEFAULT_REDIS_QUEUES_URL when config file is absent
-          super
-        end
-
-        def instrumentation_class
-          ::Gitlab::Instrumentation::Redis::Cache
+          'redis://localhost:6380'
         end
       end
     end
diff --git a/lib/gitlab/redis/queues.rb b/lib/gitlab/redis/queues.rb
index 42d5167beb3..a0777510cd5 100644
--- a/lib/gitlab/redis/queues.rb
+++ b/lib/gitlab/redis/queues.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-# please require all dependencies below:
+# We need this require for MailRoom
 require_relative 'wrapper' unless defined?(::Gitlab::Redis::Wrapper)
 
 module Gitlab
@@ -8,29 +8,10 @@ module Gitlab
     class Queues < ::Gitlab::Redis::Wrapper
       SIDEKIQ_NAMESPACE = 'resque:gitlab'
       MAILROOM_NAMESPACE = 'mail_room:gitlab'
-      DEFAULT_REDIS_QUEUES_URL = 'redis://localhost:6381'
-      REDIS_QUEUES_CONFIG_ENV_VAR_NAME = 'GITLAB_REDIS_QUEUES_CONFIG_FILE'
 
       class << self
         def default_url
-          DEFAULT_REDIS_QUEUES_URL
-        end
-
-        def config_file_name
-          # if ENV set for this class, use it even if it points to a file does not exist
-          file_name = ENV[REDIS_QUEUES_CONFIG_ENV_VAR_NAME]
-          return file_name if file_name
-
-          # otherwise, if config files exists for this class, use it
-          file_name = config_file_path('redis.queues.yml')
-          return file_name if File.file?(file_name)
-
-          # this will force use of DEFAULT_REDIS_QUEUES_URL when config file is absent
-          super
-        end
-
-        def instrumentation_class
-          ::Gitlab::Instrumentation::Redis::Queues
+          'redis://localhost:6381'
         end
       end
     end
diff --git a/lib/gitlab/redis/shared_state.rb b/lib/gitlab/redis/shared_state.rb
index 2848c9f0b59..8bd831741f3 100644
--- a/lib/gitlab/redis/shared_state.rb
+++ b/lib/gitlab/redis/shared_state.rb
@@ -1,8 +1,5 @@
 # frozen_string_literal: true
 
-# please require all dependencies below:
-require_relative 'wrapper' unless defined?(::Gitlab::Redis::Wrapper)
-
 module Gitlab
   module Redis
     class SharedState < ::Gitlab::Redis::Wrapper
@@ -10,29 +7,10 @@ module Gitlab
       USER_SESSIONS_NAMESPACE = 'session:user:gitlab'
       USER_SESSIONS_LOOKUP_NAMESPACE = 'session:lookup:user:gitlab'
       IP_SESSIONS_LOOKUP_NAMESPACE = 'session:lookup:ip:gitlab2'
-      DEFAULT_REDIS_SHARED_STATE_URL = 'redis://localhost:6382'
-      REDIS_SHARED_STATE_CONFIG_ENV_VAR_NAME = 'GITLAB_REDIS_SHARED_STATE_CONFIG_FILE'
 
       class << self
         def default_url
-          DEFAULT_REDIS_SHARED_STATE_URL
-        end
-
-        def config_file_name
-          # if ENV set for this class, use it even if it points to a file does not exist
-          file_name = ENV[REDIS_SHARED_STATE_CONFIG_ENV_VAR_NAME]
-          return file_name if file_name
-
-          # otherwise, if config files exists for this class, use it
-          file_name = config_file_path('redis.shared_state.yml')
-          return file_name if File.file?(file_name)
-
-          # this will force use of DEFAULT_REDIS_SHARED_STATE_URL when config file is absent
-          super
-        end
-
-        def instrumentation_class
-          ::Gitlab::Instrumentation::Redis::SharedState
+          'redis://localhost:6382'
         end
       end
     end
diff --git a/lib/gitlab/redis/wrapper.rb b/lib/gitlab/redis/wrapper.rb
index 94ab67ef08a..ea0802ffbdc 100644
--- a/lib/gitlab/redis/wrapper.rb
+++ b/lib/gitlab/redis/wrapper.rb
@@ -1,16 +1,16 @@
 # frozen_string_literal: true
 
 # This file should only be used by sub-classes, not directly by any clients of the sub-classes
-# please require all dependencies below:
+
+# Explicitly load parts of ActiveSupport because MailRoom does not load
+# Rails.
 require 'active_support/core_ext/hash/keys'
 require 'active_support/core_ext/module/delegation'
+require 'active_support/core_ext/string/inflections'
 
 module Gitlab
   module Redis
     class Wrapper
-      DEFAULT_REDIS_URL = 'redis://localhost:6379'
-      REDIS_CONFIG_ENV_VAR_NAME = 'GITLAB_REDIS_CONFIG_FILE'
-
       class << self
         delegate :params, :url, to: :new
 
@@ -52,32 +52,35 @@ module Gitlab
         end
 
         def default_url
-          DEFAULT_REDIS_URL
+          raise NotImplementedError
         end
 
-        # Return the absolute path to a Rails configuration file
-        #
-        # We use this instead of `Rails.root` because for certain tasks
-        # utilizing these classes, `Rails` might not be available.
         def config_file_path(filename)
-          File.expand_path("../../../config/#{filename}", __dir__)
+          path = File.join(rails_root, 'config', filename)
+          return path if File.file?(path)
         end
 
-        def config_file_name
-          # if ENV set for wrapper class, use it even if it points to a file does not exist
-          file_name = ENV[REDIS_CONFIG_ENV_VAR_NAME]
-          return file_name unless file_name.nil?
+        # We need this local implementation of Rails.root because MailRoom
+        # doesn't load Rails.
+        def rails_root
+          File.expand_path('../../..', __dir__)
+        end
 
-          # otherwise, if config files exists for wrapper class, use it
-          file_name = config_file_path('resque.yml')
-          return file_name if File.file?(file_name)
+        def config_file_name
+          [
+            ENV["GITLAB_REDIS_#{store_name.underscore.upcase}_CONFIG_FILE"],
+            config_file_path("redis.#{store_name.underscore}.yml"),
+            ENV['GITLAB_REDIS_CONFIG_FILE'],
+            config_file_path('resque.yml')
+          ].compact.first
+        end
 
-          # nil will force use of DEFAULT_REDIS_URL when config file is absent
-          nil
+        def store_name
+          name.demodulize
         end
 
         def instrumentation_class
-          raise NotImplementedError
+          "::Gitlab::Instrumentation::Redis::#{store_name}".constantize
         end
       end