From 36e716cdc1beeb3557a02cffe8f5b625c0724034 Mon Sep 17 00:00:00 2001
From: Xavier Noria <fxn@hashref.com>
Date: Mon, 9 Aug 2021 13:15:10 +0200
Subject: [PATCH] Implement support for before_remove_const in zeitwek mode

This needs Zeitwerk 2.5 for the on_unload hook.
---
 activerecord/test/cases/base_test.rb          | 19 +++++++------
 .../dependencies/zeitwerk_integration.rb      |  8 +++++-
 .../application/zeitwerk_integration_test.rb  | 28 +++++++++++++++++++
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/activerecord/test/cases/base_test.rb b/activerecord/test/cases/base_test.rb
index 1fa3684e6c..547be94899 100644
--- a/activerecord/test/cases/base_test.rb
+++ b/activerecord/test/cases/base_test.rb
@@ -1266,17 +1266,18 @@ class BasicsTest < ActiveRecord::TestCase
     assert_equal c1, c2
   end
 
-  def test_current_scope_is_reset
-    Object.const_set :UnloadablePost, Class.new(ActiveRecord::Base)
-    UnloadablePost.current_scope = UnloadablePost.all
+  def test_before_remove_const_resets_the_current_scope
+    # Done this way because a class cannot be defined in a method using the
+    # class keyword.
+    Object.const_set(:ReloadableModel, Class.new(ActiveRecord::Base))
+    ReloadableModel.current_scope = ReloadableModel.all
+    assert_not_nil ActiveRecord::Scoping::ScopeRegistry.current_scope(ReloadableModel) # precondition
 
-    UnloadablePost.unloadable
-    klass = UnloadablePost
-    assert_not_nil ActiveRecord::Scoping::ScopeRegistry.current_scope(klass)
-    ActiveSupport::Dependencies.remove_unloadable_constants!
-    assert_nil ActiveRecord::Scoping::ScopeRegistry.current_scope(klass)
+    ReloadableModel.before_remove_const
+
+    assert_nil ActiveRecord::Scoping::ScopeRegistry.current_scope(ReloadableModel)
   ensure
-    Object.class_eval { remove_const :UnloadablePost } if defined?(UnloadablePost)
+    Object.send(:remove_const, :ReloadableModel)
   end
 
   def test_marshal_round_trip
diff --git a/activesupport/lib/active_support/dependencies/zeitwerk_integration.rb b/activesupport/lib/active_support/dependencies/zeitwerk_integration.rb
index f7aec0ffb3..fee2cc63d4 100644
--- a/activesupport/lib/active_support/dependencies/zeitwerk_integration.rb
+++ b/activesupport/lib/active_support/dependencies/zeitwerk_integration.rb
@@ -89,7 +89,13 @@ module ActiveSupport
               autoloader.do_not_eager_load(autoload_path) unless eager_load?(autoload_path)
             end
 
-            Rails.autoloaders.main.enable_reloading if enable_reloading
+            if enable_reloading
+              Rails.autoloaders.main.enable_reloading
+              Rails.autoloaders.main.on_unload do |_cpath, value, _abspath|
+                value.before_remove_const if value.respond_to?(:before_remove_const)
+              end
+            end
+
             Rails.autoloaders.each(&:setup)
           end
 
diff --git a/railties/test/application/zeitwerk_integration_test.rb b/railties/test/application/zeitwerk_integration_test.rb
index 5453f69cf4..99337282a8 100644
--- a/railties/test/application/zeitwerk_integration_test.rb
+++ b/railties/test/application/zeitwerk_integration_test.rb
@@ -341,6 +341,34 @@ class ZeitwerkIntegrationTest < ActiveSupport::TestCase
     assert_equal :no_op, deps.unhook!
   end
 
+  test "reloading invokes before_remove_const" do
+    $before_remove_const_invoked = false
+
+    app_file "app/models/foo.rb", <<~RUBY
+      # While the most common use case is classes/modules, the contract does not
+      # require values to be so. Let's weaken the test down to Object.new.
+      Foo = Object.new
+      def Foo.before_remove_const
+        $before_remove_const_invoked = true
+      end
+    RUBY
+
+    app_file "app/models/bar.rb", <<~RUBY
+      # This object does not implement before_remove_const. We define it to make
+      # sure reloading does not raise. That is, it does not blindly invoke the
+      # hook on all unloaded objects.
+      Bar = Object.new
+    RUBY
+
+    boot
+
+    assert Foo
+    assert Bar
+    ActiveSupport::Dependencies.clear
+
+    assert $before_remove_const_invoked
+  end
+
   test "autoloaders.logger=" do
     boot