Simplify the Rack middleware to write request log immediately

Previously, we would only setup a `::Rack::BodyProxy` to eventually log
the request once the response is deivered to the client. We did this
mainly in order to have a more "accurate" duration in the logs, similar
to what the `::Rack::CommonLogger` middleware does.

This has some significant disadvantages though:

* If an exception is raised by an outer middleware after we returned,
the response is never delivered and the log is lost.
* If a timeout occured (and e.g. a Unicorn worker is killed or a
`Rack::Timeout` stroke) before we return the full response, the log is
lost.
* If the client goes away before the request is finished, the app server
might decide no just throw away the response.

Besides addressing these issues, logging directly after the request
leaves our middleware scope makes the whole process much easier to
reason about.

lib/rackstash/rack/middleware.rb

@@ -203,28 +203,23 @@ module Rackstash
         env['rack.logger'.freeze] = @logger
         env['rack.errors'.freeze] = Rackstash::Rack::Errors.new(@logger)
 
-        @logger.push_buffer(buffering: @buffering, allow_silent: true)
         begin
-          @logger.timestamp
-          on_request(env)
+          @logger.push_buffer(buffering: @buffering, allow_silent: true)
+          begin
+            @logger.timestamp
+            on_request(env)
 
-          response = @app.call(env)
-        rescue Exception => exception
-          buffer = @logger.pop_buffer
-          if buffer
-            begin
-              on_error(buffer, env, began_at, exception)
-            ensure
-              buffer.flush
+            @app.call(env).tap do |response|
+              on_response(env, began_at, response)
             end
+          rescue Exception => exception
+            on_error(env, began_at, exception)
+            raise
           end
-          raise
+        ensure
+          buffer = @logger.pop_buffer
+          buffer.flush if buffer
         end
-
-        buffer = @logger.pop_buffer
-        on_response(buffer, env, began_at, response) if buffer
-
-        response
       end
 
       private
@@ -250,8 +245,6 @@ module Rackstash
         @logger.tag(@request_tags, scope: request) unless @request_tags.nil?
       end
 
-      # @param buffer [Rackstash::Buffer] the log {Buffer} which was captured
-      #   the logs during the request in the lower layers.
       # @param env [Hash] the Rack environment
       # @param began_at [Float] a timestamp denoting the start of the request.
       #   It can be used to get a request duration by subtracting it from
@@ -261,55 +254,36 @@ module Rackstash
       #   the response `headers`, and the response `body`.
       #
       # @return [void]
-      def on_response(buffer, env, began_at, response)
+      def on_response(env, began_at, response)
         response[1] = ::Rack::Utils::HeaderHash.new(response[1])
-        response[2] = ::Rack::BodyProxy.new(response[2]) do
-          begin
-            log_response(buffer, env, began_at, response[0], response[1])
-          rescue Exception => exception
-            on_error(buffer, env, began_at, exception)
-            raise
-          ensure
-            buffer.flush
-          end
-        end
+        log_response(env, began_at, response[0], response[1])
 
         response
       end
 
-      # Log data from the response to the buffer for the request.
+      # Log data from the response to the log buffer of the request.
       #
-      # Note that we are directly writing to the `buffer` here, not the
-      # `@logger`. This is necessary because `log_response` is called very late
-      # during response processing by the application server where the buffer is
-      # already poped from the stack.
-      #
-      # @param buffer (see #on_response)
       # @param env (see #on_response)
       # @param began_at (see #on_response)
       # @param status [Integer] the numeric response status
       # @param headers [::Rack::Utils::HeaderHash] the response headers
       #
       # @return [void]
-      def log_response(buffer, env, began_at, status, headers)
+      def log_response(env, began_at, status, headers)
         fields = {
           FIELD_STATUS => ::Rack::Utils.status_code(status),
           FIELD_DURATION => (clock_time - began_at).round(ISO8601_PRECISION)
         }
-        buffer.fields.merge!(fields, force: false)
+        @logger.fields.merge!(fields, force: false)
+        @logger.fields.deep_merge!(
+          @response_fields,
+          scope: headers,
+          force: false
+        ) unless @response_fields.nil?
 
