Allow async queries in transactional fixtures

This is possible since rails#50999.
Now in transactional fixtures, the other threads do checkout the
same connection as the same thread, it's just synchronized around
performing the query.

This allow to properly test async queries without having to disable
transactional fixtures.

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,13 @@
+*   Make Active Record asynchronous queries compatible with transactional fixtures.
+
+    Previously transactional fixtures would disable asynchronous queries, because transactional
+    fixtures impose all queries use the same connection.
+
+    Now asynchronous queries will use the connection pinned by transactional fixtures, and behave
+    much closer to production.
+
+    *Jean Boussier*
+
 *   Deserialize binary data before decrypting
 
     This ensures that we call `PG::Connection.unescape_bytea` on PostgreSQL before decryption.

activerecord/lib/active_record/asynchronous_queries_tracker.rb

@@ -1,29 +1,30 @@
 # frozen_string_literal: true
 
+require "concurrent/atomic/atomic_boolean"
+require "concurrent/atomic/read_write_lock"
+
 module ActiveRecord
   class AsynchronousQueriesTracker # :nodoc:
-    module NullSession # :nodoc:
-      class << self
-        def active?
-          true
-        end
-
-        def finalize
-        end
-      end
-    end
-
     class Session # :nodoc:
       def initialize
-        @active = true
+        @active = Concurrent::AtomicBoolean.new(true)
+        @lock = Concurrent::ReadWriteLock.new
       end
 
       def active?
-        @active
+        @active.true?
       end
 
-      def finalize
-        @active = false
+      def synchronize(&block)
+        @lock.with_read_lock(&block)
+      end
+
+      def finalize(wait = false)
+        @active.make_false
+        if wait
+          # Wait until all thread with a read lock are done
+          @lock.with_write_lock { }
+        end
       end
     end
 
@@ -33,28 +34,31 @@ def install_executor_hooks(executor = ActiveSupport::Executor)
       end
 
       def run
-        ActiveRecord::Base.asynchronous_queries_tracker.start_session
+        ActiveRecord::Base.asynchronous_queries_tracker.tap(&:start_session)
       end
 
       def complete(asynchronous_queries_tracker)
         asynchronous_queries_tracker.finalize_session
       end
     end
 
-    attr_reader :current_session
-
     def initialize
-      @current_session = NullSession
+      @stack = []
+    end
+
+    def current_session
+      @stack.last or raise ActiveRecordError, "Can't perform asynchronous queries without a query session"
     end
 
     def start_session
-      @current_session = Session.new
-      self
+      session = Session.new
+      @stack << session
     end
 
-    def finalize_session
-      @current_session.finalize
-      @current_session = NullSession
+    def finalize_session(wait = false)
+      session = @stack.pop
+      session&.finalize(wait)
+      self
     end
   end
 end

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

@@ -683,7 +683,7 @@ def select(sql, name = nil, binds = [], prepare: false, async: false, allow_retr
               binds,
               prepare: prepare,
             )
-            if supports_concurrent_connections? && current_transaction.closed?
+            if supports_concurrent_connections? && !current_transaction.joinable?
               future_result.schedule!(ActiveRecord::Base.asynchronous_queries_session)
             else
               future_result.execute!(self)

activerecord/lib/active_record/future_result.rb

@@ -100,17 +100,21 @@ def cancel
     def execute_or_skip
       return unless pending?
 
-      @pool.with_connection do |connection|
-        return unless @mutex.try_lock
-        begin
-          if pending?
-            @event_buffer = EventBuffer.new(self, @instrumenter)
-            connection.with_instrumenter(@event_buffer) do
-              execute_query(connection, async: true)
+      @session.synchronize do
+        return unless pending?
+
+        @pool.with_connection do |connection|
+          return unless @mutex.try_lock
+          begin
+            if pending?
+              @event_buffer = EventBuffer.new(self, @instrumenter)
+              connection.with_instrumenter(@event_buffer) do
+                execute_query(connection, async: true)
+              end
             end
+          ensure
+            @mutex.unlock
           end
-        ensure
-          @mutex.unlock
         end
       end
     end

activerecord/lib/active_record/relation.rb

@@ -1124,9 +1124,6 @@ def delete_by(*args)
     # for queries to actually be executed concurrently. Otherwise it defaults to
     # executing them in the foreground.
     #
-    # +load_async+ will also fall back to executing in the foreground in the test environment when transactional
-    # fixtures are enabled.
-    #
     # If the query was actually executed in the background, the Active Record logs will show
     # it by prefixing the log line with <tt>ASYNC</tt>:
     #
@@ -1136,7 +1133,7 @@ def load_async
         return load if !c.async_enabled?
 
         unless loaded?
