Merge pull request #22428 from thedarkone/subs-notif

Subscribing to notifications while inside the said instrumented section

3 files changed, 24 insertions, 4 deletions

diff --git a/activesupport/lib/active_support/notifications/fanout.rb b/activesupport/lib/active_support/notifications/fanout.rb
index 7798c7ec60..c53f9c1039 100644
--- a/activesupport/lib/active_support/notifications/fanout.rb
+++ b/activesupport/lib/active_support/notifications/fanout.rb
@@ -42,8 +42,8 @@ module ActiveSupport
         listeners_for(name).each { |s| s.start(name, id, payload) }
       end
 
-      def finish(name, id, payload)
-        listeners_for(name).each { |s| s.finish(name, id, payload) }
+      def finish(name, id, payload, listeners = listeners_for(name))
+        listeners.each { |s| s.finish(name, id, payload) }
       end
 
       def publish(name, *args)
diff --git a/activesupport/lib/active_support/notifications/instrumenter.rb b/activesupport/lib/active_support/notifications/instrumenter.rb
index 075ddc2382..67f2ee1a7f 100644
--- a/activesupport/lib/active_support/notifications/instrumenter.rb
+++ b/activesupport/lib/active_support/notifications/instrumenter.rb
@@ -15,14 +15,15 @@ module ActiveSupport
       # and publish it. Notice that events get sent even if an error occurs
       # in the passed-in block.
       def instrument(name, payload={})
-        start name, payload
+        # some of the listeners might have state
+        listeners_state = start name, payload
         begin
           yield payload
         rescue Exception => e
           payload[:exception] = [e.class.name, e.message]
           raise e
         ensure
-          finish name, payload
+          finish_with_state listeners_state, name, payload
         end
       end
 
@@ -36,6 +37,10 @@ module ActiveSupport
         @notifier.finish name, @id, payload
       end
 
+      def finish_with_state(listeners_state, name, payload)
+        @notifier.finish name, @id, payload, listeners_state
+      end
+
       private
 
       def unique_id
diff --git a/activesupport/test/notifications_test.rb b/activesupport/test/notifications_test.rb
index f729f0a95b..d9cc392ac9 100644
--- a/activesupport/test/notifications_test.rb
+++ b/activesupport/test/notifications_test.rb
@@ -42,6 +42,21 @@ module Notifications
       ActiveSupport::Notifications.instrument(name)
       assert_equal expected, events
     end
+
+    def test_subsribing_to_instrumentation_while_inside_it
+      # the repro requires that there are no evented subscribers for the "foo" event,
+      # so we have to duplicate some of the setup code
+      old_notifier = ActiveSupport::Notifications.notifier
+      ActiveSupport::Notifications.notifier = ActiveSupport::Notifications::Fanout.new
+
+      ActiveSupport::Notifications.subscribe('foo', TestSubscriber.new)
+
+      ActiveSupport::Notifications.instrument('foo') do
+        ActiveSupport::Notifications.subscribe('foo') {}
+      end
+    ensure
+      ActiveSupport::Notifications.notifier = old_notifier
+    end
   end
 
   class UnsubscribeTest < TestCase