Filter reloaded classes in Class#subclasses and Class#descendants core exts

When calling `#descendants` on a non-reloadable class with reloadable
descendants, it can return classes that were unloaded but not yet
garbage collected.

`DescendantsTracker` have been dealing with this for a long time
but the core_ext versions of these methods were never made reloader
aware.

This also refactor `DescendantsTracker` to not be so different when
there is no native `#subclasses` method.

Author: byroot

activesupport/CHANGELOG.md

@@ -1,3 +1,13 @@
+*   `Class#subclasses` and `Class#descendants` now automatically filter reloaded classes.
+
+    Previously they could return old implementations of reloadable classes that have been
+    dereferenced but not yet garbage collected.
+
+    They now automatically filter such classes like `DescendantTracker#subclasses` and
+    `DescendantTracker#descendants`.
+
+    *Jean Boussier*
+
 *   `Rails.error.report` now marks errors as reported to avoid reporting them twice.
 
     In some cases, users might want to report errors explicitly with some extra context

activesupport/lib/active_support/core_ext/class/subclasses.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "active_support/ruby_features"
+require "active_support/descendants_tracker"
 
 class Class
   if ActiveSupport::RubyFeatures::CLASS_SUBCLASSES
@@ -26,16 +27,18 @@ def descendants
         k.singleton_class? || k == self
       end
     end
+
+    # Returns an array with the direct children of +self+.
+    #
+    #   class Foo; end
+    #   class Bar < Foo; end
+    #   class Baz < Bar; end
+    #
+    #   Foo.subclasses # => [Bar]
+    def subclasses
+      descendants.select { |descendant| descendant.superclass == self }
+    end
   end
 
-  # Returns an array with the direct children of +self+.
-  #
-  #   class Foo; end
-  #   class Bar < Foo; end
-  #   class Baz < Bar; end
-  #
-  #   Foo.subclasses # => [Bar]
-  def subclasses
-    descendants.select { |descendant| descendant.superclass == self }
-  end unless ActiveSupport::RubyFeatures::CLASS_SUBCLASSES
+  prepend ActiveSupport::DescendantsTracker::ReloadedClassesFiltering
 end

activesupport/lib/active_support/descendants_tracker.rb

@@ -5,184 +5,122 @@
 
 module ActiveSupport
   # This module provides an internal implementation to track descendants
-  # which is faster than iterating through ObjectSpace.
+  # which is faster than iterating through +ObjectSpace+.
+  #
+  # However Ruby 3.1 provide a fast native +Class#subclasses+ method,
+  # so if you know your code won't be executed on older rubies, including
+  # +ActiveSupport::DescendantsTracker+ does not provide any benefit.
   module DescendantsTracker
-    class << self
-      def direct_descendants(klass)
-        ActiveSupport::Deprecation.warn(<<~MSG)
-          ActiveSupport::DescendantsTracker.direct_descendants is deprecated and will be removed in Rails 7.1.
-          Use ActiveSupport::DescendantsTracker.subclasses instead.
-        MSG
-        subclasses(klass)
-      end
-    end
-
     @clear_disabled = false
 
-    if RubyFeatures::CLASS_SUBCLASSES
-      @@excluded_descendants = if RUBY_ENGINE == "ruby"
-        # On MRI `ObjectSpace::WeakMap` keys are weak references.
-        # So we can simply use WeakMap as a `Set`.
-        ObjectSpace::WeakMap.new
-      else
-        # On TruffleRuby `ObjectSpace::WeakMap` keys are strong references.
-        # So we use `object_id` as a key and the actual object as a value.
-        #
-        # JRuby for now doesn't have Class#descendant, but when it will, it will likely
-        # have the same WeakMap semantic than Truffle so we future proof this as much as possible.
-        class WeakSet # :nodoc:
-          def initialize
-            @map = ObjectSpace::WeakMap.new
-          end
+    if RUBY_ENGINE == "ruby"
+      # On MRI `ObjectSpace::WeakMap` keys are weak references.
+      # So we can simply use WeakMap as a `Set`.
+      class WeakSet < ObjectSpace::WeakMap # :nodoc:
+        alias_method :to_a, :keys
 
