From c1a7c2e2bfd21eec6e8cec8c869e3488a9306118 Mon Sep 17 00:00:00 2001
From: Tom de Bruijn <tom@tomdebruijn.com>
Date: Tue, 2 Jul 2024 15:45:20 +0200
Subject: [PATCH] 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: 7da96e79684f5af2f696238f6e3159de17980d6f

> 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>
---
 ...umentation-for-streaming-rack-responses.md |  21 ++
 lib/appsignal.rb                              |   1 +
 lib/appsignal/rack/abstract_middleware.rb     |  16 +-
 lib/appsignal/rack/body_wrapper.rb            | 143 ++++++++++
 spec/lib/appsignal/rack/body_wrapper_spec.rb  | 263 ++++++++++++++++++
 5 files changed, 439 insertions(+), 5 deletions(-)
 create mode 100644 .changesets/add-instrumentation-for-streaming-rack-responses.md
 create mode 100644 lib/appsignal/rack/body_wrapper.rb
 create mode 100644 spec/lib/appsignal/rack/body_wrapper_spec.rb

diff --git a/.changesets/add-instrumentation-for-streaming-rack-responses.md b/.changesets/add-instrumentation-for-streaming-rack-responses.md
new file mode 100644
index 000000000..8d562cd85
--- /dev/null
+++ b/.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
+```
diff --git a/lib/appsignal.rb b/lib/appsignal.rb
index 29b7a7faf..dde30175a 100644
--- a/lib/appsignal.rb
+++ b/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"
diff --git a/lib/appsignal/rack/abstract_middleware.rb b/lib/appsignal/rack/abstract_middleware.rb
index 802f34b65..6dd0bfb65 100644
--- a/lib/appsignal/rack/abstract_middleware.rb
+++ b/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,16 +76,22 @@ 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
@@ -93,7 +99,7 @@ def instrument_app_call(env)
       #
       # @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
diff --git a/lib/appsignal/rack/body_wrapper.rb b/lib/appsignal/rack/body_wrapper.rb
new file mode 100644
index 000000000..737b3bf7d
--- /dev/null
+++ b/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
diff --git a/spec/lib/appsignal/rack/body_wrapper_spec.rb b/spec/lib/appsignal/rack/body_wrapper_spec.rb
new file mode 100644
index 000000000..3ac346271
--- /dev/null
+++ b/spec/lib/appsignal/rack/body_wrapper_spec.rb
@@ -0,0 +1,263 @@
+describe Appsignal::Rack::BodyWrapper do
+  let(:transaction) { http_request_transaction }
+  before do
+    start_agent
+    set_current_transaction(transaction)
+  end
+
+  describe "with a body that supports all possible features" do
+    it "reduces the supported methods to just each()" do
+      # which is the safest thing to do, since the body is likely broken
+      fake_body = double(
+        :each => nil,
+        :call => nil,
+        :to_ary => [],
+        :to_path => "/tmp/foo.bin",
+        :close => nil
+      )
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to respond_to(:each)
+      expect(wrapped).to_not respond_to(:to_ary)
+      expect(wrapped).to_not respond_to(:call)
+      expect(wrapped).to respond_to(:close)
+    end
+  end
+
+  describe "with a body only supporting each()" do
+    it "wraps with appropriate class" do
+      fake_body = double(:each => nil)
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to respond_to(:each)
+      expect(wrapped).to_not respond_to(:to_ary)
+      expect(wrapped).to_not respond_to(:call)
+      expect(wrapped).to respond_to(:close)
+    end
+
+    it "reads out the body in full using each" do
+      fake_body = double
+      expect(fake_body).to receive(:each).once.and_yield("a").and_yield("b").and_yield("c")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect { |b| wrapped.each(&b) }.to yield_successive_args("a", "b", "c")
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#each)"
+      )
+    end
+
+    it "returns an Enumerator if each() gets called without a block" do
+      fake_body = double
+      expect(fake_body).to receive(:each).once.and_yield("a").and_yield("b").and_yield("c")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      enum = wrapped.each
+      expect(enum).to be_kind_of(Enumerator)
+      expect { |b| enum.each(&b) }.to yield_successive_args("a", "b", "c")
+
+      expect(transaction).to_not include_event("name" => "process_response_body.rack")
+    end
+
+    it "sets the exception raised inside each() on the transaction" do
+      fake_body = double
+      expect(fake_body).to receive(:each).once.and_raise(ExampleException, "error message")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect do
+        expect { |b| wrapped.each(&b) }.to yield_control
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+
+    it "closes the body and tracks an instrumentation event when it gets closed" do
+      fake_body = double(:close => nil)
+      expect(fake_body).to receive(:each).once.and_yield("a").and_yield("b").and_yield("c")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect { |b| wrapped.each(&b) }.to yield_successive_args("a", "b", "c")
+      wrapped.close
+
+      expect(transaction).to include_event("name" => "close_response_body.rack")
+    end
+  end
+
+  describe "with a body supporting both each() and call" do
+    it "wraps with the wrapper that conceals call() and exposes each" do
+      fake_body = double
+      allow(fake_body).to receive(:each)
+      allow(fake_body).to receive(:call)
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to respond_to(:each)
+      expect(wrapped).to_not respond_to(:to_ary)
+      expect(wrapped).to_not respond_to(:call)
+      expect(wrapped).to_not respond_to(:to_path)
+      expect(wrapped).to respond_to(:close)
+    end
+  end
+
+  describe "with a body supporting both to_ary and each" do
+    let(:fake_body) { double(:each => nil, :to_ary => []) }
+
+    it "wraps with appropriate class" do
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to respond_to(:each)
+      expect(wrapped).to respond_to(:to_ary)
+      expect(wrapped).to_not respond_to(:call)
+      expect(wrapped).to_not respond_to(:to_path)
+      expect(wrapped).to respond_to(:close)
+    end
+
+    it "reads out the body in full using each" do
+      expect(fake_body).to receive(:each).once.and_yield("a").and_yield("b").and_yield("c")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect { |b| wrapped.each(&b) }.to yield_successive_args("a", "b", "c")
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#each)"
+      )
+    end
+
+    it "sets the exception raised inside each() into the Appsignal transaction" do
+      expect(fake_body).to receive(:each).once.and_raise(ExampleException, "error message")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect do
+        expect { |b| wrapped.each(&b) }.to yield_control
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+
+    it "reads out the body in full using to_ary" do
+      expect(fake_body).to receive(:to_ary).and_return(["one", "two", "three"])
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped.to_ary).to eq(["one", "two", "three"])
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#to_ary)"
+      )
+    end
+
+    it "sends the exception raised inside to_ary() into the Appsignal and closes transaction" do
+      fake_body = double
+      allow(fake_body).to receive(:each)
+      expect(fake_body).to receive(:to_ary).once.and_raise(ExampleException, "error message")
+      expect(fake_body).to_not receive(:close) # Per spec we expect the body has closed itself
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect do
+        wrapped.to_ary
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+  end
+
+  describe "with a body supporting both to_path and each" do
+    let(:fake_body) { double(:each => nil, :to_path => nil) }
+
+    it "wraps with appropriate class" do
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to respond_to(:each)
+      expect(wrapped).to_not respond_to(:to_ary)
+      expect(wrapped).to_not respond_to(:call)
+      expect(wrapped).to respond_to(:to_path)
+      expect(wrapped).to respond_to(:close)
+    end
+
+    it "reads out the body in full using each()" do
+      expect(fake_body).to receive(:each).once.and_yield("a").and_yield("b").and_yield("c")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect { |b| wrapped.each(&b) }.to yield_successive_args("a", "b", "c")
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#each)"
+      )
+    end
+
+    it "sets the exception raised inside each() into the Appsignal transaction" do
+      expect(fake_body).to receive(:each).once.and_raise(ExampleException, "error message")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect do
+        expect { |b| wrapped.each(&b) }.to yield_control
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+
+    it "sets the exception raised inside to_path() into the Appsignal transaction" do
+      allow(fake_body).to receive(:to_path).once.and_raise(ExampleException, "error message")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect do
+        wrapped.to_path
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+
+    it "exposes to_path to the sender" do
+      allow(fake_body).to receive(:to_path).and_return("/tmp/file.bin")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped.to_path).to eq("/tmp/file.bin")
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#to_path)"
+      )
+    end
+  end
+
+  describe "with a body only supporting call()" do
+    let(:fake_body) { double(:call => nil) }
+
+    it "wraps with appropriate class" do
+      wrapped = described_class.wrap(fake_body, transaction)
+      expect(wrapped).to_not respond_to(:each)
+      expect(wrapped).to_not respond_to(:to_ary)
+      expect(wrapped).to respond_to(:call)
+      expect(wrapped).to_not respond_to(:to_path)
+      expect(wrapped).to respond_to(:close)
+    end
+
+    it "passes the stream into the call() of the body" do
+      fake_rack_stream = double("stream")
+      expect(fake_body).to receive(:call).with(fake_rack_stream)
+
+      wrapped = described_class.wrap(fake_body, transaction)
+      wrapped.call(fake_rack_stream)
+
+      expect(transaction).to include_event(
+        "name" => "process_response_body.rack",
+        "title" => "Process Rack response body (#call)"
+      )
+    end
+
+    it "sets the exception raised inside call() into the Appsignal transaction" do
+      fake_rack_stream = double
+      allow(fake_body).to receive(:call)
+        .with(fake_rack_stream)
+        .and_raise(ExampleException, "error message")
+
+      wrapped = described_class.wrap(fake_body, transaction)
+
+      expect do
+        wrapped.call(fake_rack_stream)
+      end.to raise_error(ExampleException, "error message")
+
+      expect(transaction).to have_error("ExampleException", "error message")
+    end
+  end
+end