Instrument Rack response bodies

.changesets/add-instrumentation-for-streaming-rack-responses.md

@@ -0,0 +1,21 @@
+---
+bump: "minor"
+type: "add"
+---
+
+Add instrumentation to Rack responses, including streaming responses. New `process_response_body.rack` and `close_response_body.rack` events will be shown in the event timeline. These events show how long it takes to complete responses, depending on the response implementation, and when the response is closed.
+
+This Sinatra route with a streaming response will be better instrumented, for example:
+
+```ruby
+get "/stream" do
+  stream do |out|
+    sleep 1
+    out << "1"
+    sleep 1
+    out << "2"
+    sleep 1
+    out << "3"
+  end
+end
+```

.changesets/deprecate-appsignal--rack--streaminglistener-middleware.md

@@ -0,0 +1,7 @@
+---
+bump: patch
+type: deprecate
+---
+
+Deprecate the `Appsignal::Rack::StreamingListener` middleware. Use the `Appsignal::Rack::InstrumentationMiddleware` middleware instead.
+

lib/appsignal.rb

@@ -325,6 +325,7 @@ def const_missing(name)
 require "appsignal/marker"
 require "appsignal/garbage_collection"
 require "appsignal/rack"
+require "appsignal/rack/body_wrapper"
 require "appsignal/rack/abstract_middleware"
 require "appsignal/rack/instrumentation_middleware"
 require "appsignal/rack/event_handler"

lib/appsignal/rack/abstract_middleware.rb

@@ -58,7 +58,7 @@ def call(env)
                 wrapped_instrumentation
               )
             else
-              instrument_app_call(request.env)
+              instrument_app_call(request.env, transaction)
             end
           ensure
             add_transaction_metadata_after(transaction, request)
@@ -81,24 +81,36 @@ def call(env)
       # stack and will report the exception and complete the transaction.
       #
       # @see {#instrument_app_call_with_exception_handling}
-      def instrument_app_call(env)
+      def instrument_app_call(env, transaction)
         if @instrument_event_name
           Appsignal.instrument(@instrument_event_name) do
-            @app.call(env)
+            call_app(env, transaction)
           end
         else
-          @app.call(env)
+          call_app(env, transaction)
         end
       end
 
+      def call_app(env, transaction)
+        status, headers, obody = @app.call(env)
+        body =
+          if obody.is_a? Appsignal::Rack::BodyWrapper
+            obody
+          else
+            # Instrument response body and closing of the response body
+            Appsignal::Rack::BodyWrapper.wrap(obody, transaction)
+          end
+        [status, headers, body]
+      end
+
       # Instrument the request fully. This is used by the top instrumentation
       # middleware in the middleware stack. Unlike
       # {#instrument_app_call} this will report any exceptions being
       # raised.
       #
       # @see {#instrument_app_call}
       def instrument_app_call_with_exception_handling(env, transaction, wrapped_instrumentation)
-        instrument_app_call(env)
+        instrument_app_call(env, transaction)
       rescue Exception => error # rubocop:disable Lint/RescueException
         report_errors =
           if @report_errors == DEFAULT_ERROR_REPORTING

