let EXPLAIN use a thread locals registry [John J. Wang & Xavier Noria]

Closes #10198.

3 files changed, 41 insertions, 9 deletions

diff --git a/activerecord/lib/active_record/explain.rb b/activerecord/lib/active_record/explain.rb
index 15736575a2..e65dab07ba 100644
--- a/activerecord/lib/active_record/explain.rb
+++ b/activerecord/lib/active_record/explain.rb
@@ -1,22 +1,22 @@
 require 'active_support/lazy_load_hooks'
+require 'active_record/explain_registry'
 
 module ActiveRecord
   module Explain
-    # Relation#explain needs to be able to collect the queries.
+    # Executes the block with the collect flag enabled. Queries are collected
+    # asynchronously by the subscriber and returned.
     def collecting_queries_for_explain # :nodoc:
-      current = Thread.current
-      original, current[:available_queries_for_explain] = current[:available_queries_for_explain], []
+      ExplainRegistry.collect = true
       yield
-      return current[:available_queries_for_explain]
+      ExplainRegistry.queries
     ensure
-      # Note that the return value above does not depend on this assignment.
-      current[:available_queries_for_explain] = original
+      ExplainRegistry.reset
     end
 
     # Makes the adapter execute EXPLAIN for the tuples of queries and bindings.
     # Returns a formatted string ready to be logged.
     def exec_explain(queries) # :nodoc:
-      str = queries && queries.map do |sql, bind|
+      str = queries.map do |sql, bind|
         [].tap do |msg|
           msg << "EXPLAIN for: #{sql}"
           unless bind.empty?
@@ -31,6 +31,7 @@ module ActiveRecord
       def str.inspect
         self
       end
+
       str
     end
   end
diff --git a/activerecord/lib/active_record/explain_registry.rb b/activerecord/lib/active_record/explain_registry.rb
new file mode 100644
index 0000000000..b84a692b87
--- /dev/null
+++ b/activerecord/lib/active_record/explain_registry.rb
@@ -0,0 +1,30 @@
+require 'active_support/per_thread_registry'
+
+module ActiveRecord
+  # This is a thread locals registry for EXPLAIN. For example
+  #
+  #   ActiveRecord::ExplainRegistry.queries
+  #
+  # returns the collected queries local to the current thread.
+  #
+  # See the documentation of <tt>ActiveSupport::PerThreadRegistry</tt>
+  # for further details.
+  class ExplainRegistry
+    extend ActiveSupport::PerThreadRegistry
+
+    attr_accessor :queries, :collect
+
+    def initialize
+      reset
+    end
+
+    def collect?
+      @collect
+    end
+
+    def reset
+      @collect = false
+      @queries = []
+    end
+  end
+end
\ No newline at end of file
diff --git a/activerecord/lib/active_record/explain_subscriber.rb b/activerecord/lib/active_record/explain_subscriber.rb
index 0f927496fb..a3bc56d600 100644
--- a/activerecord/lib/active_record/explain_subscriber.rb
+++ b/activerecord/lib/active_record/explain_subscriber.rb
@@ -1,4 +1,5 @@
 require 'active_support/notifications'
+require 'active_record/explain_registry'
 
 module ActiveRecord
   class ExplainSubscriber # :nodoc:
@@ -7,8 +8,8 @@ module ActiveRecord
     end
 
     def finish(name, id, payload)
-      if queries = Thread.current[:available_queries_for_explain]
-        queries << payload.values_at(:sql, :binds) unless ignore_payload?(payload)
+      if ExplainRegistry.collect? && !ignore_payload?(payload)
+        ExplainRegistry.queries << payload.values_at(:sql, :binds)
       end
     end