Merge pull request #1434 from rack/rack-response-buffered

Lazily initialize the response body and only buffer it if required.

3 files changed, 107 insertions, 39 deletions

diff --git a/lib/rack/mock.rb b/lib/rack/mock.rb
index f2b94832..3feaedd9 100644
--- a/lib/rack/mock.rb
+++ b/lib/rack/mock.rb
@@ -169,6 +169,8 @@ module Rack
       @cookies = parse_cookies_from_header
 
       super(body, status, headers)
+
+      buffered_body!
     end
 
     def =~(other)
@@ -190,7 +192,13 @@ module Rack
       #     ...
       #     res.body.should == "foo!"
       #   end
-      super.join
+      buffer = String.new
+
+      super.each do |chunk|
+        buffer << chunk
+      end
+
+      return buffer
     end
 
     def empty?
diff --git a/lib/rack/response.rb b/lib/rack/response.rb
index f39848cc..11f296b4 100644
--- a/lib/rack/response.rb
+++ b/lib/rack/response.rb
@@ -26,32 +26,30 @@ module Rack
     alias headers header
 
     CHUNKED = 'chunked'
-    STATUS_WITH_NO_ENTITY_BODY = {
-      204 => true,
-      304 => true
-    }.freeze
+    STATUS_WITH_NO_ENTITY_BODY = Utils::STATUS_WITH_NO_ENTITY_BODY
 
-    def initialize(body = [], status = 200, header = {})
+    def initialize(body = nil, status = 200, header = {})
       @status = status.to_i
       @header = Utils::HeaderHash.new(header)
 
-      @writer  = lambda { |x| @body << x }
-      @block   = nil
-      @length  = 0
+      @writer = self.method(:append)
 
-      @body = []
+      @block = nil
+      @length = 0
 
-      if body.respond_to? :to_str
-        write body.to_str
-      elsif body.respond_to?(:each)
-        body.each { |part|
-          write part.to_s
-        }
+      # Keep track of whether we have expanded the user supplied body.
+      if body.nil?
+        @body = []
+        @buffered = true
+      elsif body.respond_to? :to_str
+        @body = [body]
+        @buffered = true
       else
-        raise TypeError, "stringable or iterable required"
+        @body = body
+        @buffered = false
       end
 
-      yield self  if block_given?
+      yield self if block_given?
     end
 
     def redirect(target, status = 302)
@@ -64,40 +62,45 @@ module Rack
     end
 
     def finish(&block)
-      @block = block
-
-      if STATUS_WITH_NO_ENTITY_BODY.key?(status.to_i)
+      if STATUS_WITH_NO_ENTITY_BODY[status.to_i]
         delete_header CONTENT_TYPE
         delete_header CONTENT_LENGTH
         close
         [status.to_i, header, []]
       else
-        [status.to_i, header, self]
+        if block_given?
+          @block = block
+          [status.to_i, header, self]
+        else
+          [status.to_i, header, @body]
+        end
       end
     end
+
     alias to_a finish           # For *response
 
     def each(&callback)
       @body.each(&callback)
-      @writer = callback
-      @block.call(self)  if @block
+      @buffered = true
+
+      if @block
+        @writer = callback
+        @block.call(self)
+      end
     end
 
     # Append to body and update Content-Length.
     #
     # NOTE: Do not mix #write and direct #body access!
     #
-    def write(str)
-      s = str.to_s
-      @length += s.bytesize unless chunked?
-      @writer.call s
+    def write(chunk)
+      buffered_body!
 
-      set_header(CONTENT_LENGTH, @length.to_s) unless chunked?
-      str
+      @writer.call(chunk.to_s)
     end
 
     def close
-      body.close if body.respond_to?(:close)
+      @body.close if @body.respond_to?(:close)
     end
 
     def empty?
@@ -216,6 +219,38 @@ module Rack
       def etag= v
         set_header ETAG, v
       end
+
+    protected
+
+      def buffered_body!
+        return if @buffered
+
+        if @body.is_a?(Array)
+          # The user supplied body was an array:
+          @body = @body.dup
+        else
+          # Turn the user supplied body into a buffered array:
+          body = @body
+          @body = Array.new
+
+          body.each do |part|
+            @writer.call(part.to_s)
+          end
+        end
+        
+        @buffered = true
+      end
+
+      def append(chunk)
+        @body << chunk
+        
+        unless chunked?
+          @length += chunk.bytesize
+          set_header(CONTENT_LENGTH, @length.to_s)
+        end
+
+        return chunk
+      end
     end
 
     include Helpers
diff --git a/test/spec_response.rb b/test/spec_response.rb
index 6958f429..c0736c73 100644
--- a/test/spec_response.rb
+++ b/test/spec_response.rb
@@ -307,9 +307,9 @@ describe Rack::Response do
     header['Content-Length'].must_be_nil
 
     lambda {
-      Rack::Response.new(Object.new)
-    }.must_raise(TypeError).
-      message.must_match(/stringable or iterable required/)
+      Rack::Response.new(Object.new).each{}
+    }.must_raise(NoMethodError).
+      message.must_match(/undefined method .each. for/)
   end
 
   it "knows if it's empty" do
@@ -433,6 +433,28 @@ describe Rack::Response do
     res.headers["Content-Length"].must_equal "8"
   end
 
+  it "does not wrap body" do
+    body = Object.new
+    res = Rack::Response.new(body)
+
+    # It was passed through unchanged:
+    res.finish.last.must_equal body
+  end
+
+  it "does wraps body when using #write" do
+    body = ["Foo"]
+    res = Rack::Response.new(body)
+
+    # Write something using the response object:
+    res.write("Bar")
+
+    # The original body was not modified:
+    body.must_equal ["Foo"]
+
+    # But a new buffered body was created:
+    res.finish.last.must_equal ["Foo", "Bar"]
+  end
+
   it "calls close on #body" do
     res = Rack::Response.new
     res.body = StringIO.new
@@ -452,16 +474,19 @@ describe Rack::Response do
     b.wont_equal res.body
 
     res.body = StringIO.new
-    res.status = 205
+    res.status = 304
     _, _, b = res.finish
-    res.body.wont_be :closed?
+    res.body.must_be :closed?
     b.wont_equal res.body
+  end
+
+  it "doesn't call close on #body when 205" do
+    res = Rack::Response.new
 
     res.body = StringIO.new
-    res.status = 304
+    res.status = 205
     _, _, b = res.finish
-    res.body.must_be :closed?
-    b.wont_equal res.body
+    res.body.wont_be :closed?
   end
 
   it "flatten doesn't cause infinite loop" do