Add event object subscriptions to AS::Notifications

activesupport/lib/active_support/notifications/fanout.rb

@@ -71,11 +71,23 @@ def wait
       module Subscribers # :nodoc:
         def self.new(pattern, listener)
           if listener.respond_to?(:start) && listener.respond_to?(:finish)
-            subscriber = Evented.new pattern, listener
+            subscriber_class = Evented
           else
-            subscriber = Timed.new pattern, listener
+            if listener.respond_to?(:arity) && listener.arity == 1
+              subscriber_class = EventObject
+            else
+              subscriber_class = Timed
+            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 +142,27 @@ def finish(name, id, payload)
           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

activesupport/lib/active_support/notifications/instrumenter.rb

@@ -69,26 +69,34 @@ def initialize(name, start, ending, transaction_id, payload)
         @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

activesupport/test/notifications_test.rb

@@ -26,6 +26,42 @@ def event(*args)
     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"