Merge pull request #33451 from rails/event-object-subscription

Add event object subscriptions to AS::Notifications

4 files changed, 120 insertions, 4 deletions

diff --git a/activesupport/CHANGELOG.md b/activesupport/CHANGELOG.md
index 10a25a120d..e789d7fa20 100644
--- a/activesupport/CHANGELOG.md
+++ b/activesupport/CHANGELOG.md
@@ -1,3 +1,37 @@
+*   Add "event object" support to the notification system.
+    Before this change, end users were forced to create hand made arsenal
+    event objects on their own, like this:
+
+        ActiveSupport::Notifications.subscribe('wait') do |*args|
+          @event = ActiveSupport::Notifications::Event.new(*args)
+        end
+        
+        ActiveSupport::Notifications.instrument('wait') do
+          sleep 1
+        end
+        
+        @event.duration # => 1000.138
+
+    After this change, if the block passed to `subscribe` only takes one
+    parameter, the framework will yield an event object to the block.  Now
+    end users are no longer required to make their own:
+
+        ActiveSupport::Notifications.subscribe('wait') do |event|
+          @event = event
+        end
+        
+        ActiveSupport::Notifications.instrument('wait') do
+          sleep 1
+        end
+        
+        p @event.allocations # => 7
+        p @event.cpu_time    # => 0.256
+        p @event.idle_time   # => 1003.2399
+
+    Now you can enjoy event objects without making them yourself.  Neat!
+
+    *Aaron "t.lo" Patterson*
+
 *   Add cpu_time, idle_time, and allocations to Event
 
     *Eileen M. Uchitelle*, *Aaron Patterson*
diff --git a/activesupport/lib/active_support/notifications/fanout.rb b/activesupport/lib/active_support/notifications/fanout.rb
index 25aab175b4..4e4ca70942 100644
--- a/activesupport/lib/active_support/notifications/fanout.rb
+++ b/activesupport/lib/active_support/notifications/fanout.rb
@@ -70,12 +70,29 @@ module ActiveSupport
 
       module Subscribers # :nodoc:
         def self.new(pattern, listener)
+          subscriber_class = Timed
+
           if listener.respond_to?(:start) && listener.respond_to?(:finish)
-            subscriber = Evented.new pattern, listener
+            subscriber_class = Evented
           else
-            subscriber = Timed.new pattern, listener
+            # Doing all this to detect a block like `proc { |x| }` vs
+            # `proc { |*x| }` or `proc { |**x| }`
+            if listener.respond_to?(:parameters)
+              params = listener.parameters
+              if params.length == 1 && params.first.first == :opt
+                subscriber_class = EventObject
+              end
+            end
           end
 
+          wrap_all pattern, subscriber_class.new(pattern, listener)
+        end
+
+        def self.event_object_subscriber(pattern, block)
+          wrap_all pattern, EventObject.new(pattern, block)
+        end
+
+        def self.wrap_all(pattern, subscriber)
           unless pattern
             AllMessages.new(subscriber)
           else
@@ -130,6 +147,27 @@ module ActiveSupport
           end
         end
 
+        class EventObject < Evented
+          def start(name, id, payload)
+            stack = Thread.current[:_event_stack] ||= []
+            event = build_event name, id, payload
+            event.start!
+            stack.push event
+          end
+
+          def finish(name, id, payload)
+            stack = Thread.current[:_event_stack]
+            event = stack.pop
+            event.finish!
+            @delegate.call event
+          end
+
+          private
+            def build_event(name, id, payload)
+              ActiveSupport::Notifications::Event.new name, nil, nil, id, payload
+            end
+        end
+
         class AllMessages # :nodoc:
           def initialize(delegate)
             @delegate = delegate
diff --git a/activesupport/lib/active_support/notifications/instrumenter.rb b/activesupport/lib/active_support/notifications/instrumenter.rb
index ddcd1661e4..455b7a44a6 100644
--- a/activesupport/lib/active_support/notifications/instrumenter.rb
+++ b/activesupport/lib/active_support/notifications/instrumenter.rb
@@ -69,26 +69,34 @@ module ActiveSupport
         @allocation_count_finish = 0
       end
 
+      # Record information at the time this event starts
       def start!
         @time = now
         @cpu_time_start = now_cpu
         @allocation_count_start = now_allocations
       end
 
+      # Record information at the time this event finishes
       def finish!
-        @end = now
         @cpu_time_finish = now_cpu
+        @end = now
         @allocation_count_finish = now_allocations
       end
 
+      # Returns the CPU time (in milliseconds) passed since the call to
+      # +start!+ and the call to +finish!+
       def cpu_time
-        @cpu_time_finish - @cpu_time_start
+        (@cpu_time_finish - @cpu_time_start) * 1000
       end
 
+      # Returns the idle time time (in milliseconds) passed since the call to
+      # +start!+ and the call to +finish!+
       def idle_time
         duration - cpu_time
       end
 
+      # Returns the number of allocations made since the call to +start!+ and
+      # the call to +finish!+
       def allocations
         @allocation_count_finish - @allocation_count_start
       end
diff --git a/activesupport/test/notifications_test.rb b/activesupport/test/notifications_test.rb
index d035f993f7..62817a839a 100644
--- a/activesupport/test/notifications_test.rb
+++ b/activesupport/test/notifications_test.rb
@@ -26,6 +26,42 @@ module Notifications
     end
   end
 
+  class SubscribeEventObjects < TestCase
+    def test_subscribe_events
+      events = []
+      @notifier.subscribe do |event|
+        events << event
+      end
+
+      ActiveSupport::Notifications.instrument("foo")
+      event = events.first
+      assert event, "should have an event"
+      assert_operator event.allocations, :>, 0
+      assert_operator event.cpu_time, :>, 0
+      assert_operator event.idle_time, :>, 0
+      assert_operator event.duration, :>, 0
+    end
+
+    def test_subscribe_via_top_level_api
+      old_notifier = ActiveSupport::Notifications.notifier
+      ActiveSupport::Notifications.notifier = ActiveSupport::Notifications::Fanout.new
+
+      event = nil
+      ActiveSupport::Notifications.subscribe("foo") do |e|
+        event = e
+      end
+
+      ActiveSupport::Notifications.instrument("foo") do
+        100.times { Object.new } # allocate at least 100 objects
+      end
+
+      assert event
+      assert_operator event.allocations, :>=, 100
+    ensure
+      ActiveSupport::Notifications.notifier = old_notifier
+    end
+  end
+
   class SubscribedTest < TestCase
     def test_subscribed
       name     = "foo"