Decouple transactional fixtures and active connections

Ref: rails#50793

Transactional fixtures are currently tightly coupled with the pool
active connection. It assumes calling `pool.connection` will memoize
the checked out connection and leverage that to start a transaction
on it and ensure all subsequent accesses will get the same connection.

To allow to remove checkout caching (or make it optional), we first
must decouple transactional fixtures to not rely on it.

The idea is to behave similarly, but store the connection in
the pool as a special "pinned" connection, and not as the regular
active connection.

This allows to always return the same pinned connection,
but without necessarily assigning it as the active connection.

Additionally, this pinning impact all threads and fibers, so
that all threads have a consistent view of the database state.

Author: byroot

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -162,34 +162,56 @@ def initialize(pool_config)
         @threads_blocking_new_connections = 0
 
         @available = ConnectionLeasingQueue.new self
-
-        @lock_thread = false
+        @pinned_connection = nil
 
         @async_executor = build_async_executor
 
         @reaper = Reaper.new(self, db_config.reaping_frequency)
         @reaper.run
       end
 
-      def lock_thread=(lock_thread)
-        if lock_thread
-          @lock_thread = ActiveSupport::IsolatedExecutionState.context
-        else
-          @lock_thread = nil
-        end
-
-        if (active_connection = @thread_cached_conns[connection_cache_key(current_thread)])
-          active_connection.lock_thread = @lock_thread
-        end
-      end
-
       # Retrieve the connection associated with the current thread, or call
       # #checkout to obtain one if necessary.
       #
       # #connection can be called any number of times; the connection is
       # held in a cache keyed by a thread.
       def connection
-        @thread_cached_conns[connection_cache_key(current_thread)] ||= checkout
+        @thread_cached_conns[connection_cache_key(ActiveSupport::IsolatedExecutionState.context)] ||= checkout
+      end
+
+      def pin_connection!(lock_thread) # :nodoc:
+        raise "There is already a pinned connection" if @pinned_connection
+
+        @pinned_connection = (@thread_cached_conns[connection_cache_key(ActiveSupport::IsolatedExecutionState.context)] || checkout)
+        # Any leased connection must be in @connections otherwise
+        # some methods like #connected? won't behave correctly
+        unless @connections.include?(@pinned_connection)
+          @connections << @pinned_connection
+        end
+
+        @pinned_connection.lock_thread = ActiveSupport::IsolatedExecutionState.context if lock_thread
+        @pinned_connection.verify! # eagerly validate the connection
+        @pinned_connection.begin_transaction joinable: false, _lazy: false
+      end
+
+      def unpin_connection! # :nodoc:
+        raise "There isn't a pinned connection #{object_id}" unless @pinned_connection
+
+        clean = true
+        @pinned_connection.lock.synchronize do
+          connection, @pinned_connection = @pinned_connection, nil
+          if connection.transaction_open?
+            connection.rollback_transaction
+          else
+            # Something committed or rolled back the transaction
+            clean = false
+            connection.reset!
+          end
+          connection.lock_thread = nil
+          checkin(connection)
+        end
+
+        clean
       end
 
       def connection_class # :nodoc:
@@ -204,7 +226,7 @@ def connection_class # :nodoc:
       # #connection or #with_connection methods. Connections obtained through
       # #checkout will not be detected by #active_connection?
       def active_connection?
-        @thread_cached_conns[connection_cache_key(current_thread)]
+        @thread_cached_conns[connection_cache_key(ActiveSupport::IsolatedExecutionState.context)]
       end
 
       # Signal that the thread is finished with the current connection.
@@ -276,6 +298,7 @@ def disconnect(raise_on_acquisition_timeout = true)
               conn.disconnect!
             end
             @connections = []
+            @thread_cached_conns.clear
             @available.clear
           end
         end
@@ -360,9 +383,19 @@ def clear_reloadable_connections!
       # Raises:
       # - ActiveRecord::ConnectionTimeoutError no connection can be obtained from the pool.
       def checkout(checkout_timeout = @checkout_timeout)