lib/appsignal/rack/body_wrapper.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+module Appsignal
+  module Rack
+    # @api private
+    class BodyWrapper
+      def self.wrap(original_body, appsignal_transaction)
+        # The logic of how Rack treats a response body differs based on which methods
+        # the body responds to. This means that to support the Rack 3.x spec in full
+        # we need to return a wrapper which matches the API of the wrapped body as closely
+        # as possible. Pick the wrapper from the most specific to the least specific.
+        # See https://github.com/rack/rack/blob/main/SPEC.rdoc#the-body-
+        #
+        # What is important is that our Body wrapper responds to the same methods Rack
+        # (or a webserver) would be checking and calling, and passes through that functionality
+        # to the original body.
+        #
+        # This comment https://github.com/rails/rails/pull/49627#issuecomment-1769802573
+        # is of particular interest to understand why this has to be somewhat complicated.
+        if original_body.respond_to?(:to_path)
+          PathableBodyWrapper.new(original_body, appsignal_transaction)
+        elsif original_body.respond_to?(:to_ary)
+          ArrayableBodyWrapper.new(original_body, appsignal_transaction)
+        elsif !original_body.respond_to?(:each) && original_body.respond_to?(:call)
+          # This body only supports #call, so we must be running a Rack 3 application
+          # It is possible that a body exposes both `each` and `call` in the hopes of
+          # being backwards-compatible with both Rack 3.x and Rack 2.x, however
+          # this is not going to work since the SPEC says that if both are available,
+          # `each` should be used and `call` should be ignored.
+          # So for that case we can drop to our default EnumerableBodyWrapper
+          CallableBodyWrapper.new(original_body, appsignal_transaction)
+        else
+          EnumerableBodyWrapper.new(original_body, appsignal_transaction)
+        end
+      end
+
+      def initialize(body, appsignal_transaction)
+        @body_already_closed = false
+        @body = body
+        @transaction = appsignal_transaction
+      end
+
+      # This must be present in all Rack bodies and will be called by the serving adapter
+      def close
+        # The @body_already_closed check is needed so that if `to_ary`
+        # of the body has already closed itself (as prescribed) we do not
+        # attempt to close it twice
+        if !@body_already_closed && @body.respond_to?(:close)
+          Appsignal.instrument("close_response_body.rack") { @body.close }
+        end
+        @body_already_closed = true
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction.set_error(error)
+        raise error
+      end
+    end
+
+    # The standard Rack body wrapper which exposes "each" for iterating
+    # over the response body. This is supported across all 3 major Rack
+    # versions.
+    #
+    # @api private
+    class EnumerableBodyWrapper < BodyWrapper
+      def each(&blk)
+        # This is a workaround for the Rails bug when there was a bit too much
+        # eagerness in implementing to_ary, see:
+        # https://github.com/rails/rails/pull/44953
+        # https://github.com/rails/rails/pull/47092
+        # https://github.com/rails/rails/pull/49627
+        # https://github.com/rails/rails/issues/49588
+        # While the Rack SPEC does not mandate `each` to be callable
+        # in a blockless way it is still a good idea to have it in place.
+        return enum_for(:each) unless block_given?
+
+        Appsignal.instrument("process_response_body.rack", "Process Rack response body (#each)") do
+          @body.each(&blk)
+        end
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction.set_error(error)
+        raise error
+      end
+    end
+
+    # The callable response bodies are a new Rack 3.x feature, and would not work
+    # with older Rack versions. They must not respond to `each` because
+    # "If it responds to each, you must call each and not call". This is why
+    # it inherits from BodyWrapper directly and not from EnumerableBodyWrapper
+    #
+    # @api private
+    class CallableBodyWrapper < BodyWrapper
+      def call(stream)
+        # `stream` will be closed by the app we are calling, no need for us
+        # to close it ourselves
+        Appsignal.instrument("process_response_body.rack", "Process Rack response body (#call)") do
+          @body.call(stream)
+        end
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction.set_error(error)
+        raise error
+      end
+    end
+
+    # "to_ary" takes precedence over "each" and allows the response body
+    # to be read eagerly. If the body supports that method, it takes precedence
+    # over "each":
+    # "Middleware may call to_ary directly on the Body and return a new Body in its place"
+    # One could "fold" both the to_ary API and the each() API into one Body object, but
+    # to_ary must also call "close" after it executes - and in the Rails implementation
+    # this pecularity was not handled properly.
+    #
+    # @api private
+    class ArrayableBodyWrapper < EnumerableBodyWrapper
+      def to_ary
+        @body_already_closed = true
+        Appsignal.instrument(
+          "process_response_body.rack",
+          "Process Rack response body (#to_ary)"
+        ) do
+          @body.to_ary
+        end
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction.set_error(error)
+        raise error
+      end
+    end
+
+    # Having "to_path" on a body allows Rack to serve out a static file, or to
+    # pass that file to the downstream webserver for sending using X-Sendfile
+    class PathableBodyWrapper < EnumerableBodyWrapper
+      def to_path
+        Appsignal.instrument(
+          "process_response_body.rack",
+          "Process Rack response body (#to_path)"
+        ) do
+          @body.to_path
+        end
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction.set_error(error)
+        raise error
+      end
+    end
+  end
+end

lib/appsignal/rack/streaming_listener.rb

@@ -1,73 +1,27 @@
 # frozen_string_literal: true
 
+Appsignal::Utils::StdoutAndLoggerMessage.warning \
+  "The constant Appsignal::Rack::StreamingListener has been deprecated. " \
+    "Please update the constant name to " \
+    "Appsignal::Rack::InstrumentationMiddleware."
+
 module Appsignal
   module Rack
