Merge pull request rails#52806 from Shopify/transactional-async-queries

Allow async queries in 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