Instrument Rack response bodies

Instrument what happens when a Rack response body is read and closed. We
instrument these events by wrapping the response body in the appropriate
response body BodyWrapper subclass, depending on the response object
type.

This change was previously sent in as PR #1037 and reverted in #1052. We
saw issues with multiple requests reported in the same transaction. This
problem occurred when there was an error in the middleware stack, and
the BodyWrapper never closed the response body.

I've removed any transaction complete logic from the BodyWrapper in the
original PR. We now have a Rack EventHandler that ensures the
transaction is always closed per request, even when such an error
scenario occurs again.

The only scenario in which we don't support this response body
instrumentation is when no EventHandler is present in the middleware
stack. This level of support is acceptable to me. We want people to use
the EventHandler. Most of our automatic instrumentations are updated to
add it to the middleware stack.

From the original commit: 7da96e7

> Some work might be getting done within a Rack response body. For
> example, when ActionController::Streaming is used, or when a Rack app
> elects to stream a response.
>
> The Rack SPEC doc defines the behavior in sufficient detail to wrap
> this into the current AppSignal transaction.
>
> Sadly, there is some work involved in supporting all the right
> methods, so just "one-size-fits-all" wrapper will not quite work.

Co-authored-by: Julik Tarkhanov <me@julik.nl>

.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
+```

lib/appsignal.rb

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

lib/appsignal/rack/abstract_middleware.rb

@@ -53,7 +53,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)
@@ -76,24 +76,30 @@ def call(env)
       # the 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_span_name
           Appsignal.instrument(@instrument_span_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)
+        # Instrument response body and closing of the response body
+        [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
+      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