-    # Appsignal module that tracks exceptions in Streaming rack responses.
+    # Instrumentation middleware that tracks exceptions in streaming Rack
+    # responses.
     #
     # @api private
-    class StreamingListener
+    class StreamingListener < AbstractMiddleware
       def initialize(app, options = {})
-        Appsignal.internal_logger.debug "Initializing Appsignal::Rack::StreamingListener"
-        @app = app
-        @options = options
-      end
-
-      def call(env)
-        if Appsignal.active?
-          call_with_appsignal_monitoring(env)
-        else
-          @app.call(env)
-        end
+        options[:instrument_event_name] ||= "process_streaming_request.rack"
+        super
       end
 
-      def call_with_appsignal_monitoring(env)
-        request = ::Rack::Request.new(env)
-        transaction = Appsignal::Transaction.create(
-          SecureRandom.uuid,
-          Appsignal::Transaction::HTTP_REQUEST,
-          request
-        )
+      def add_transaction_metadata_after(transaction, request)
+        transaction.set_action_if_nil(request.env["appsignal.action"])
 
-        # Instrument a `process_action`, to set params/action name
-        status, headers, body =
-          Appsignal.instrument("process_action.rack") do
-            @app.call(env)
-          rescue Exception => e # rubocop:disable Lint/RescueException
-            transaction.set_error(e)
-            raise e
-          ensure
-            transaction.set_action_if_nil(env["appsignal.action"])
-            transaction.set_metadata("path", request.path)
-            transaction.set_metadata("method", request.request_method)
-            transaction.set_http_or_background_queue_start
-          end
-
-        # Wrap the result body with our StreamWrapper
-        [status, headers, StreamWrapper.new(body, transaction)]
+        super
       end
     end
   end
-
-  class StreamWrapper
-    def initialize(stream, transaction)
-      @stream = stream
-      @transaction = transaction
-    end
-
-    def each(&block)
-      @stream.each(&block)
-    rescue Exception => e # rubocop:disable Lint/RescueException
-      @transaction.set_error(e)
-      raise e
-    end
-
-    def close
-      @stream.close if @stream.respond_to?(:close)
-    rescue Exception => e # rubocop:disable Lint/RescueException
-      @transaction.set_error(e)
-      raise e
-    ensure
-      Appsignal::Transaction.complete_current!
-    end
-  end
 end

spec/lib/appsignal/rack/abstract_middleware_spec.rb

@@ -45,6 +45,11 @@ def make_request_with_error(error_class, error_message)
         expect(last_transaction).to have_namespace(Appsignal::Transaction::HTTP_REQUEST)
       end
 
+      it "wraps the response body in a BodyWrapper subclass" do
+        _status, _headers, body = make_request
+        expect(body).to be_kind_of(Appsignal::Rack::BodyWrapper)
+      end
+
       context "without an error" do
         before { make_request }
 
@@ -326,8 +331,10 @@ def filtered_params
       end
 
       context "with parent instrumentation" do
+        let(:transaction) { http_request_transaction }
         before do
-          env[Appsignal::Rack::APPSIGNAL_TRANSACTION] = http_request_transaction
+          env[Appsignal::Rack::APPSIGNAL_TRANSACTION] = transaction
+          set_current_transaction(transaction)
         end
 
         it "uses the existing transaction" do
@@ -336,6 +343,35 @@ def filtered_params
           expect { make_request }.to_not(change { created_transactions.count })
         end
 
+        it "wraps the response body in a BodyWrapper subclass" do
+          _status, _headers, body = make_request
+          expect(body).to be_kind_of(Appsignal::Rack::BodyWrapper)
+
+          body.to_ary
+          response_events =
+            last_transaction.to_h["events"].count do |event|
+              event["name"] == "process_response_body.rack"
+            end
+          expect(response_events).to eq(1)
+        end
+
+        context "when response body is already a BodyWrapper subclass" do
+          let(:body) { Appsignal::Rack::BodyWrapper.wrap(["hello!"], transaction) }
+          let(:app) { DummyApp.new { [200, {}, body] } }
+
+          it "doesn't wrap the body again" do
+            _status, _headers, body = make_request
+            expect(body).to eq(body)
+
+            body.to_ary
+            response_events =
+              last_transaction.to_h["events"].count do |event|
+                event["name"] == "process_response_body.rack"
+              end
+            expect(response_events).to eq(1)
+          end
+        end
+
         context "with error" do
           let(:app) { lambda { |_env| raise ExampleException, "error message" } }