Make DescendantsTracker thread safe and optimize the #descendants method.

3 files changed, 51 insertions, 30 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
 
diff --git a/activesupport/test/descendants_tracker_test_cases.rb b/activesupport/test/descendants_tracker_test_cases.rb
index 066ec8549b..8f8feb462c 100644
--- a/activesupport/test/descendants_tracker_test_cases.rb
+++ b/activesupport/test/descendants_tracker_test_cases.rb
@@ -1,3 +1,5 @@
+require 'set'
+
 module DescendantsTrackerTestCases
   class Parent
     extend ActiveSupport::DescendantsTracker
@@ -18,15 +20,15 @@ module DescendantsTrackerTestCases
   ALL = [Parent, Child1, Child2, Grandchild1, Grandchild2]
 
   def test_descendants
-    assert_equal [Child1, Grandchild1, Grandchild2, Child2], Parent.descendants
-    assert_equal [Grandchild1, Grandchild2], Child1.descendants
-    assert_equal [], Child2.descendants
+    assert_equal_sets [Child1, Grandchild1, Grandchild2, Child2], Parent.descendants
+    assert_equal_sets [Grandchild1, Grandchild2], Child1.descendants
+    assert_equal_sets [], Child2.descendants
   end
 
   def test_direct_descendants
-    assert_equal [Child1, Child2], Parent.direct_descendants
-    assert_equal [Grandchild1, Grandchild2], Child1.direct_descendants
-    assert_equal [], Child2.direct_descendants
+    assert_equal_sets [Child1, Child2], Parent.direct_descendants
+    assert_equal_sets [Grandchild1, Grandchild2], Child1.direct_descendants
+    assert_equal_sets [], Child2.direct_descendants
   end
 
   def test_clear
@@ -40,6 +42,10 @@ module DescendantsTrackerTestCases
 
   protected
 
+  def assert_equal_sets(expected, actual)
+    Set.new(expected) == Set.new(actual)
+  end
+
   def mark_as_autoloaded(*klasses)
     # If ActiveSupport::Dependencies is not loaded, forget about autoloading.
     # This allows using AS::DescendantsTracker without AS::Dependencies.
diff --git a/activesupport/test/descendants_tracker_with_autoloading_test.rb b/activesupport/test/descendants_tracker_with_autoloading_test.rb
index 9180f1f977..ab1be296f8 100644
--- a/activesupport/test/descendants_tracker_with_autoloading_test.rb
+++ b/activesupport/test/descendants_tracker_with_autoloading_test.rb
@@ -18,17 +18,17 @@ class DescendantsTrackerWithAutoloadingTest < ActiveSupport::TestCase
   def test_clear_with_autoloaded_children_and_granchildren
     mark_as_autoloaded Child1, Grandchild1, Grandchild2 do
       ActiveSupport::DescendantsTracker.clear
-      assert_equal [Child2], Parent.descendants
-      assert_equal [], Child2.descendants
+      assert_equal_sets [Child2], Parent.descendants
+      assert_equal_sets [], Child2.descendants
     end
   end
 
   def test_clear_with_autoloaded_granchildren
     mark_as_autoloaded Grandchild1, Grandchild2 do
       ActiveSupport::DescendantsTracker.clear
-      assert_equal [Child1, Child2], Parent.descendants
-      assert_equal [], Child1.descendants
-      assert_equal [], Child2.descendants
+      assert_equal_sets [Child1, Child2], Parent.descendants
+      assert_equal_sets [], Child1.descendants
+      assert_equal_sets [], Child2.descendants
     end
   end
 end