Improve Rails/Rack request instrumentation

Rework the Rails instrumentation to use the Rack::Events module to start
and complete the transaction at the very start and end of the request.
A new middleware is added at the very top of the middleware stack to
start and complete the transaction.

This allows us to instrument more of the request's duration, as we're
now at the very start and end of the request handling. We will now
report the time the request has spend handling any other middleware.
This will be visible with a new `process_request.rack` event.

To make sure we always close the transaction, listen for the `on_finish`
event, but also register a last resort "after_reply" callback. This
callback is supported by web servers like Puma and Unicorn. If it's
called when a transaction is active, it will close it, preventing
transactions to span multiple requests. If no transaction is active, it
will do nothing.

Moving forward, I'd like to update our Rack, Sinatra, Padrino, Grape,
etc. instrumentations to use this EventHandler middleware too, which
will be a bit more work, because some require users to manually add
middleware. When they all use the same EventHandler, we have one place
to start and complete transactions, and issues like
#289 cannot happen
again in the future.

Previously, our only Rails instrumentation was positioned after the
DebugExceptions middleware, just in time to catch any errors before the
Rails error middleware showed a nice error/debug page. This wouldn't
show the complete picture in request duration.

The original Rails instrumentation middleware will remain to report
errors and Rails specific metadata. This cannot be part of the
EventHandler, without overloading it with responsibilities, and for
Rails error reporting we need to be right after the DebugExceptions
middleware before errors are no longer bubbled up in the stack.

Instead of the transaction ID being the request ID, set a new
`request_id` tag to not break existing behavior.

Part of #329

.changesets/report-rails-request-middleware-in-request-duration.md

@@ -0,0 +1,6 @@
+---
+bump: minor
+type: change
+---
+
+Report the time spent in Rails middleware as part of the request duration. The AppSignal Rack middleware is now higher in the middleware stack and reports more time of the request to give insights in how long other middleware took. This is reported under the new `process_request.rack` event in the event timeline.

lib/appsignal.rb

@@ -325,5 +325,6 @@ def const_missing(name)
 require "appsignal/transaction"
 require "appsignal/version"
 require "appsignal/rack/generic_instrumentation"
+require "appsignal/rack/event_handler"
 require "appsignal/transmitter"
 require "appsignal/heartbeat"

lib/appsignal/integrations/railtie.rb

@@ -29,6 +29,11 @@ def self.initialize_appsignal(app)
         # Start logger
         Appsignal.start_logger
 
