Revert "Monitor serving of Rack response bodies"

lib/appsignal.rb

@@ -305,6 +305,5 @@ 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/generic_instrumentation.rb

@@ -16,9 +16,7 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+          @app.call(env)
         end
       end
 
@@ -29,30 +27,19 @@ def call_with_appsignal_monitoring(env)
           Appsignal::Transaction::HTTP_REQUEST,
           request
         )
-        # We need to complete the transaction if there is an exception inside the `call`
-        # of the app. If there isn't one and the app returns us a Rack response triplet, we let
-        # the BodyWrapper complete the transaction when #close gets called on it
-        # (guaranteed by the webserver)
-        complete_transaction_without_body = false
         begin
           Appsignal.instrument("process_action.generic") do
-            status, headers, obody = @app.call(env)
-            [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
+            @app.call(env)
           end
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
-          complete_transaction_without_body = true
           raise error
         ensure
-          default_action = env["appsignal.route"] || env["appsignal.action"] || "unknown"
-          transaction.set_action_if_nil(default_action)
+          transaction.set_action_if_nil(env["appsignal.route"] || "unknown")
           transaction.set_metadata("path", request.path)
           transaction.set_metadata("method", request.request_method)
           transaction.set_http_or_background_queue_start
-
-          # Transaction gets completed when the body gets read out, except in cases when
-          # the app failed before returning us the Rack response triplet.
-          Appsignal::Transaction.complete_current! if complete_transaction_without_body
+          Appsignal::Transaction.complete_current!
         end
       end
     end

lib/appsignal/rack/rails_instrumentation.rb

@@ -16,9 +16,7 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+          @app.call(env)
         end
       end
 
@@ -30,17 +28,10 @@ def call_with_appsignal_monitoring(env)
           request,
           :params_method => :filtered_parameters
         )
-        # We need to complete the transaction if there is an exception exception inside the `call`
-        # of the app. If there isn't one and the app returns us a Rack response triplet, we let
-        # the BodyWrapper complete the transaction when #close gets called on it
-        # (guaranteed by the webserver)
-        complete_transaction_without_body = false
         begin
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
+          @app.call(env)
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
-          complete_transaction_without_body = true
           raise error
         ensure
           controller = env["action_controller.instance"]
@@ -54,10 +45,7 @@ def call_with_appsignal_monitoring(env)
           rescue => error
             Appsignal.internal_logger.error("Unable to report HTTP request method: '#{error}'")
           end
-
-          # Transaction gets completed when the body gets read out, except in cases when
-          # the app failed before returning us the Rack response triplet.
-          Appsignal::Transaction.complete_current! if complete_transaction_without_body
+          Appsignal::Transaction.complete_current!
         end
       end
 

lib/appsignal/rack/sinatra_instrumentation.rb

@@ -42,9 +42,7 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+          @app.call(env)
         end
       end
 
@@ -58,19 +56,12 @@ def call_with_appsignal_monitoring(env)
           request,
           options
         )
-        # We need to complete the transaction if there is an exception exception inside the `call`
-        # of the app. If there isn't one and the app returns us a Rack response triplet, we let
-        # the BodyWrapper complete the transaction when #close gets called on it
-        # (guaranteed by the webserver)
-        complete_transaction_without_body = false
         begin
           Appsignal.instrument("process_action.sinatra") do
-            status, headers, obody = @app.call(env)
-            [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
+            @app.call(env)
           end
         rescue Exception => error # rubocop:disable Lint/RescueException
           transaction.set_error(error)
-          complete_transaction_without_body = true
           raise error
         ensure
           # If raise_error is off versions of Sinatra don't raise errors, but store
@@ -82,10 +73,7 @@ def call_with_appsignal_monitoring(env)
           transaction.set_metadata("path", request.path)
           transaction.set_metadata("method", request.request_method)
           transaction.set_http_or_background_queue_start
-
-          # Transaction gets completed when the body gets read out, except in cases when
-          # the app failed before returning us the Rack response triplet.
-          Appsignal::Transaction.complete_current! if complete_transaction_without_body
+          Appsignal::Transaction.complete_current!
         end
       end
 

lib/appsignal/rack/streaming_listener.rb

@@ -16,9 +16,7 @@ def call(env)
         if Appsignal.active?
           call_with_appsignal_monitoring(env)
         else
-          nil_transaction = Appsignal::Transaction::NilTransaction.new
-          status, headers, obody = @app.call(env)
-          [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, nil_transaction)]
+          @app.call(env)
         end
       end
 
@@ -30,35 +28,46 @@ def call_with_appsignal_monitoring(env)
           request
         )
 
-        # We need to complete the transaction if there is an exception exception inside the `call`
-        # of the app. If there isn't one and the app returns us a Rack response triplet, we let
-        # the BodyWrapper complete the transaction when #close gets called on it
-        # (guaranteed by the webserver)
-        complete_transaction_without_body = false
-
         # Instrument a `process_action`, to set params/action name
-        begin
+        status, headers, body =
           Appsignal.instrument("process_action.rack") do
-            status, headers, obody = @app.call(env)
-            [status, headers, Appsignal::Rack::BodyWrapper.wrap(obody, transaction)]
+            @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
-        rescue Exception => error # rubocop:disable Lint/RescueException
-          transaction.set_error(error)
-          complete_transaction_without_body = true
-          raise error
-        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
 
-          # Transaction gets completed when the body gets read out, except in cases when
-          # the app failed before returning us the Rack response triplet.
-          Appsignal::Transaction.complete_current! if complete_transaction_without_body
-        end
+        # Wrap the result body with our StreamWrapper
+        [status, headers, StreamWrapper.new(body, transaction)]
       end
     end
   end
 
-  StreamWrapper = Rack::EnumerableBodyWrapper
+  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