Fix DescendantTracker.clear on Ruby 3.1

Previously I assumed it was useless, however I was wrong.

The method is called by the reloader to give the illusion that
the GC is precise. Meaning a class that will be unloaded is
immediately made invisible without waiting for it to be garbage collected.

This is easy to do up to Ruby 3.0 because `DescendantTracker` keeps
a map of all tracked classes.

However on 3.1 we need to use the inverse strategy, we keep a WeakMap
of all the classes we cleared, and we filter the return value of `descendants`
and `subclasses`.

Since `clear` is private API and is only used when reloading is enabled,
to reduce the performance impact in production mode, we entirely remove
this behavior when `config.cache_classes` is enabled.

Author: byroot

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

@@ -20,7 +20,7 @@ def descendants
     ObjectSpace.each_object(singleton_class).reject do |k|
       k.singleton_class? || k == self
     end
-  end unless ActiveSupport::RubyFeatures::CLASS_DESCENDANTS
+  end unless ActiveSupport::RubyFeatures::CLASS_SUBCLASSES
 
   # Returns an array with the direct children of +self+.
   #

activesupport/lib/active_support/descendants_tracker.rb

@@ -17,8 +17,44 @@ def direct_descendants(klass)
       end
     end
 
-    if RubyFeatures::CLASS_DESCENDANTS
+    @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
+          def initialize
+            @map = ObjectSpace::WeakMap.new
+          end
+
+          def [](object)
+            @map.key?(object.object_id)
+          end
+
+          def []=(object, _present)
+            @map[object_id] = object
+          end
+        end
+      end
+
       class << self
+        def disable_clear! # :nodoc:
+          unless @clear_disabled
+            @clear_disabled = true
+            remove_method(:subclasses)
+            remove_method(:descendants)
+            @@excluded_descendants = nil
+          end
+        end
+
         def subclasses(klass)
           klass.subclasses
         end
@@ -27,19 +63,32 @@ def descendants(klass)
           klass.descendants
         end
 
-        def clear(only: nil) # :nodoc:
-          # noop
+        def clear(classes) # :nodoc:
+          raise "DescendantsTracker.clear was disabled because config.cache_classes = true" if @clear_disabled
+
+          classes.each do |klass|
+            @@excluded_descendants[klass] = true
+            klass.descendants.each do |descendant|
+              @@excluded_descendants[descendant] = true
+            end
+          end
         end
 
         def native? # :nodoc:
           true
         end
       end
 
-      unless RubyFeatures::CLASS_SUBCLASSES
-        def subclasses
-          descendants.select { |descendant| descendant.superclass == self }
-        end
+      def subclasses
+        subclasses = super
+        subclasses.reject! { |d| @@excluded_descendants[d] }
+        subclasses
+      end
+
+      def descendants
+        descendants = super
+        descendants.reject! { |d| @@excluded_descendants[d] }
+        descendants
       end
 
       def direct_descendants
@@ -53,6 +102,10 @@ def direct_descendants
       @@direct_descendants = {}
 
       class << self
+        def disable_clear! # :nodoc:
+          @clear_disabled = true
+        end
+
         def subclasses(klass)
           descendants = @@direct_descendants[klass]
           descendants ? descendants.to_a : []
@@ -64,18 +117,15 @@ def descendants(klass)
           arr
         end
 
-        def clear(only: nil) # :nodoc:
-          if only.nil?
-            @@direct_descendants.clear
-            return
-          end
+        def clear(classes) # :nodoc:
+          raise "DescendantsTracker.clear was disabled because config.cache_classes = true" if @clear_disabled
 
           @@direct_descendants.each do |klass, direct_descendants_of_klass|
-            if only.member?(klass)
+            if classes.member?(klass)
               @@direct_descendants.delete(klass)
             else
               direct_descendants_of_klass.reject! do |direct_descendant_of_class|
-                only.member?(direct_descendant_of_class)
+                classes.member?(direct_descendant_of_class)
               end
             end
           end

activesupport/lib/active_support/ruby_features.rb

@@ -2,7 +2,6 @@
 
 module ActiveSupport
   module RubyFeatures # :nodoc:
-    CLASS_DESCENDANTS = Class.method_defined?(:descendants) # RUBY_VERSION >= "3.1"
     CLASS_SUBCLASSES = Class.method_defined?(:subclasses) # RUBY_VERSION >= "3.1"
   end
 end

activesupport/test/descendants_tracker_test.rb

@@ -11,7 +11,6 @@ class DescendantsTrackerTest < ActiveSupport::TestCase
       @original_state.each { |k, v| @original_state[k] = v.dup }
     end
 
-    ActiveSupport::DescendantsTracker.clear
     eval <<~RUBY
       class Parent
         extend ActiveSupport::DescendantsTracker
@@ -90,17 +89,8 @@ class Grandchild2 < Child1
     end
   end
 
-  test ".clear deletes all state" do
-    ActiveSupport::DescendantsTracker.clear
-    if ActiveSupport::DescendantsTracker.class_variable_defined?(:@@direct_descendants)
-      assert_empty ActiveSupport::DescendantsTracker.class_variable_get(:@@direct_descendants)
-    end
-  end
-
-  test ".clear(only) deletes the given classes only" do
-    skip "Irrelevant for native Class#descendants" if ActiveSupport::DescendantsTracker.native?
-
-    ActiveSupport::DescendantsTracker.clear(only: Set[Child2, Grandchild1])
+  test ".clear(classes) deletes the given classes only" do
+    ActiveSupport::DescendantsTracker.clear(Set[Child2, Grandchild1])
 
     assert_equal_sets [Child1, Grandchild2], Parent.descendants
     assert_equal_sets [Grandchild2], Child1.descendants

railties/lib/rails/application/finisher.rb

@@ -176,9 +176,7 @@ def self.complete(_state)
       initializer :set_clear_dependencies_hook, group: :all do |app|
         callback = lambda do
           # Order matters.
-          ActiveSupport::DescendantsTracker.clear(
-            only: ActiveSupport::Dependencies._autoloaded_tracked_classes
-          )
+          ActiveSupport::DescendantsTracker.clear(ActiveSupport::Dependencies._autoloaded_tracked_classes)
           ActiveSupport::Dependencies.clear
         end
 
@@ -194,6 +192,7 @@ def self.complete(_state)
 
         if config.cache_classes
           # No reloader
+          ActiveSupport::DescendantsTracker.disable_clear!
         elsif config.reload_classes_only_on_change
           reloader = config.file_watcher.new(*watchable_args, &callback)
           reloaders << reloader