-          def [](object)
-            @map.key?(object.object_id)
-          end
-
-          def []=(object, _present)
-            @map[object.object_id] = object
-          end
+        def <<(object)
+          self[object] = true
         end
-        WeakSet.new
       end
-
-      class << self
-        def disable_clear! # :nodoc:
-          unless @clear_disabled
-            @clear_disabled = true
-            remove_method(:subclasses)
-            @@excluded_descendants = nil
-          end
+    else
+      # On TruffleRuby `ObjectSpace::WeakMap` keys are strong references.
+      # So we use `object_id` as a key and the actual object as a value.
+      #
+      # JRuby for now doesn't have Class#descendant, but when it will, it will likely
+      # have the same WeakMap semantic than Truffle so we future proof this as much as possible.
+      class WeakSet # :nodoc:
+        def initialize
+          @map = ObjectSpace::WeakMap.new
         end
 
-        def subclasses(klass)
-          klass.subclasses
+        def [](object)
+          @map.key?(object.object_id)
         end
+        alias_method :include?, :[]
 
-        def descendants(klass)
-          klass.descendants
+        def []=(object, _present)
+          @map[object.object_id] = object
         end
 
-        def clear(classes) # :nodoc:
-          raise "DescendantsTracker.clear was disabled because config.enable_reloading is false" if @clear_disabled
-
-          classes.each do |klass|
-            @@excluded_descendants[klass] = true
-            klass.descendants.each do |descendant|
-              @@excluded_descendants[descendant] = true
-            end
-          end
+        def to_a
+          @map.values
         end
 
-        def native? # :nodoc:
-          true
+        def <<(object)
+          self[object] = true
         end
       end
+    end
+    @excluded_descendants = WeakSet.new
 
+    module ReloadedClassesFiltering # :nodoc:
       def subclasses
-        subclasses = super
-        subclasses.reject! { |d| @@excluded_descendants[d] }
-        subclasses
+        DescendantsTracker.reject!(super)
       end
 
       def descendants
-        subclasses.concat(subclasses.flat_map(&:descendants))
+        DescendantsTracker.reject!(super)
       end
+    end
 
-      def direct_descendants
-        ActiveSupport::Deprecation.warn(<<~MSG)
-          ActiveSupport::DescendantsTracker#direct_descendants is deprecated and will be removed in Rails 7.1.
-          Use #subclasses instead.
-        MSG
-        subclasses
-      end
-    else
-      @@direct_descendants = {}
+    include ReloadedClassesFiltering
 
-      class << self
-        def disable_clear! # :nodoc:
+    class << self
+      def disable_clear! # :nodoc:
+        unless @clear_disabled
           @clear_disabled = true
+          ReloadedClassesFiltering.remove_method(:subclasses)
+          ReloadedClassesFiltering.remove_method(:descendants)
+          @excluded_descendants = nil
         end
+      end
 
-        def subclasses(klass)
-          descendants = @@direct_descendants[klass]
-          descendants ? descendants.to_a : []
-        end
+      def clear(classes) # :nodoc:
+        raise "DescendantsTracker.clear was disabled because config.enable_reloading is false" if @clear_disabled
 
-        def descendants(klass)
-          arr = []
-          accumulate_descendants(klass, arr)
-          arr
-        end
-
-        def clear(classes) # :nodoc:
-          raise "DescendantsTracker.clear was disabled because config.enable_reloading is false" if @clear_disabled
-
-          @@direct_descendants.each do |klass, direct_descendants_of_klass|
-            if classes.member?(klass)
-              @@direct_descendants.delete(klass)
-            else
-              direct_descendants_of_klass.reject! do |direct_descendant_of_class|
-                classes.member?(direct_descendant_of_class)
-              end
-            end
+        classes.each do |klass|
+          @excluded_descendants << klass
+          klass.descendants.each do |descendant|
+            @excluded_descendants << descendant
           end
         end
