Merge pull request rails#51091 from Shopify/pool-schema-cache

Refactor SchemaCache to hold a ConnectionPool
juliusdelta · Feb 15, 2024 · ed76d0f · ed76d0f
2 parents 024f2de + 1af13c6
commit ed76d0f
Show file tree

Hide file tree

Showing 22 changed files with 189 additions and 147 deletions.
diff --git a/activerecord/lib/active_record/attribute_methods/primary_key.rb b/activerecord/lib/active_record/attribute_methods/primary_key.rb
@@ -109,7 +109,7 @@ def get_primary_key(base_name) # :nodoc:
             elsif base_name && primary_key_prefix_type == :table_name_with_underscore
               base_name.foreign_key
             elsif ActiveRecord::Base != self && table_exists?
-              connection.schema_cache.primary_keys(table_name)
+              schema_cache.primary_keys(table_name)
             else
               "id"
             end

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -36,6 +36,7 @@ def schema_reflection
         SchemaReflection.new(nil)
       end
 
+      def schema_cache; end
       def connection_class; end
       def checkin(_); end
       def remove(_); end
@@ -117,7 +118,7 @@ class ConnectionPool
       include ConnectionAdapters::AbstractPool
 
       attr_accessor :automatic_reconnect, :checkout_timeout
-      attr_reader :db_config, :size, :reaper, :pool_config, :async_executor, :role, :shard
+      attr_reader :db_config, :size, :reaper, :pool_config, :async_executor, :role, :shard, :schema_cache
 
       delegate :schema_reflection, :schema_reflection=, :server_version, to: :pool_config
 
@@ -166,6 +167,8 @@ def initialize(pool_config)
 
         @async_executor = build_async_executor
 
+        @schema_cache = BoundSchemaReflection.new(schema_reflection, self)
+
         @reaper = Reaper.new(self, db_config.reaping_frequency)
         @reaper.run
       end
@@ -588,6 +591,7 @@ def attempt_to_checkout_all_existing_connections(raise_on_acquisition_timeout =
             loop do
               synchronize do
                 return if collected_conns.size == @connections.size && @now_connecting == 0
+
                 remaining_timeout = timeout_time - Process.clock_gettime(Process::CLOCK_MONOTONIC)
                 remaining_timeout = 0 if remaining_timeout < 0
                 conn = checkout_for_exclusive_access(remaining_timeout)

diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -50,7 +50,9 @@ def pool=(value)
         @schema_cache = nil
         @pool = value
 
-        @pool.schema_reflection.load!(self) if ActiveRecord.lazily_load_schema_cache
+        if @pool && ActiveRecord.lazily_load_schema_cache
+          @pool.schema_reflection.load!(@pool)
+        end
       end
 
       set_callback :checkin, :after, :enable_lazy_transactions!
@@ -321,7 +323,7 @@ def shard
       end
 
       def schema_cache
-        @schema_cache ||= BoundSchemaReflection.new(@pool.schema_reflection, self)
+        @pool.schema_cache || (@schema_cache ||= BoundSchemaReflection.for_lone_connection(@pool.schema_reflection, self))
       end
 
       # this method must only be called while holding connection pool's mutex