Add ActionDispatch::Notifications middleware.

10 files changed, 122 insertions, 43 deletions

diff --git a/actionpack/lib/action_dispatch.rb b/actionpack/lib/action_dispatch.rb
index 763a090de2..082562d921 100644
--- a/actionpack/lib/action_dispatch.rb
+++ b/actionpack/lib/action_dispatch.rb
@@ -46,6 +46,7 @@ module ActionDispatch
     autoload :Cookies
     autoload :Flash
     autoload :Head
+    autoload :Notifications
     autoload :ParamsParser
     autoload :Rescue
     autoload :ShowExceptions
diff --git a/actionpack/lib/action_dispatch/middleware/callbacks.rb b/actionpack/lib/action_dispatch/middleware/callbacks.rb
index 5ec406e134..d07841218a 100644
--- a/actionpack/lib/action_dispatch/middleware/callbacks.rb
+++ b/actionpack/lib/action_dispatch/middleware/callbacks.rb
@@ -45,8 +45,6 @@ module ActionDispatch
         run_callbacks(:prepare) if @prepare_each_request
         @app.call(env)
       end
-    ensure
-      ActiveSupport::Notifications.instrument "action_dispatch.callback"
     end
   end
 end
diff --git a/actionpack/lib/action_dispatch/middleware/notifications.rb b/actionpack/lib/action_dispatch/middleware/notifications.rb
new file mode 100644
index 0000000000..01d2cbb435
--- /dev/null
+++ b/actionpack/lib/action_dispatch/middleware/notifications.rb
@@ -0,0 +1,24 @@
+module ActionDispatch
+  # Provide notifications in the middleware stack. Notice that for the before_dispatch
+  # and after_dispatch notifications, we just send the original env, so we don't pile
+  # up large env hashes in the queue. However, in exception cases, the whole env hash
+  # is actually useful, so we send it all.
+  class Notifications
+    def initialize(app)
+      @app = app
+    end
+
+    def call(stack_env)
+      env = stack_env.dup
+      ActiveSupport::Notifications.instrument("action_dispatch.before_dispatch", :env => env)
+
+      ActiveSupport::Notifications.instrument!("action_dispatch.after_dispatch", :env => env) do
+        @app.call(stack_env)
+      end
+    rescue Exception => exception
+      ActiveSupport::Notifications.instrument('action_dispatch.exception',
+        :env => stack_env, :exception => exception)
+      raise exception
+    end
+  end
+end
\ No newline at end of file
diff --git a/actionpack/lib/action_dispatch/middleware/show_exceptions.rb b/actionpack/lib/action_dispatch/middleware/show_exceptions.rb
index 56fd4f45c0..0977f5c60a 100644
--- a/actionpack/lib/action_dispatch/middleware/show_exceptions.rb
+++ b/actionpack/lib/action_dispatch/middleware/show_exceptions.rb
@@ -61,11 +61,8 @@ module ActionDispatch
     def call(env)
       @app.call(env)
     rescue Exception => exception
-      ActiveSupport::Notifications.instrument 'action_dispatch.show_exception',
-        :env => env, :exception => exception do
-        raise exception if env['action_dispatch.show_exceptions'] == false
-        render_exception(env, exception)
-      end
+      raise exception if env['action_dispatch.show_exceptions'] == false
+      render_exception(env, exception)
     end
 
     private
diff --git a/actionpack/test/dispatch/callbacks_test.rb b/actionpack/test/dispatch/callbacks_test.rb
index f3ea5209f4..9df882ce75 100644
--- a/actionpack/test/dispatch/callbacks_test.rb
+++ b/actionpack/test/dispatch/callbacks_test.rb
@@ -85,18 +85,6 @@ class DispatcherTest < Test::Unit::TestCase
     assert_equal 4, Foo.b
   end
 
-  def test_should_send_an_instrumentation_callback_for_async_processing
-    ActiveSupport::Notifications.expects(:instrument).with("action_dispatch.callback")
-    dispatch
-  end
-
-  def test_should_send_an_instrumentation_callback_for_async_processing_even_on_failure
-    ActiveSupport::Notifications.notifier.expects(:publish)
-    assert_raise RuntimeError do
-      dispatch { |env| raise "OMG" }
-    end
-  end
-
   private
 
     def dispatch(cache_classes = true, &block)
