Cache PostgresqlAdapter type map

This executes a query to the DB every time a connection is opened.
Currently the query returns 800 rows so caching this can improve
performance.

config/initializers/active_record_ping.rb → config/initializers/active_record_postgresql_adapter.rb

 # frozen_string_literal: true
 
-# # frozen_string_literal: true
-
 if Gitlab::Utils.to_boolean(ENV['ENABLE_ACTIVERECORD_EMPTY_PING'], default: true)
   ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Gitlab::Database::PostgresqlAdapter::EmptyQueryPing)
 end
+
+if Gitlab::Utils.to_boolean(ENV['ENABLE_ACTIVERECORD_TYPEMAP_CACHE'], default: false)
+  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Gitlab::Database::PostgresqlAdapter::TypeMapCache)
+end

lib/gitlab/database/postgresql_adapter/type_map_cache.rb

+# frozen_string_literal: true
+
+# Caches loading of additional types from the DB
+# https://github.com/rails/rails/blob/v6.0.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L521-L589
+
+# rubocop:disable Gitlab/ModuleWithInstanceVariables
+
+module Gitlab
+  module Database
+    module PostgresqlAdapter
+      module TypeMapCache
+        extend ActiveSupport::Concern
+
+        TYPE_MAP_CACHE_MONITOR = ::Monitor.new
+
+        class_methods do
+          def type_map_cache
+            TYPE_MAP_CACHE_MONITOR.synchronize do
+              @type_map_cache ||= {}
+            end
+          end
+        end
+
+        def initialize_type_map(map = type_map)
+          TYPE_MAP_CACHE_MONITOR.synchronize do
+            cached_type_map = self.class.type_map_cache[@connection_parameters.hash]
+            break @type_map = cached_type_map if cached_type_map
+
+            super
+            self.class.type_map_cache[@connection_parameters.hash] = map
+          end
+        end
+
+        def reload_type_map
+          TYPE_MAP_CACHE_MONITOR.synchronize do
+            self.class.type_map_cache[@connection_parameters.hash] = nil
+          end
+
+          super
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::PostgresqlAdapter::TypeMapCache do
+  let(:db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'test', name: 'primary').configuration_hash }
+  let(:adapter_class) do
+    Class.new(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+  end
+
+  describe '#initialize_type_map' do
+    it 'caches loading of types in memory' do
+      initialize_connection.disconnect!
+      recorder_without_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! }
+
+      expect(recorder_without_cache.log).to include(a_string_matching(/FROM pg_type/)).twice
+
+      adapter_class.prepend(described_class)
+
+      initialize_connection.disconnect!
+      recorder_with_cache = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection.disconnect! }
+
+      expect(recorder_with_cache.count).to be < recorder_without_cache.count
+
+      # There's still one pg_type query left here because `#add_pg_decoders` executes another pg_type query
+      # in https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L912.
+      # This query is much cheaper because it only returns very few records.
+      expect(recorder_with_cache.log).to include(a_string_matching(/FROM pg_type/)).once
+    end
+
+    it 'only reuses the cache if the connection parameters are exactly the same' do
+      adapter_class.prepend(described_class)
+
+      initialize_connection.disconnect!
+
+      other_config = db_config.dup
+      other_config[:connect_timeout] = db_config[:connect_timeout].to_i + 10
+
+      recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { initialize_connection(other_config).disconnect! }
+
+      expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).twice
+    end
+  end
+
+  describe '#reload_type_map' do
+    it 'clears the cache and executes the type map query again' do
+      adapter_class.prepend(described_class)
+
+      initialize_connection.disconnect!
+
+      connection = initialize_connection
+      recorder = ActiveRecord::QueryRecorder.new(skip_schema_queries: false) { connection.reload_type_map }
+
+      expect(recorder.log).to include(a_string_matching(/FROM pg_type/)).once
+    end
+  end
+
+  # Based on https://github.com/rails/rails/blob/v6.1.3.2/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L36-L41
+  def initialize_connection(config = db_config)
+    conn_params = config.symbolize_keys.compact
+
+    conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
+    conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]
+
+    valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:requiressl]
+    conn_params.slice!(*valid_conn_param_keys)
+
+    adapter_class.new(
+      adapter_class.new_client(conn_params),
+      ActiveRecord::Base.logger,
+      conn_params,
+      config
+    )
+  end
+end

spec/support/helpers/query_recorder.rb

@@ -2,15 +2,16 @@
 
 module ActiveRecord
   class QueryRecorder
-    attr_reader :log, :skip_cached, :cached, :data
+    attr_reader :log, :skip_cached, :skip_schema_queries, :cached, :data
 
     UNKNOWN = %w[unknown unknown].freeze
 
-    def initialize(skip_cached: true, log_file: nil, query_recorder_debug: false, &block)
+    def initialize(skip_cached: true, skip_schema_queries: true, log_file: nil, query_recorder_debug: false, &block)
       @data = Hash.new { |h, k| h[k] = { count: 0, occurrences: [], backtrace: [], durations: [] } }
       @log = []
       @cached = []
       @skip_cached = skip_cached
+      @skip_schema_queries = skip_schema_queries
       @query_recorder_debug = ENV['QUERY_RECORDER_DEBUG'] || query_recorder_debug
       @log_file = log_file
       record(&block) if block_given?
@@ -79,7 +80,7 @@ module ActiveRecord
 
       if values[:cached] && skip_cached
         @cached << values[:sql]
-      elsif !values[:name]&.include?("SCHEMA")
+      elsif !skip_schema_queries || !values[:name]&.include?("SCHEMA")
         backtrace = @query_recorder_debug ? show_backtrace(values, duration) : nil
         @log << values[:sql]
         store_sql_by_source(values: values, duration: duration, backtrace: backtrace)