Monitor serving of Rack response bodies

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 actually defines the behavior in sufficient detail to wrap this
into the same 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

lib/appsignal.rb

@@ -305,5 +305,6 @@ def collect_environment_metadata
 require "appsignal/integrations/railtie" if defined?(::Rails)
 require "appsignal/transaction"
 require "appsignal/version"
+require "appsignal/rack/body_wrapper"
 require "appsignal/rack/generic_instrumentation"
 require "appsignal/transmitter"

lib/appsignal/rack/body_wrapper.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+module Appsignal
+  # @api private
+  module Rack
+    class BodyWrapper
+      def self.wrap(original_body, appsignal_transaction_or_nil)
+        # 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 can be done using delegation via i.e. SimpleDelegate
+        # but we also need "close" to get called correctly so that the Appsignal transaction
+        # gets completed - which will not happen, for example, when #to_ary gets called
+        # just on the delegated Rack 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_or_nil)
+        elsif original_body.respond_to?(:to_ary)
+          ArrayableBodyWrapper.new(original_body, appsignal_transaction_or_nil)
+        elsif !original_body.respond_to?(:each) && original_body.respond_to?(:call)
+          CallableBodyWrapper.new(original_body, appsignal_transaction_or_nil)
+        else
+          EnumerableBodyWrapper.new(original_body, appsignal_transaction_or_nil)
+        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
+        @body.close if !@body_already_closed && @body.respond_to?(:close)
+        @body_already_closed = true
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction&.set_error(error)
+        raise error
+      ensure
+        @transaction&.complete
+      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?
+
+        @body.each(&blk)
+      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
+        @body.call(stream)
+      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
+        @body.to_ary
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction&.set_error(error)
+        raise error
+      ensure
+        close
+      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
+        @body.to_path
+      rescue Exception => error # rubocop:disable Lint/RescueException
+        @transaction&.set_error(error)
+        raise error
+      end
+    end
+  end
+end

lib/appsignal/rack/generic_instrumentation.rb

@@ -16,7 +16,8 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          @app.call(env)
+          status, headers, obody = @app.call(env)
+          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, _transaction = nil)]
         end
       end
 
@@ -29,17 +30,19 @@ def call_with_appsignal_monitoring(env)
         )
         begin
           Appsignal.instrument("process_action.generic") do
-            @app.call(env)
+            status, headers, obody = @app.call(env)
+            [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
           end
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
           raise error
         ensure
-          transaction.set_action_if_nil(env["appsignal.route"] || "unknown")
+          default_action = env["appsignal.route"] || env["appsignal.action"] || "unknown"
+          transaction.set_action_if_nil(default_action)
           transaction.set_metadata("path", request.path)
           transaction.set_metadata("method", request.request_method)
           transaction.set_http_or_background_queue_start
-          Appsignal::Transaction.complete_current!
+          # Transaction gets completed when the body gets read out
         end
       end
     end

lib/appsignal/rack/rails_instrumentation.rb

@@ -16,7 +16,8 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          @app.call(env)
+          status, headers, obody = @app.call(env)
+          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, _transaction = nil)]
         end
       end
 
@@ -29,7 +30,8 @@ def call_with_appsignal_monitoring(env)
           :params_method => :filtered_parameters
         )
         begin
-          @app.call(env)
+          status, headers, obody = @app.call(env)
+          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
           raise error
@@ -45,7 +47,7 @@ def call_with_appsignal_monitoring(env)
           rescue => error
             Appsignal.internal_logger.error("Unable to report HTTP request method: '#{error}'")
           end
-          Appsignal::Transaction.complete_current!
+          # Transaction gets completed when the body gets read out
         end
       end
 

lib/appsignal/rack/streaming_listener.rb

@@ -16,7 +16,8 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          @app.call(env)
+          status, headers, obody = @app.call(env)
+          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, _transaction = nil)]
         end
       end
 
@@ -43,31 +44,10 @@ def call_with_appsignal_monitoring(env)
           end
 
         # Wrap the result body with our StreamWrapper
-        [status, headers, StreamWrapper.new(body, transaction)]
+        [status, headers, Appsignal::Rack::BodyWrapper.wrap(body, transaction)]
       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
+  StreamWrapper = Rack::EnumerableBodyWrapper
 end