-          result = exec_main_query(async: c.current_transaction.closed?)
+          result = exec_main_query(async: !c.current_transaction.joinable?)
 
           if result.is_a?(Array)
             @records = result

activerecord/lib/active_record/test_fixtures.rb

@@ -137,12 +137,15 @@ def setup_fixtures(config = ActiveRecord::Base)
           invalidate_already_loaded_fixtures
           @loaded_fixtures = load_fixtures(config)
         end
+        setup_asynchronous_queries_session
 
         # Instantiate fixtures for every test if requested.
         instantiate_fixtures if use_instantiated_fixtures
       end
 
       def teardown_fixtures
+        teardown_asynchronous_queries_session
+
         # Rollback changes if a transaction is active.
         if run_in_transaction?
           teardown_transactional_fixtures
@@ -154,6 +157,14 @@ def teardown_fixtures
         ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
       end
 
+      def setup_asynchronous_queries_session
+        @_async_queries_session = ActiveRecord::Base.asynchronous_queries_tracker.start_session
+      end
+
+      def teardown_asynchronous_queries_session
+        ActiveRecord::Base.asynchronous_queries_tracker.finalize_session(true) if @_async_queries_session
+      end
+
       def invalidate_already_loaded_fixtures
         @@already_loaded_fixtures.clear
       end
@@ -190,6 +201,7 @@ def setup_transactional_fixtures
 
       def teardown_transactional_fixtures
         ActiveSupport::Notifications.unsubscribe(@connection_subscriber) if @connection_subscriber
+
         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.

activerecord/test/cases/asynchronous_queries_test.rb

@@ -6,8 +6,6 @@
 
 module AsynchronousQueriesSharedTests
   def test_async_select_failure
-    ActiveRecord::Base.asynchronous_queries_tracker.start_session
-
     if in_memory_db?
       assert_raises ActiveRecord::StatementInvalid do
         @connection.select_all "SELECT * FROM does_not_exists", async: true
@@ -19,13 +17,9 @@ def test_async_select_failure
         future_result.result
       end
     end
-  ensure
-    ActiveRecord::Base.asynchronous_queries_tracker.finalize_session
   end
 
   def test_async_query_from_transaction
-    ActiveRecord::Base.asynchronous_queries_tracker.start_session
-
     assert_nothing_raised do
       @connection.select_all "SELECT * FROM posts", async: true
     end
@@ -37,20 +31,15 @@ def test_async_query_from_transaction
         end
       end
     end
-  ensure
-    ActiveRecord::Base.asynchronous_queries_tracker.finalize_session
   end
 
   def test_async_query_cache
-    ActiveRecord::Base.asynchronous_queries_tracker.start_session
-
     @connection.enable_query_cache!
 
     @connection.select_all "SELECT * FROM posts"
     result = @connection.select_all "SELECT * FROM posts", async: true
     assert_equal ActiveRecord::FutureResult::Complete, result.class
   ensure
-    ActiveRecord::Base.asynchronous_queries_tracker.finalize_session
     @connection.disable_query_cache!
   end
 
@@ -103,7 +92,6 @@ def setup
   end
 
   def test_async_select_all
-    ActiveRecord::Base.asynchronous_queries_tracker.start_session
     status = {}
 
     subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |event|
@@ -125,7 +113,6 @@ def test_async_select_all
     assert_kind_of ActiveRecord::Result, future_result.result
     assert_equal @connection.supports_concurrent_connections?, status[:async]
   ensure
-    ActiveRecord::Base.asynchronous_queries_tracker.finalize_session
     ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
   end
 end

activerecord/test/cases/relation/load_async_test.rb

@@ -10,8 +10,6 @@ module ActiveRecord
   class LoadAsyncTest < ActiveRecord::TestCase
     include WaitForAsyncTestHelper
 
-    self.use_transactional_tests = false
-
     fixtures :posts, :comments, :categories, :categories_posts
 
     def test_scheduled?
@@ -240,8 +238,6 @@ def test_load_async_count_with_query_cache
 
   class LoadAsyncNullExecutorTest < ActiveRecord::TestCase
     unless in_memory_db?
-      self.use_transactional_tests = false
-
       fixtures :posts, :comments
 
       def setup
@@ -364,8 +360,6 @@ class LoadAsyncMultiThreadPoolExecutorTest < ActiveRecord::TestCase
     unless in_memory_db?
       include WaitForAsyncTestHelper
 
-      self.use_transactional_tests = false
-
       fixtures :posts, :comments
 
       def setup
@@ -505,8 +499,6 @@ class LoadAsyncMixedThreadPoolExecutorTest < ActiveRecord::TestCase
     unless in_memory_db?
       include WaitForAsyncTestHelper
 
-      self.use_transactional_tests = false
-
       fixtures :posts, :comments, :other_dogs
 
       def setup