Merge pull request rails#51508 from ioquatix/actiondispatch-response-body-to_ary

Don't buffer things that shouldn't be buffered.

Author: matthewd

actionpack/lib/action_controller/metal/live.rb

@@ -171,12 +171,6 @@ def initialize(response)
         @ignore_disconnect = false
       end
 
-      # ActionDispatch::Response delegates #to_ary to the internal
-      # ActionDispatch::Response::Buffer, defining #to_ary is an indicator that the
-      # response body can be buffered and/or cached by Rack middlewares, this is not
-      # the case for Live responses so we undefine it for this Buffer subclass.
-      undef_method :to_ary
-
       def write(string)
         unless @response.committed?
           @response.headers["Cache-Control"] ||= "no-cache"

actionpack/lib/action_dispatch/http/response.rb

@@ -105,10 +105,22 @@ def initialize(response, buf)
         @str_body = nil
       end
 
+      BODY_METHODS = { to_ary: true }
+
+      def respond_to?(method, include_private = false)
+        if BODY_METHODS.key?(method)
+          @buf.respond_to?(method)
+        else
+          super
+        end
+      end
+
       def to_ary
-        @buf.respond_to?(:to_ary) ?
-          @buf.to_ary :
-          @buf.each
+        if @str_body
+          [body]
+        else
+          @buf = @buf.to_ary
+        end
       end
 
       def body
@@ -328,7 +340,13 @@ def message
     # Returns the content of the response as a string. This contains the contents of
     # any calls to `render`.
     def body
-      @stream.body
+      if @stream.respond_to?(:to_ary)
+        @stream.to_ary.join
+      elsif @stream.respond_to?(:body)
+        @stream.body
+      else
+        @stream
+      end
     end
 
     def write(string)
@@ -337,11 +355,16 @@ def write(string)
 
     # Allows you to manually set or override the response body.
     def body=(body)
-      if body.respond_to?(:to_path)
-        @stream = body
-      else
-        synchronize do
-          @stream = build_buffer self, munge_body_object(body)
+      # Prevent ActionController::Metal::Live::Response from committing the response prematurely.
+      synchronize do
+        if body.respond_to?(:to_str)
+          @stream = build_buffer(self, [body])
+        elsif body.respond_to?(:to_path)
+          @stream = body
+        elsif body.respond_to?(:to_ary)
+          @stream = build_buffer(self, body)
+        else
+          @stream = body
         end
       end
     end
@@ -482,10 +505,6 @@ def build_buffer(response, body)
       Buffer.new response, body
     end
 
-    def munge_body_object(body)
-      body.respond_to?(:each) ? body : [body]
-    end
-
     def assign_default_content_type_and_charset!
       return if media_type
 
@@ -499,6 +518,8 @@ def initialize(response)
         @response = response
       end
 
+      attr :response
+
       def close
         # Rack "close" maps to Response#abort, and **not** Response#close (which is used
         # when the controller's finished writing)

actionpack/test/controller/new_base/render_streaming_test.rb

@@ -45,44 +45,78 @@ def explicit_cache
 
   class StreamingTest < Rack::TestCase
     test "rendering with streaming enabled at the class level" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/hello_world")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["Hello world", ", I'm here!"], body
+
       get "/render_streaming/basic/hello_world"
       assert_body "Hello world, I'm here!"
-      assert_streaming!
     end
 
     test "rendering with streaming given to render" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/explicit")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["Hello world", ", I'm here!"], body
+
       get "/render_streaming/basic/explicit"
       assert_body "Hello world, I'm here!"
-      assert_streaming!
+      assert_cache_control!
     end
 
     test "rendering with streaming do not override explicit cache control given to render" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/explicit_cache")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["Hello world", ", I'm here!"], body
+
       get "/render_streaming/basic/explicit_cache"
       assert_body "Hello world, I'm here!"
-      assert_streaming! "private"
+      assert_cache_control! "private"
     end
 
     test "rendering with streaming no layout" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/no_layout")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["Hello world"], body
+
       get "/render_streaming/basic/no_layout"
       assert_body "Hello world"
-      assert_streaming!
+      assert_cache_control!
     end
 
     test "skip rendering with streaming at render level" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/skip")
