fire a notification when the request stops / starts

2 files changed, 88 insertions, 2 deletions

diff --git a/railties/lib/rails/rack/logger.rb b/railties/lib/rails/rack/logger.rb
index 78b4e68ce4..6ed6722c44 100644
--- a/railties/lib/rails/rack/logger.rb
+++ b/railties/lib/rails/rack/logger.rb
@@ -1,12 +1,17 @@
 require 'active_support/core_ext/time/conversions'
 require 'active_support/core_ext/object/blank'
+require 'active_support/log_subscriber'
+require 'action_dispatch/http/request'
+require 'rack/body_proxy'
 
 module Rails
   module Rack
     # Sets log tags, logs the request, calls the app, and flushes the logs.
     class Logger < ActiveSupport::LogSubscriber
       def initialize(app, taggers = nil)
-        @app, @taggers = app, taggers || []
+        @app          = app
+        @taggers      = taggers || []
+        @instrumenter = ActiveSupport::Notifications.instrumenter
       end
 
       def call(env)
@@ -28,8 +33,14 @@ module Rails
           logger.debug ''
         end
 
+        @instrumenter.start 'action_dispatch.request', request: request
         logger.info started_request_message(request)
-        @app.call(env)
+        resp = @app.call(env)
+        resp[2] = ::Rack::BodyProxy.new(resp[2]) { finish(request) }
+        resp
+      rescue
+        finish(request)
+        raise
       ensure
         ActiveSupport::LogSubscriber.flush_all!
       end
@@ -58,6 +69,10 @@ module Rails
 
       private
 
+      def finish(request)
+        @instrumenter.finish 'action_dispatch.request', request: request
+      end
+
       def development?
         Rails.env.development?
       end
diff --git a/railties/test/rack_logger_test.rb b/railties/test/rack_logger_test.rb
new file mode 100644
index 0000000000..3a9392fd26
--- /dev/null
+++ b/railties/test/rack_logger_test.rb
@@ -0,0 +1,71 @@
+require 'active_support/testing/autorun'
+require 'active_support/test_case'
+require 'rails/rack/logger'
+require 'logger'
+
+module Rails
+  module Rack
+    class LoggerTest < ActiveSupport::TestCase
+      class TestLogger < Rails::Rack::Logger
+        NULL = ::Logger.new File::NULL
+
+        attr_reader :logger
+
+        def initialize(logger = NULL, taggers = nil, &block)
+          super(->(_) { block.call; [200, {}, []] }, taggers)
+          @logger = logger
+        end
+
+        def development?; false; end
+      end
+
+      class Subscriber < Struct.new(:starts, :finishes)
+        def initialize(starts = [], finishes = [])
+          super
+        end
+
+        def start(name, id, payload)
+          starts << [name, id, payload]
+        end
+
+        def finish(name, id, payload)
+          finishes << [name, id, payload]
+        end
+      end
+
+      attr_reader :subscriber, :notifier
+
+      def setup
+        @subscriber = Subscriber.new
+        @notifier = ActiveSupport::Notifications.notifier
+        notifier.subscribe 'action_dispatch.request', subscriber
+      end
+
+      def teardown
+        notifier.unsubscribe subscriber
+      end
+
+      def test_notification
+        logger = TestLogger.new { }
+
+        assert_difference('subscriber.starts.length') do
+          assert_difference('subscriber.finishes.length') do
+            logger.call('REQUEST_METHOD' => 'GET').last.close
+          end
+        end
+      end
+
+      def test_notification_on_raise
+        logger = TestLogger.new { raise }
+
+        assert_difference('subscriber.starts.length') do
+          assert_difference('subscriber.finishes.length') do
+            assert_raises(RuntimeError) do
+              logger.call 'REQUEST_METHOD' => 'GET'
+            end
+          end
+        end
+      end
+    end
+  end
+end