Fix active record instrumentation not thread safe:

- When using Active Record async feature, the instrumentation was
  not thread safe.

  ```ruby
  Post.count_async

  ActiveSupport::Notifications.subscribed(->(*) { }, "active_record.sql") do
    Post.count
  end
  ```

  With the right race condition, the subscription would never fire up.
  This is due to temporarly overriding the instrumentation instance
  variable on the connection.

Author: Edouard-chin

activerecord/CHANGELOG.md

@@ -1,3 +1,18 @@
+*   Fix SQL notifications sometimes not sent when using async queries.
+
+    ```ruby
+    Post.async_count
+    ActiveSupport::Notifications.subscribed(->(*) { "Will never reach here" }) do
+      Post.count
+    end
+    ```
+
+    In rare circumstances and under the right race condition, Active Support notifications
+    would no longer be dispatched after using an asynchronous query.
+    This is now fixed.
+
+    *Edouard Chin*
+
 *   Eliminate queries loading dumped schema cache on Postgres
 
     Improve resiliency by avoiding needing to open a database connection to load the

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -150,7 +150,6 @@ def initialize(config_or_deprecated_connection, deprecated_logger = nil, depreca
         end
 
         @owner = nil
-        @instrumenter = ActiveSupport::Notifications.instrumenter
         @pool = ActiveRecord::ConnectionAdapters::NullPool.new
         @idle_since = Process.clock_gettime(Process::CLOCK_MONOTONIC)
         @visitor = arel_visitor
@@ -191,19 +190,6 @@ def lock_thread=(lock_thread) # :nodoc:
         end
       end
 
-      EXCEPTION_NEVER = { Exception => :never }.freeze # :nodoc:
-      EXCEPTION_IMMEDIATE = { Exception => :immediate }.freeze # :nodoc:
-      private_constant :EXCEPTION_NEVER, :EXCEPTION_IMMEDIATE
-      def with_instrumenter(instrumenter, &block) # :nodoc:
-        Thread.handle_interrupt(EXCEPTION_NEVER) do
-          previous_instrumenter = @instrumenter
-          @instrumenter = instrumenter
-          Thread.handle_interrupt(EXCEPTION_IMMEDIATE, &block)
-        ensure
-          @instrumenter = previous_instrumenter
-        end
-      end
-
       def check_if_write_query(sql) # :nodoc:
         if preventing_writes? && write_query?(sql)
           raise ActiveRecord::ReadOnlyError, "Write query attempted while in readonly mode: #{sql}"
@@ -1146,7 +1132,7 @@ def translate_exception_class(native_error, sql, binds)
         end
 
         def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &block) # :doc:
-          @instrumenter.instrument(
+          instrumenter.instrument(
             "sql.active_record",
             sql:               sql,
             name:              name,
@@ -1163,6 +1149,10 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &bl
           raise ex.set_query(sql, binds)
         end
 
+        def instrumenter # :nodoc:
+          ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] ||= ActiveSupport::Notifications.instrumenter
+        end
+
         def translate_exception(exception, message:, sql:, binds:)
           # override in derived class
           case exception

activerecord/lib/active_record/future_result.rb

@@ -108,9 +108,9 @@ def execute_or_skip
           begin
             if pending?
               @event_buffer = EventBuffer.new(self, @instrumenter)
-              connection.with_instrumenter(@event_buffer) do
-                execute_query(connection, async: true)
-              end
+              ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] = @event_buffer
+
+              execute_query(connection, async: true)
             end
           ensure
             @mutex.unlock

activerecord/test/cases/relation/load_async_test.rb

@@ -5,6 +5,7 @@
 require "models/category"
 require "models/comment"
 require "models/other_dog"
+require "concurrent/atomic/count_down_latch"
 
 module ActiveRecord
   class LoadAsyncTest < ActiveRecord::TestCase
@@ -157,6 +158,42 @@ def test_load_async_from_transaction
       assert_equal ["In Transaction"], posts.map(&:title).uniq
     end
 
+    def test_load_async_instrumentation_is_thread_safe
+      skip unless ActiveRecord::Base.connection.async_enabled?
+
+      begin
+        latch1 = Concurrent::CountDownLatch.new
+        latch2 = Concurrent::CountDownLatch.new
+
+        old_log = ActiveRecord::Base.connection.method(:log)
+        ActiveRecord::Base.connection.singleton_class.undef_method(:log)
+
+        ActiveRecord::Base.connection.singleton_class.define_method(:log) do |*args, **kwargs, &block|
+          unless kwargs[:async]
+            return old_log.call(*args, **kwargs, &block)
+          end
+
+          latch1.count_down
+          latch2.wait
+          old_log.call(*args, **kwargs, &block)
+        end
+
+        Post.async_count
+        latch1.wait
+
+        notification_called = false
+        ActiveSupport::Notifications.subscribed(->(*) { notification_called = true }, "sql.active_record") do
+          Post.count
+        end
+
+        assert(notification_called)
+      ensure
+        latch2.count_down
+        ActiveRecord::Base.connection.singleton_class.undef_method(:log)
+        ActiveRecord::Base.connection.singleton_class.define_method(:log, old_log)
+      end
+    end
+
     def test_eager_loading_query
       expected_records = Post.where(author_id: 1).eager_load(:comments).to_a