Add instrument! to notifications which adds the result to the payload.

3 files changed, 21 insertions, 3 deletions

diff --git a/activesupport/lib/active_support/notifications.rb b/activesupport/lib/active_support/notifications.rb
index 3e96decb8c..a1383bb478 100644
--- a/activesupport/lib/active_support/notifications.rb
+++ b/activesupport/lib/active_support/notifications.rb
@@ -45,7 +45,7 @@ module ActiveSupport
     class << self
       attr_writer :notifier
       delegate :publish, :subscribe, :to => :notifier
-      delegate :instrument, :to => :instrumenter
+      delegate :instrument, :instrument!, :to => :instrumenter
 
       def notifier
         @notifier ||= Notifier.new
diff --git a/activesupport/lib/active_support/notifications/instrumenter.rb b/activesupport/lib/active_support/notifications/instrumenter.rb
index af35eb92cb..85c5199946 100644
--- a/activesupport/lib/active_support/notifications/instrumenter.rb
+++ b/activesupport/lib/active_support/notifications/instrumenter.rb
@@ -11,13 +11,21 @@ module ActiveSupport
         @notifier = notifier
       end
 
-      def instrument(name, payload={})
+      # Instrument the given block by measuring the time taken to execute it
+      # and publish it.
+      def instrument(name, payload={}, add_result=false)
         time = Time.now
         result = yield if block_given?
+        payload.merge!(:result => result) if add_result
         @notifier.publish(name, time, Time.now, @id, payload)
         result
       end
 
+      # The same as instrument, but adds the result as payload.
+      def instrument!(name, payload={}, &block)
+        instrument(name, payload, true, &block)
+      end
+
       private
         def unique_id
           SecureRandom.hex(10)
diff --git a/activesupport/test/notifications_test.rb b/activesupport/test/notifications_test.rb
index 226e12ff75..3690d723fa 100644
--- a/activesupport/test/notifications_test.rb
+++ b/activesupport/test/notifications_test.rb
@@ -83,10 +83,20 @@ module Notifications
   end
 
   class InstrumentationTest < TestCase
-    delegate :instrument, :to => ActiveSupport::Notifications
+    delegate :instrument, :instrument!, :to => ActiveSupport::Notifications
 
     def test_instrument_returns_block_result
       assert_equal 2, instrument(:awesome) { 1 + 1 }
+      drain
+    end
+
+    def test_instrument_with_band_adds_result_to_payload
+      assert_equal 2, instrument!(:awesome) { 1 + 1 }
+      drain
+
+      assert_equal 1, @events.size
+      assert_equal :awesome, @events.first.name
+      assert_equal Hash[:result => 2], @events.first.payload
     end
 
     def test_instrumenter_exposes_its_id