Make DescendantsTracker thread safe and optimize the #descendants method.

1 files changed, 34 insertions, 19 deletions

diff --git a/activesupport/lib/active_support/descendants_tracker.rb b/activesupport/lib/active_support/descendants_tracker.rb
index e2a8b4d4e3..27861e01d0 100644
--- a/activesupport/lib/active_support/descendants_tracker.rb
+++ b/activesupport/lib/active_support/descendants_tracker.rb
@@ -2,35 +2,50 @@ module ActiveSupport
   # This module provides an internal implementation to track descendants
   # which is faster than iterating through ObjectSpace.
   module DescendantsTracker
-    @@direct_descendants = Hash.new { |h, k| h[k] = [] }
+    @@direct_descendants = {}
 
-    def self.direct_descendants(klass)
-      @@direct_descendants[klass]
-    end
+    class << self
+      def direct_descendants(klass)
+        @@direct_descendants[klass] || []
+      end
 
-    def self.descendants(klass)
-      @@direct_descendants[klass].inject([]) do |descendants, _klass|
-        descendants << _klass
-        descendants.concat _klass.descendants
+      def descendants(klass)
+        arr = []
+        accumulate_descendants(klass, arr)
+        arr
       end
-    end
 
-    def self.clear
-      if defined? ActiveSupport::Dependencies
-        @@direct_descendants.each do |klass, descendants|
-          if ActiveSupport::Dependencies.autoloaded?(klass)
-            @@direct_descendants.delete(klass)
-          else
-            descendants.reject! { |v| ActiveSupport::Dependencies.autoloaded?(v) }
+      def clear
+        if defined? ActiveSupport::Dependencies
+          @@direct_descendants.each do |klass, descendants|
+            if ActiveSupport::Dependencies.autoloaded?(klass)
+              @@direct_descendants.delete(klass)
+            else
+              descendants.reject! { |v| ActiveSupport::Dependencies.autoloaded?(v) }
+            end
           end
+        else
+          @@direct_descendants.clear
+        end
+      end
+
+      # This is the only method that is not thread safe, but is only ever called
+      # during the eager loading phase.
+      def store_inherited(klass, descendant)
+        (@@direct_descendants[klass] ||= []) << descendant
+      end
+
+      private
+      def accumulate_descendants(klass, acc)
+        if direct_descendants = @@direct_descendants[klass]
+          acc.concat(direct_descendants)
+          direct_descendants.each { |direct_descendant| accumulate_descendants(direct_descendant, acc) }
         end
-      else
-        @@direct_descendants.clear
       end
     end
 
     def inherited(base)
-      self.direct_descendants << base
+      DescendantsTracker.store_inherited(self, base)
       super
     end