+      status, _, body = app.call(env)
+      assert_equal 200, status
+      assert_chunks ["Hello world, I'm here!"], body
+
       get "/render_streaming/basic/skip"
       assert_body "Hello world, I'm here!"
     end
 
     test "rendering with layout exception" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/layout_exception")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["<body class=\"", "\"><script>window.location = \"/500.html\"</script></html>"], body
+
       get "/render_streaming/basic/layout_exception"
       assert_body "<body class=\"\"><script>window.location = \"/500.html\"</script></html>"
-      assert_streaming!
+      assert_cache_control!
     end
 
     test "rendering with template exception" do
+      env = Rack::MockRequest.env_for("/render_streaming/basic/template_exception")
+      status, headers, body = app.call(env)
+      assert_streaming!(status, headers, body)
+      assert_chunks ["\"><script>window.location = \"/500.html\"</script></html>"], body
+
       get "/render_streaming/basic/template_exception"
       assert_body "\"><script>window.location = \"/500.html\"</script></html>"
-      assert_streaming!
+      assert_cache_control!
     end
 
     test "rendering with template exception logs the exception" do
@@ -98,9 +132,28 @@ class StreamingTest < Rack::TestCase
       end
     end
 
-    def assert_streaming!(cache = "no-cache")
-      assert_status 200
-      assert_equal cache, headers["cache-control"]
+    def assert_streaming!(status, headers, body)
+      assert_equal 200, status
+
+      # It should not have a content length
+      assert_nil headers["content-length"]
+
+      # The body should not respond to `#to_ary`
+      assert_not_respond_to body, :to_ary
+    end
+
+    def assert_cache_control!(value = "no-cache", headers: self.headers)
+      assert_equal value, headers["cache-control"]
+    end
+
+    def assert_chunks(expected, body)
+      index = 0
+      body.each do |chunk|
+        assert_equal expected[index], chunk
+        index += 1
+      end
+
+      assert_equal expected.size, index
     end
   end
 end

actionpack/test/dispatch/response_test.rb

@@ -42,7 +42,12 @@ def test_write_after_close
   def test_each_isnt_called_if_str_body_is_written
     # Controller writes and reads response body
     each_counter = 0
-    @response.body = Object.new.tap { |o| o.singleton_class.define_method(:each) { |&block| each_counter += 1; block.call "foo" } }
+
+    @response.body = Object.new.tap do |object|
+      object.singleton_class.define_method(:each) { |&block| each_counter += 1; block.call "foo" }
+      object.singleton_class.define_method(:to_ary) { enum_for(:each).to_a }
+    end
+
     @response["X-Foo"] = @response.body
 
     assert_equal 1, each_counter, "#each was not called once"
@@ -647,4 +652,17 @@ class ResponseIntegrationTest < ActionDispatch::IntegrationTest
     assert_equal("utf-8", @response.charset)
     assert_equal("012345678910", @response.body)
   end
+
+  test "response does not buffer enumerator body" do
+    # This is an enumerable body, and it should not be buffered:
+    body = Enumerator.new do |enumerator|
+      enumerator << "Hello World"
+    end
+
+    # The response created here should not attempt to buffer the body:
+    response = ActionDispatch::Response.new(200, { "content-type" => "text/plain" }, body)
+
+    # The body should be the same enumerator object, i.e. it should be passed through unchanged:
+    assert_equal body, response.body
+  end
 end

actionview/lib/action_view/renderer/streaming_template_renderer.rb

@@ -25,6 +25,13 @@ def each(&block)
         self
       end
 
+      # Returns the complete body as a string.
+      def body
+        buffer = String.new
+        each { |part| buffer << part }
+        buffer
+      end
+
       private
         # This is the same logging logic as in ShowExceptions middleware.
         def log_error(exception)
@@ -42,7 +49,7 @@ def log_error(exception)
     # object that responds to each. This object is initialized with a block
     # that knows how to render the template.
     def render_template(view, template, layout_name = nil, locals = {}) # :nodoc:
-      return [super.body] unless layout_name && template.supports_streaming?
+      return [super.body] unless template.supports_streaming?
 
       locals ||= {}
       layout   = find_layout(layout_name, locals.keys, [formats.first])

actionview/lib/action_view/template/raw_file.rb

@@ -20,6 +20,10 @@ def identifier
       def render(*args)
         ::File.read(@filename)
       end
+
+      def supports_streaming?
+        false
+      end
     end
   end
 end