-
-        def native? # :nodoc:
-          false
-        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] ||= DescendantsArray.new) << descendant
-        end
-
-        private
-          def accumulate_descendants(klass, acc)
-            if direct_descendants = @@direct_descendants[klass]
-              direct_descendants.each do |direct_descendant|
-                acc << direct_descendant
-                accumulate_descendants(direct_descendant, acc)
-              end
-            end
-          end
       end
 
-      def inherited(base)
-        DescendantsTracker.store_inherited(self, base)
-        super
+      def reject!(classes) # :nodoc:
+        if @excluded_descendants
+          classes.reject! { |d| @excluded_descendants.include?(d) }
+        end
+        classes
       end
+    end
 
-      def direct_descendants
-        ActiveSupport::Deprecation.warn(<<~MSG)
-          ActiveSupport::DescendantsTracker#direct_descendants is deprecated and will be removed in Rails 7.1.
-          Use #subclasses instead.
-        MSG
-        DescendantsTracker.subclasses(self)
-      end
+    def descendants
+      subclasses = self.subclasses
+      subclasses.concat(subclasses.flat_map(&:descendants))
+    end
 
-      def subclasses
-        DescendantsTracker.subclasses(self)
-      end
+    if RubyFeatures::CLASS_SUBCLASSES
+      class << self
+        def subclasses(klass)
+          klass.subclasses
+        end
 
-      def descendants
-        DescendantsTracker.descendants(self)
+        def descendants(klass)
+          klass.descendants
+        end
       end
-
+    else
       # DescendantsArray is an array that contains weak references to classes.
+      # Note: DescendantsArray is redundant with WeakSet, however WeakSet when used
+      # on Ruby 2.7 or 3.0 can trigger a Ruby crash: https://bugs.ruby-lang.org/issues/18928
       class DescendantsArray # :nodoc:
         include Enumerable
 
         def initialize
           @refs = []
         end
 
-        def initialize_copy(orig)
-          @refs = @refs.dup
-        end
-
         def <<(klass)
           @refs << WeakRef.new(klass)
         end
@@ -213,6 +151,39 @@ def reject!
           end
         end
       end
+
+      @direct_descendants = {}
+
+      def inherited(base) # :nodoc:
+        DescendantsTracker.store_inherited(self, base)
+        super
+      end
+
+      class << self
+        def subclasses(klass)
+          descendants = @direct_descendants[klass]
+          descendants ? DescendantsTracker.reject!(descendants.to_a) : []
+        end
+
+        def descendants(klass)
+          subclasses = self.subclasses(klass)
+          subclasses.concat(subclasses.flat_map { |k| descendants(k) })
+        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) # :nodoc:
+          (@direct_descendants[klass] ||= DescendantsArray.new) << descendant
+        end
+      end
+
+      def subclasses
+        DescendantsTracker.subclasses(self)
+      end
+
+      def descendants
+        DescendantsTracker.descendants(self)
+      end
     end
   end
 end

activesupport/test/core_ext/class_test.rb

@@ -2,6 +2,7 @@
 
 require_relative "../abstract_unit"
 require "active_support/core_ext/class"
+require "active_support/descendants_tracker"
 require "set"
 
 class ClassTest < ActiveSupport::TestCase
@@ -37,4 +38,18 @@ def test_subclasses_excludes_singleton_classes
     klass = Parent.new.singleton_class
     assert_not Parent.subclasses.include?(klass), "subclasses should not include singleton classes"
   end
+
+  def test_subclasses_exclude_reloaded_classes
+    subclass = Class.new(Parent)
+    assert_includes Parent.subclasses, subclass
+    ActiveSupport::DescendantsTracker.clear(Set[subclass])
+    assert_not_includes Parent.subclasses, subclass
+  end
+
+  def test_descendants_exclude_reloaded_classes
+    subclass = Class.new(Parent)
+    assert_includes Parent.descendants, subclass
+    ActiveSupport::DescendantsTracker.clear(Set[subclass])
+    assert_not_includes Parent.descendants, subclass
+  end
 end