Allow lazy load hooks to be executed only once

Provide run_once: true option to on_load in case you want a hook only to be executed once. This may be useful in cases where executing a hook several times may have undesired side effects

2 files changed, 63 insertions, 8 deletions

diff --git a/activesupport/lib/active_support/lazy_load_hooks.rb b/activesupport/lib/active_support/lazy_load_hooks.rb
index c23b319046..c124416595 100644
--- a/activesupport/lib/active_support/lazy_load_hooks.rb
+++ b/activesupport/lib/active_support/lazy_load_hooks.rb
@@ -27,11 +27,17 @@ module ActiveSupport
       base.class_eval do
         @load_hooks = Hash.new { |h, k| h[k] = [] }
         @loaded     = Hash.new { |h, k| h[k] = [] }
+        @run_once   = Hash.new { |h, k| h[k] = [] }
       end
     end
 
     # Declares a block that will be executed when a Rails component is fully
     # loaded.
+    #
+    # Options:
+    #
+    # * <tt>:yield</tt> - Yields the object that run_load_hooks to +block+.
+    # * <tt>:run_once</tt> - Given +block+ will run only once.
     def on_load(name, options = {}, &block)
       @loaded[name].each do |base|
         execute_hook(base, options, block)
@@ -40,20 +46,32 @@ module ActiveSupport
       @load_hooks[name] << [block, options]
     end
 
-    def execute_hook(base, options, block)
-      if options[:yield]
-        block.call(base)
-      else
-        base.instance_eval(&block)
-      end
-    end
-
     def run_load_hooks(name, base = Object)
       @loaded[name] << base
       @load_hooks[name].each do |hook, options|
         execute_hook(base, options, hook)
       end
     end
+
+    private
+
+      def with_execution_control(name, block, once)
+        unless @run_once[name].include?(block)
+          @run_once[name] << block if once
+
+          yield
+        end
+      end
+
+      def execute_hook(base, options, block)
+        with_execution_control(name, block, options[:run_once]) do
+          if options[:yield]
+            block.call(base)
+          else
+            base.instance_eval(&block)
+          end
+        end
+      end
   end
 
   extend LazyLoadHooks
diff --git a/activesupport/test/lazy_load_hooks_test.rb b/activesupport/test/lazy_load_hooks_test.rb
index 9c9264e8fc..c161005100 100644
--- a/activesupport/test/lazy_load_hooks_test.rb
+++ b/activesupport/test/lazy_load_hooks_test.rb
@@ -20,6 +20,18 @@ class LazyLoadHooksTest < ActiveSupport::TestCase
     assert_equal 7, i
   end
 
+  def test_basic_hook_with_two_registrations_only_once
+    i = 0
+    ActiveSupport.on_load(:basic_hook_with_two_once, run_once: true) do
+      i += incr
+    end
+    assert_equal 0, i
+    ActiveSupport.run_load_hooks(:basic_hook_with_two_once, FakeContext.new(2))
+    assert_equal 2, i
+    ActiveSupport.run_load_hooks(:basic_hook_with_two_once, FakeContext.new(5))
+    assert_equal 2, i
+  end
+
   def test_hook_registered_after_run
     i = 0
     ActiveSupport.run_load_hooks(:registered_after)
@@ -37,6 +49,15 @@ class LazyLoadHooksTest < ActiveSupport::TestCase
     assert_equal 7, i
   end
 
+  def test_hook_registered_after_run_with_two_registrations_only_once
+    i = 0
+    ActiveSupport.run_load_hooks(:registered_after_with_two_once, FakeContext.new(2))
+    ActiveSupport.run_load_hooks(:registered_after_with_two_once, FakeContext.new(5))
+    assert_equal 0, i
+    ActiveSupport.on_load(:registered_after_with_two_once, run_once: true) { i += incr }
+    assert_equal 2, i
+  end
+
   def test_hook_registered_interleaved_run_with_two_registrations
     i = 0
     ActiveSupport.run_load_hooks(:registered_interleaved_with_two, FakeContext.new(2))
@@ -47,6 +68,22 @@ class LazyLoadHooksTest < ActiveSupport::TestCase
     assert_equal 7, i
   end
 
+  def test_hook_registered_interleaved_run_with_two_registrations_once
+    i = 0
+    ActiveSupport
+      .run_load_hooks(:registered_interleaved_with_two_once, FakeContext.new(2))
+    assert_equal 0, i
+
+    ActiveSupport.on_load(:registered_interleaved_with_two_once, run_once: true) do
+      i += incr
+    end
+    assert_equal 2, i
+
+    ActiveSupport
+      .run_load_hooks(:registered_interleaved_with_two_once, FakeContext.new(5))
+    assert_equal 2, i
+  end
+
   def test_hook_receives_a_context
     i = 0
     ActiveSupport.on_load(:contextual) { i += incr }