-        connection = checkout_and_verify(acquire_connection(checkout_timeout))
-        connection.lock_thread = @lock_thread
-        connection
+        if @pinned_connection
+          synchronize do
+            @pinned_connection.verify!
+            # Any leased connection must be in @connections otherwise
+            # some methods like #connected? won't behave correctly
+            unless @connections.include?(@pinned_connection)
+              @connections << @pinned_connection
+            end
+          end
+          @pinned_connection
+        else
+          checkout_and_verify(acquire_connection(checkout_timeout))
+        end
       end
 
       # Check-in a database connection back into the pool, indicating that you
@@ -371,6 +404,8 @@ def checkout(checkout_timeout = @checkout_timeout)
       # +conn+: an AbstractAdapter object, which was obtained by earlier by
       # calling #checkout on this pool.
       def checkin(conn)
+        return if @pinned_connection.equal?(conn)
+
         conn.lock.synchronize do
           synchronize do
             remove_connection_from_thread_cache conn
@@ -379,7 +414,6 @@ def checkin(conn)
               conn.expire
             end
 
-            conn.lock_thread = nil
             @available.add conn
           end
         end
@@ -533,10 +567,6 @@ def connection_cache_key(thread)
           thread
         end
 
-        def current_thread
-          @lock_thread || ActiveSupport::IsolatedExecutionState.context
-        end
-
         # Take control of all existing connections so a "group" action such as
         # reload/disconnect can be performed safely. It is no longer enough to
         # wrap it in +synchronize+ because some pool's actions are allowed

activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb

@@ -36,17 +36,17 @@ def initialize(*)
         end
 
         def enable_query_cache!
-          @query_cache_enabled[connection_cache_key(current_thread)] = true
+          @query_cache_enabled[connection_cache_key(ActiveSupport::IsolatedExecutionState.context)] = true
           connection.enable_query_cache! if active_connection?
         end
 
         def disable_query_cache!
-          @query_cache_enabled.delete connection_cache_key(current_thread)
+          @query_cache_enabled.delete connection_cache_key(ActiveSupport::IsolatedExecutionState.context)
           connection.disable_query_cache! if active_connection?
         end
 
         def query_cache_enabled
-          @query_cache_enabled[connection_cache_key(current_thread)]
+          @query_cache_enabled[connection_cache_key(ActiveSupport::IsolatedExecutionState.context)]
         end
       end
 

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -174,19 +174,13 @@ def initialize(config_or_deprecated_connection, deprecated_logger = nil, depreca
         @verified = false
       end
 
-      THREAD_LOCK = ActiveSupport::Concurrency::ThreadLoadInterlockAwareMonitor.new
-      private_constant :THREAD_LOCK
-
-      FIBER_LOCK = ActiveSupport::Concurrency::LoadInterlockAwareMonitor.new
-      private_constant :FIBER_LOCK
-
       def lock_thread=(lock_thread) # :nodoc:
         @lock =
         case lock_thread
         when Thread
-          THREAD_LOCK
+          ActiveSupport::Concurrency::ThreadLoadInterlockAwareMonitor.new
         when Fiber
-          FIBER_LOCK
+          ActiveSupport::Concurrency::LoadInterlockAwareMonitor.new
         else
           ActiveSupport::Concurrency::NullLock
         end

activerecord/lib/active_record/test_fixtures.rb

@@ -131,8 +131,8 @@ def setup_fixtures(config = ActiveRecord::Base)
         end
 
         @fixture_cache = {}
-        @fixture_connections = {}
         @fixture_cache_key = [self.class.fixture_table_names.dup, self.class.fixture_paths.dup, self.class.fixture_class_names.dup]
+        @fixture_connection_pools = []
         @@already_loaded_fixtures ||= {}
         @connection_subscriber = nil
         @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
@@ -163,20 +163,24 @@ def teardown_fixtures
           teardown_transactional_fixtures
         else
           ActiveRecord::FixtureSet.reset_cache
+          invalidate_already_loaded_fixtures
         end
 
         ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
       end
 
+      def invalidate_already_loaded_fixtures
+        @@already_loaded_fixtures.clear
+      end
+
       def setup_transactional_fixtures
         setup_shared_connection_pool
 
         # Begin transactions for connections already established
-        @fixture_connections = ActiveRecord::Base.connection_handler.connection_pool_list(:writing).to_h do |pool|
-          pool.lock_thread = true if lock_threads
-          connection = pool.connection
-          transaction = connection.begin_transaction(joinable: false, _lazy: false)
-          [connection, transaction]
+        @fixture_connection_pools = ActiveRecord::Base.connection_handler.connection_pool_list(:writing)
+        @fixture_connection_pools.each do |pool|
+          pool.pin_connection!(lock_threads)
+          pool.connection
         end
 
         # When connections are established in the future, begin a transaction too
@@ -185,19 +189,14 @@ def setup_transactional_fixtures
           shard = payload[:shard] if payload.key?(:shard)
 
           if connection_name
-            begin
-              connection = ActiveRecord::Base.connection_handler.retrieve_connection(connection_name, shard: shard)
-              connection.connect! # eagerly validate the connection
-            rescue ConnectionNotEstablished
-              connection = nil
-            end
-
-            if connection
+            pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_name, shard: shard)
+            if pool
               setup_shared_connection_pool
 
