Merge pull request rails#50999 from Shopify/refactor-transactional-fixtures

Decouple transactional fixtures and active connections

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