-        unless @response_fields.nil?
-          buffer.fields.deep_merge!(
-            @response_fields,
-            scope: headers,
-            force: false
-          )
-        end
-        buffer.tag(@response_tags, scope: headers) unless @response_tags.nil?
+        @logger.tag(@response_tags, scope: headers) unless @response_tags.nil?
       end
 
-      # @param buffer [Rackstash::Buffer] the log {Buffer} which was captured
-      #   the logs during the request in the lower layers.
       # @param env [Hash] The Rack environment
       # @param began_at [Float] a timestamp denoting the start of the request.
       #   It can be used to get a request duration by subtracting it from
@@ -317,19 +291,19 @@ module Rackstash
       # @param exception [Exception] a rescued exception
       #
       # @return [void]
-      def on_error(buffer, env, began_at, exception)
-        log_error(buffer, env, began_at, exception)
+      def on_error(env, began_at, exception)
+        log_error(env, began_at, exception)
       end
 
       # (see #on_error)
-      def log_error(buffer, env, began_at, exception)
-        buffer.add_exception(exception, force: false)
+      def log_error(env, began_at, exception)
+        @logger.add_exception(exception, force: false)
 
         # Always set the status to 500, even if the app returned a successful
         # status. This is necessary to reflect the true status upsteam in case
         # there was an exception on adding request fields
-        buffer.fields[FIELD_STATUS] = 500
-        buffer.fields.set(FIELD_DURATION, force: false) {
+        @logger.fields[FIELD_STATUS] = 500
+        @logger.fields.set(FIELD_DURATION, force: false) {
           (clock_time - began_at).round(ISO8601_PRECISION)
         }
       end

spec/rackstash/rack/middleware_spec.rb

@@ -31,7 +31,8 @@ describe Rackstash::Rack::Middleware do
   let(:logger) { Rackstash::Logger.new ->(event) { log << event } }
 
   let(:args) { {} }
-  let(:stack) { described_class.new(app, logger, **args) }
+  let(:middleware) { described_class.new(app, logger, **args) }
+  let(:stack) { middleware }
 
   def get(path, opts = {})
     ::Rack::MockRequest.new(::Rack::Lint.new(stack)).get(path, opts)
@@ -182,9 +183,25 @@ describe Rackstash::Rack::Middleware do
 
       expect(log.last).to include('tags' => ['foo', 'text'])
     end
+
+    context 'on merge errors' do
+      it 'raises an error' do
+        args[:response_tags] = ->(_headers){ raise 'Oh No!' }
+
+        expect { get('/stuff') }.to raise_error RuntimeError, 'Oh No!'
+        expect(log.last).to include(
+          'path' => '/stuff',
+          'method' => 'GET',
+          'status' => 500,
+          'error' => 'RuntimeError',
+          'error_message' => 'Oh No!',
+          'error_trace' => %r{\A#{__FILE__}:#{__LINE__ - 9}:in}
+        )
+      end
+    end
   end
 
-  describe 'on errors' do
+  describe 'on application errors' do
     it 'logs errors' do
       expect { get('/error', params: { raise: 'Oh noes!' }) }
         .to raise_error(RuntimeError, 'Oh noes!')
@@ -256,4 +273,24 @@ describe Rackstash::Rack::Middleware do
       expect(log.last).to_not include 'foo'
     end
   end
+
+  describe 'on late middleware errors' do
+    let(:failing_middleware) {
+      lambda do |env|
+        middleware.call(env)
+        raise 'NEIN!'
+      end
+    }
+    let(:stack) { failing_middleware }
+
+    it 'logs even if a late middleware raises' do
+      expect { get('/foo') }.to raise_error(RuntimeError, 'NEIN!')
+
+      expect(log.last).to include(
+        'path' => '/foo',
+        'message' => "Request started\nNothing to do...\n"
+      )
+      expect(log.last).to_not include('error', 'error_message', 'error_trace')
+    end
+  end
 end