-              if !@fixture_connections.key?(connection)
-                connection.pool.lock_thread = true if lock_threads
-                @fixture_connections[connection] = connection.begin_transaction(joinable: false, _lazy: false)
+              unless @fixture_connection_pools.include?(pool)
+                pool.pin_connection!(lock_threads)
+                pool.connection
+                @fixture_connection_pools << pool
               end
             end
           end
@@ -206,27 +205,15 @@ def setup_transactional_fixtures
 
       def teardown_transactional_fixtures
         ActiveSupport::Notifications.unsubscribe(@connection_subscriber) if @connection_subscriber
-        @fixture_connections.each do |connection, transaction|
-          begin
-            connection.rollback_transaction(transaction)
-          rescue ActiveRecord::StatementInvalid
-            # Something commited or rolled back the transaction.
-            # We can no longer trust the database state is clean.
-            invalidate_already_loaded_fixtures
-            # We also don't know for sure the connection wasn't
-            # mutated in dangerous ways.
-            connection.disconnect!
-          end
-          connection.pool.lock_thread = false
+        unless @fixture_connection_pools.map(&:unpin_connection!).all?
+          # Something caused the transaction to be committed or rolled back
+          # We can no longer trust the database is in a clean state.
+          @@already_loaded_fixtures.clear
         end
-        @fixture_connections.clear
+        @fixture_connection_pools.clear
         teardown_shared_connection_pool
       end
 
-      def invalidate_already_loaded_fixtures
-        @@already_loaded_fixtures.clear
-      end
-
       # Shares the writing connection pool with connections on
       # other handlers.
       #

activerecord/test/cases/adapters/abstract_mysql_adapter/connection_test.rb

@@ -37,7 +37,7 @@ def test_no_automatic_reconnection_after_timeout
     assert_not_predicate @connection, :active?
   ensure
     # Repair all fixture connections so other tests won't break.
-    @fixture_connections.each_key(&:verify!)
+    @fixture_connection_pools.each { |p| p.connection.verify! }
   end
 
   def test_successful_reconnection_after_timeout_with_manual_reconnect

activerecord/test/cases/adapters/postgresql/connection_test.rb

@@ -145,7 +145,7 @@ def test_reconnection_after_actual_disconnection_with_verify
       assert_predicate @connection, :active?
     ensure
       # Repair all fixture connections so other tests won't break.
-      @fixture_connections.each_key(&:verify!)
+      @fixture_connection_pools.each { |p| p.connection.verify! }
     end
 
     def test_set_session_variable_true

activerecord/test/cases/connection_adapters/connection_handler_test.rb

@@ -6,8 +6,6 @@
 module ActiveRecord
   module ConnectionAdapters
     class ConnectionHandlerTest < ActiveRecord::TestCase
-      self.use_transactional_tests = false
-
       fixtures :people
 
       def setup
@@ -95,8 +93,6 @@ def test_fixtures_dont_raise_if_theres_no_writing_pool_config
           connection_handler = ActiveRecord::Base.connection_handler
           ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
 
-          setup_transactional_fixtures
-
           assert_nothing_raised do
             ActiveRecord::Base.connects_to(database: { reading: :arunit, writing: :arunit })
           end
@@ -105,8 +101,6 @@ def test_fixtures_dont_raise_if_theres_no_writing_pool_config
           ro_conn = ActiveRecord::Base.connection_handler.retrieve_connection("ActiveRecord::Base", role: :reading)
 
           assert_equal rw_conn, ro_conn
-
-          teardown_transactional_fixtures
         ensure
           ActiveRecord::Base.connection_handler = connection_handler
         end