+        app.middleware.insert(
+          0,
+          ::Rack::Events,
+          [Appsignal::Rack::EventHandler.new]
+        )
         app.middleware.insert_after(
           ActionDispatch::DebugExceptions,
           Appsignal::Rack::RailsInstrumentation

lib/appsignal/rack/event_handler.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module Appsignal
+  module Rack
+    APPSIGNAL_TRANSACTION = "appsignal.transaction"
+    RACK_AFTER_REPLY = "rack.after_reply"
+
+    class EventHandler
+      include ::Rack::Events::Abstract
+
+      def self.safe_execution(name)
+        yield
+      rescue => e
+        Appsignal.internal_logger.error(
+          "Error occurred in #{name}: #{e.class}: #{e}: #{e.backtrace}"
+        )
+      end
+
+      def on_start(request, _response)
+        self.class.safe_execution("Appsignal::Rack::EventHandler#on_start") do
+          transaction = Appsignal::Transaction.create(
+            SecureRandom.uuid,
+            Appsignal::Transaction::HTTP_REQUEST,
+            request
+          )
+          request.env[APPSIGNAL_TRANSACTION] = transaction
+
+          request.env[RACK_AFTER_REPLY] ||= []
+          request.env[RACK_AFTER_REPLY] << proc do
+            Appsignal::Rack::EventHandler
+              .safe_execution("Appsignal::Rack::EventHandler's after_reply") do
+              transaction.finish_event("process_request.rack", "", "")
+              transaction.set_http_or_background_queue_start
+
+              # Make sure the current transaction is always closed when the request
+              # is finished. This is a fallback for in case the `on_finish`
+              # callback is not called. This is supported by servers like Puma and
+              # Unicorn.
+              #
+              # The EventHandler.on_finish callback should be called first, this is
+              # just a fallback if that doesn't get called.
+              Appsignal::Transaction.complete_current!
+            end
+          end
+          transaction.start_event
+        end
+      end
+
+      def on_error(request, _response, error)
+        self.class.safe_execution("Appsignal::Rack::EventHandler#on_error") do
+          transaction = request.env[APPSIGNAL_TRANSACTION]
+          return unless transaction
+
+          transaction.set_error(error)
+        end
+      end
+
+      def on_finish(request, _response)
+        self.class.safe_execution("Appsignal::Rack::EventHandler#on_finish") do
+          transaction = request.env[APPSIGNAL_TRANSACTION]
+          return unless transaction
+
+          transaction.finish_event("process_request.rack", "", "")
+          transaction.set_action_if_nil("#{request.request_method} #{request.path}")
+          transaction.set_http_or_background_queue_start
+
+          # Make sure the current transaction is always closed when the request
+          # is finished
+          Appsignal::Transaction.complete_current!
+        end
+      end
+
+      private
+
+      def format_namespace(namespace)
+        if namespace == Appsignal::Transaction::HTTP_REQUEST
+          :web
+        else
+          namespace
+        end
+      end
+    end
+  end
+end

lib/appsignal/rack/rails_instrumentation.rb

@@ -22,13 +22,16 @@ def call(env)
 
       def call_with_appsignal_monitoring(env)
         request = ActionDispatch::Request.new(env)
-        transaction = Appsignal::Transaction.create(
-          request_id(env),
-          Appsignal::Transaction::HTTP_REQUEST,
-          request,
-          :params_method => :filtered_parameters
+        transaction = env.fetch(
+          Appsignal::Rack::APPSIGNAL_TRANSACTION,
+          Appsignal::Transaction::NilTransaction.new
         )
+
         begin
+          transaction.params = fetch_params(request)
+          request_id = fetch_request_id(env)
+          transaction.set_tags(:request_id => request_id) if request_id
+
           @app.call(env)
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
@@ -38,19 +41,29 @@ def call_with_appsignal_monitoring(env)
           if controller
             transaction.set_action_if_nil("#{controller.class}##{controller.action_name}")
           end
-          transaction.set_http_or_background_queue_start
           transaction.set_metadata("path", request.path)
-          begin
-            transaction.set_metadata("method", request.request_method)
-          rescue => error
-            Appsignal.internal_logger.error("Unable to report HTTP request method: '#{error}'")
-          end
-          Appsignal::Transaction.complete_current!
+          request_method = fetch_request_method(request)
+          transaction.set_metadata("method", request_method) if request_method
         end
       end
 
-      def request_id(env)
-        env["action_dispatch.request_id"] || SecureRandom.uuid
+      def fetch_request_id(env)
+        env["action_dispatch.request_id"]
+      end
+
+      def fetch_params(request)
+        request.filtered_parameters
+      rescue => error
+        # Getting params from the request has been know to fail.
+        Appsignal.internal_logger.debug "Exception while getting params: #{error}"
+        nil
+      end
+
+      def fetch_request_method(request)
+        request.request_method
+      rescue => error
+        Appsignal.internal_logger.error("Unable to report HTTP request method: '#{error}'")
+        nil
       end
     end
   end

spec/lib/appsignal/rack/event_handler_spec.rb

@@ -0,0 +1,191 @@
+describe Appsignal::Rack::EventHandler do
+  let(:queue_start_time) { fixed_time * 1_000 }
+  let(:env) do
+    {
+      "HTTP_X_REQUEST_START" => "t=#{queue_start_time.to_i}", # in milliseconds
+      "REQUEST_METHOD" => "GET",
+      "PATH_INFO" => "/path"
+    }
+  end
+  let(:request) { Rack::Request.new(env) }
+  let(:response) { nil }
+  let(:log_stream) { StringIO.new }
+  let(:log) { log_contents(log_stream) }
+  before do
+    start_agent
+    Appsignal.internal_logger = test_logger(log_stream)
+  end
+  around { |example| keep_transactions { example.run } }
+
+  def on_start
+    described_class.new.on_start(request, response)
+  end
+
+  describe "#on_start" do
+    it "creates a new transaction" do
+      expect { on_start }.to change { created_transactions.length }.by(1)
+
+      transaction = last_transaction
+      expect(transaction.to_h).to include(
+        "id" => kind_of(String),
+        "namespace" => Appsignal::Transaction::HTTP_REQUEST
+      )
+
+      expect(Appsignal::Transaction.current).to eq(last_transaction)
+    end
+
+    it "registers transaction on the request environment" do
+      on_start
+
+      expect(request.env[Appsignal::Rack::APPSIGNAL_TRANSACTION])
+        .to eq(last_transaction)
+    end
+
+    it "registers an rack.after_reply callback that completes the transaction" do
+      request.env[Appsignal::Rack::RACK_AFTER_REPLY] = []
+      expect do
+        on_start
+      end.to change { request.env[Appsignal::Rack::RACK_AFTER_REPLY].length }.by(1)
+
+      expect(Appsignal::Transaction.current).to eq(last_transaction)
+
+      callback = request.env[Appsignal::Rack::RACK_AFTER_REPLY].first
+      callback.call
+
+      expect(Appsignal::Transaction.current).to be_kind_of(Appsignal::Transaction::NilTransaction)
+
+      expect(last_transaction.ext.queue_start).to eq(queue_start_time)
+    end
+
+    it "logs errors from rack.after_reply callbacks" do
+      on_start
+
+      expect(request.env[Appsignal::Rack::APPSIGNAL_TRANSACTION])
+        .to receive(:finish_event)
+        .and_raise(ExampleStandardError, "oh no")
+      callback = request.env[Appsignal::Rack::RACK_AFTER_REPLY].first
+      callback.call
+
+      expect(log).to contains_log(
+        :error,
+        "Error occurred in Appsignal::Rack::EventHandler's after_reply: ExampleStandardError: oh no"
+      )
+    end
+
+    it "logs an error in case of an error" do
+      expect(Appsignal::Transaction)
+        .to receive(:create).and_raise(ExampleStandardError, "oh no")
+
+      on_start
+
+      expect(log).to contains_log(
+        :error,
+        "Error occurred in Appsignal::Rack::EventHandler#on_start: ExampleStandardError: oh no"
+      )
+    end
+  end
+
+  describe "#on_error" do
+    def on_error(error)
+      described_class.new.on_error(request, response, error)
+    end
+
+    it "reports the error" do
+      on_start
+      on_error(ExampleStandardError.new("the error"))
+
+      expect(last_transaction.to_h).to include(
+        "error" => {
+          "name" => "ExampleStandardError",
+          "message" => "the error",
+          "backtrace" => kind_of(String)
+        }
+      )
+    end
+
+    it "logs an error in case of an internal error" do
+      on_start
+
+      expect(request.env[Appsignal::Rack::APPSIGNAL_TRANSACTION])
+        .to receive(:set_error).and_raise(ExampleStandardError, "oh no")
+
+      on_error(ExampleStandardError.new("the error"))
+
+      expect(log).to contains_log(
+        :error,
+        "Error occurred in Appsignal::Rack::EventHandler#on_error: ExampleStandardError: oh no"
+      )
+    end
+  end
+
+  describe "#on_finish" do
+    let(:response) { Rack::Events::BufferedResponse.new(200, {}, ["body"]) }
+
+    def on_finish
+      described_class.new.on_finish(request, response)
+    end
+
+    it "doesn't do anything without a transaction" do
+      on_start
+
+      request.env[Appsignal::Rack::APPSIGNAL_TRANSACTION] = nil
+
+      on_finish
+
+      expect(last_transaction.to_h).to include(
+        "action" => nil,
+        "sample_data" => {},
+        "events" => []
+      )
+    end
+
+    it "completes the transaction" do
+      on_start
+      on_finish
+
+      expect(last_transaction.to_h).to include(
+        "action" => "GET /path"
+      )
+      expect(last_transaction.ext.queue_start).to eq(queue_start_time)
+    end
+
+    it "doesn't set the action name if already set" do
+      on_start
+      last_transaction.set_action("My action")
+      on_finish
+
+      expect(last_transaction.to_h).to include(
+        "action" => "My action"
+      )
+    end
+
+    it "finishes the process_request.rack event" do
+      on_start
+      on_finish
+
+      expect(last_transaction.to_h).to include(
+        "events" => [
+          hash_including(
+            "name" => "process_request.rack",
+            "title" => "",
+            "body" => "",
+            "body_format" => Appsignal::EventFormatter::DEFAULT
+          )
+        ]
+      )
+    end
+
+    it "logs an error in case of an error" do
+      expect(Appsignal::Transaction)
+        .to receive(:complete_current!).and_raise(ExampleStandardError, "oh no")
+
+      on_start
+      on_finish
+
+      expect(log).to contains_log(
+        :error,
+        "Error occurred in Appsignal::Rack::EventHandler#on_finish: ExampleStandardError: oh no"
+      )
+    end
+  end
+end