diff --git a/actionpack/test/dispatch/notifications_test.rb b/actionpack/test/dispatch/notifications_test.rb
new file mode 100644
index 0000000000..d834a734ef
--- /dev/null
+++ b/actionpack/test/dispatch/notifications_test.rb
@@ -0,0 +1,69 @@
+require 'abstract_unit'
+
+class NotificationsMiddlewareTest < ActionController::IntegrationTest
+  Boomer = lambda do |env|
+    req = ActionDispatch::Request.new(env)
+    case req.path
+    when "/"
+      [200, {}, []]
+    else
+      raise "puke!"
+    end
+  end
+
+  App = ActionDispatch::Notifications.new(Boomer)
+
+  def setup
+    @queue    = ActiveSupport::Notifications::Fanout.new
+    @notifier = ActiveSupport::Notifications::Notifier.new(@queue)
+    ActiveSupport::Notifications.notifier = @notifier
+
+    @events = []
+    ActiveSupport::Notifications.subscribe do |*args|
+      @events << args
+    end
+
+    @app = App
+  end
+
+  test "publishes notifications" do
+    get "/"
+    ActiveSupport::Notifications.notifier.wait
+
+    assert_equal 2, @events.size
+    before, after = @events
+
+    assert_equal 'action_dispatch.before_dispatch', before[0]
+    assert_kind_of Hash, before[4][:env]
+    assert_equal 'GET',  before[4][:env]["REQUEST_METHOD"]
+
+    assert_equal 'action_dispatch.after_dispatch', after[0]
+    assert_kind_of Hash, after[4][:env]
+    assert_equal 'GET',  after[4][:env]["REQUEST_METHOD"]
+  end
+
+  test "publishes notifications on failure" do
+    begin
+      get "/puke"
+    rescue
+    end
+
+    ActiveSupport::Notifications.notifier.wait
+
+    assert_equal 3, @events.size
+    before, after, exception = @events
+
+    assert_equal 'action_dispatch.before_dispatch', before[0]
+    assert_kind_of Hash, before[4][:env]
+    assert_equal 'GET',  before[4][:env]["REQUEST_METHOD"]
+
+    assert_equal 'action_dispatch.after_dispatch', after[0]
+    assert_kind_of Hash, after[4][:env]
+    assert_equal 'GET',  after[4][:env]["REQUEST_METHOD"]  
+
+    assert_equal 'action_dispatch.exception', exception[0]
+    assert_kind_of Hash, exception[4][:env]
+    assert_equal 'GET',  exception[4][:env]["REQUEST_METHOD"]  
+    assert_kind_of RuntimeError, exception[4][:exception]
+  end
+end
\ No newline at end of file
diff --git a/actionpack/test/dispatch/show_exceptions_test.rb b/actionpack/test/dispatch/show_exceptions_test.rb
index 951fb4a22e..9f6a93756c 100644
--- a/actionpack/test/dispatch/show_exceptions_test.rb
+++ b/actionpack/test/dispatch/show_exceptions_test.rb
@@ -104,27 +104,4 @@ class ShowExceptionsTest < ActionController::IntegrationTest
     assert_response 405
     assert_match /ActionController::MethodNotAllowed/, body
   end
-
-  test "publishes notifications" do
-    # Wait pending notifications to be published
-    ActiveSupport::Notifications.notifier.wait
-
-    @app, event = ProductionApp, nil
-    self.remote_addr = '127.0.0.1'
-
-    ActiveSupport::Notifications.subscribe('action_dispatch.show_exception') do |*args|
-      event = args
-    end
-
-    get "/"
-    assert_response 500
-    assert_match /puke/, body
-
-    ActiveSupport::Notifications.notifier.wait
-
-    assert_equal 'action_dispatch.show_exception', event.first
-    assert_kind_of Hash, event.last[:env]
-    assert_equal 'GET', event.last[:env]["REQUEST_METHOD"]
-    assert_kind_of RuntimeError, event.last[:exception]
-  end
 end
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 f3d877efe7..7c5b118ee3 100644
--- a/activesupport/lib/active_support/notifications/instrumenter.rb
+++ b/activesupport/lib/active_support/notifications/instrumenter.rb
@@ -20,6 +20,15 @@ module ActiveSupport
         result
       end
 
+      # The same as instrument, but sends the notification even if the yielded
+      # block raises an error.
+      def instrument!(name, payload={})
+        time = Time.now
+        yield(payload) if block_given?
+      ensure
+        @notifier.publish(name, time, Time.now, @id, payload)
+      end
+
       private
         def unique_id
           SecureRandom.hex(10)
diff --git a/activesupport/test/notifications_test.rb b/activesupport/test/notifications_test.rb
index c3eb1a4eb5..c41d81fe7e 100644
--- a/activesupport/test/notifications_test.rb
+++ b/activesupport/test/notifications_test.rb
@@ -90,6 +90,22 @@ module Notifications
       drain
     end
 
+    def test_instrument_with_bang_returns_result_even_on_failure
+      begin
+        instrument!(:awesome, :payload => "notifications") do
+          raise "OMG"
+        end
+        flunk
+      rescue
+      end
+
+      drain
+
+      assert_equal 1, @events.size
+      assert_equal :awesome, @events.last.name
+      assert_equal Hash[:payload => "notifications"], @events.last.payload
+    end
+
     def test_instrument_yields_the_paylod_for_further_modification
       assert_equal 2, instrument(:awesome) { |p| p[:result] = 1 + 1 }
       drain