Merge pull request rails#43674 from Shopify/activerecord-isolated-state

Use `IsolatedExecutionState` across Active Record

Author: byroot

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

@@ -81,11 +81,11 @@ def initialize
       end
 
       def prevent_writes # :nodoc:
-        Thread.current[:prevent_writes]
+        ActiveSupport::IsolatedExecutionState[:active_record_prevent_writes]
       end
 
       def prevent_writes=(prevent_writes) # :nodoc:
-        Thread.current[:prevent_writes] = prevent_writes
+        ActiveSupport::IsolatedExecutionState[:active_record_prevent_writes] = prevent_writes
       end
 
       # Prevent writing to the database regardless of role.

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -178,7 +178,7 @@ def prepared_statements?
       alias :prepared_statements :prepared_statements?
 
       def prepared_statements_disabled_cache # :nodoc:
-        Thread.current[:ar_prepared_statements_disabled_cache] ||= Set.new
+        ActiveSupport::IsolatedExecutionState[:active_record_prepared_statements_disabled_cache] ||= Set.new
       end
 
       class Version

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

@@ -78,7 +78,7 @@ def begin_isolated_db_transaction(isolation) # :nodoc:
           raise TransactionIsolationError, "SQLite3 only supports the `read_uncommitted` transaction isolation level" if isolation != :read_uncommitted
           raise StandardError, "You need to enable the shared-cache mode in SQLite mode before attempting to change the transaction isolation level" unless shared_cache?
 
-          Thread.current.thread_variable_set("read_uncommitted", @connection.get_first_value("PRAGMA read_uncommitted"))
+          ActiveSupport::IsolatedExecutionState[:active_record_read_uncommitted] = @connection.get_first_value("PRAGMA read_uncommitted")
           @connection.read_uncommitted = true
           begin_db_transaction
         end
@@ -108,7 +108,7 @@ def high_precision_current_timestamp
 
         private
           def reset_read_uncommitted
-            read_uncommitted = Thread.current.thread_variable_get("read_uncommitted")
+            read_uncommitted = ActiveSupport::IsolatedExecutionState[:active_record_read_uncommitted]
             return unless read_uncommitted
 
             @connection.read_uncommitted = read_uncommitted

activerecord/lib/active_record/connection_handling.rb

@@ -212,18 +212,16 @@ def connecting_to(role: default_role, shard: default_shard, prevent_writes: fals
     # is useful in cases you're using sharding to provide per-request
     # database isolation.
     def prohibit_shard_swapping(enabled = true)
-      prev_value = Thread.current.thread_variable_get(:prohibit_shard_swapping)
-
-      Thread.current.thread_variable_set(:prohibit_shard_swapping, enabled)
-
+      prev_value = ActiveSupport::IsolatedExecutionState[:active_record_prohibit_shard_swapping]
+      ActiveSupport::IsolatedExecutionState[:active_record_prohibit_shard_swapping] = enabled
       yield
     ensure
-      Thread.current.thread_variable_set(:prohibit_shard_swapping, prev_value)
+      ActiveSupport::IsolatedExecutionState[:active_record_prohibit_shard_swapping] = prev_value
     end
 
     # Determine whether or not shard swapping is currently prohibited
     def shard_swapping_prohibited?
-      Thread.current.thread_variable_get(:prohibit_shard_swapping)
+      ActiveSupport::IsolatedExecutionState[:active_record_prohibit_shard_swapping]
     end
 
     # Prevent writing to the database regardless of role.

activerecord/lib/active_record/core.rb

@@ -92,11 +92,11 @@ def self.application_record_class? # :nodoc:
       self.filter_attributes = []
 
       def self.connection_handler
-        Thread.current.thread_variable_get(:ar_connection_handler) || default_connection_handler
+        ActiveSupport::IsolatedExecutionState[:active_record_connection_handler] || default_connection_handler
       end
 
       def self.connection_handler=(handler)
-        Thread.current.thread_variable_set(:ar_connection_handler, handler)
+        ActiveSupport::IsolatedExecutionState[:active_record_connection_handler] = handler
       end
 
       def self.connection_handlers
@@ -131,8 +131,8 @@ def self.asynchronous_queries_session # :nodoc:
       end
 
       def self.asynchronous_queries_tracker # :nodoc:
-        Thread.current.thread_variable_get(:ar_asynchronous_queries_tracker) ||
-          Thread.current.thread_variable_set(:ar_asynchronous_queries_tracker, AsynchronousQueriesTracker.new)
+        ActiveSupport::IsolatedExecutionState[:active_record_asynchronous_queries_tracker] ||= \
+          AsynchronousQueriesTracker.new
       end
 
       # Returns the symbol representing the current connected role.
@@ -199,11 +199,11 @@ def self.current_preventing_writes
       end
 
       def self.connected_to_stack # :nodoc:
-        if connected_to_stack = Thread.current.thread_variable_get(:ar_connected_to_stack)
+        if connected_to_stack = ActiveSupport::IsolatedExecutionState[:active_record_connected_to_stack]
           connected_to_stack
         else
           connected_to_stack = Concurrent::Array.new
-          Thread.current.thread_variable_set(:ar_connected_to_stack, connected_to_stack)
+          ActiveSupport::IsolatedExecutionState[:active_record_connected_to_stack] = connected_to_stack
           connected_to_stack
         end
       end

activerecord/lib/active_record/no_touching.rb

@@ -39,7 +39,7 @@ def applied_to?(klass) # :nodoc:
 
       private
         def klasses
-          Thread.current[:no_touching_classes] ||= []
+          ActiveSupport::IsolatedExecutionState[:active_record_no_touching_classes] ||= []
         end
     end