Merge pull request #252 from httprb/zanker/stress-tests

Fixed chunk writing if the data is bugger than our buffer (#251)

Author: tarcieri

lib/http/request/writer.rb

@@ -12,6 +12,9 @@ class Writer
       # Chunked transfer encoding
       CHUNKED = "chunked".freeze
 
+      # End of a chunked transfer
+      CHUNKED_END = "#{ZERO}#{CRLF}#{CRLF}".freeze
+
       # Types valid to be used as body source
       VALID_BODY_TYPES = [String, NilClass, Enumerable]
 
@@ -40,7 +43,7 @@ def stream
       # Send headers needed to connect through proxy
       def connect_through_proxy
         add_headers
-        @socket << join_headers
+        write(join_headers)
       end
 
       # Adds the headers to the header array for the given request body we are working
@@ -65,24 +68,35 @@ def send_request_header
         add_headers
         add_body_type_headers
 
-        @socket << join_headers
+        write(join_headers)
       end
 
       def send_request_body
         if @body.is_a?(String)
-          @socket << @body
+          write(@body)
         elsif @body.is_a?(Enumerable)
           @body.each do |chunk|
-            @socket << chunk.bytesize.to_s(16) << CRLF
-            @socket << chunk << CRLF
+            write(chunk.bytesize.to_s(16) << CRLF)
+            write(chunk << CRLF)
           end
 
-          @socket << ZERO << CRLF << CRLF
+          write(CHUNKED_END)
         end
       end
 
       private
 
+      def write(data)
+        while data.present?
+          length = @socket.write(data)
+          if data.length > length
+            data = data[length..-1]
+          else
+            break
+          end
+        end
+      end
+
       def validate_body_type!
         return if VALID_BODY_TYPES.any? { |type| @body.is_a? type }
         fail RequestError, "body of wrong type: #{@body.class}"

lib/http/timeout/global.rb

@@ -61,7 +61,7 @@ def write(data)
           reset_timer
 
           begin
-            socket.write_nonblock(data)
+            return socket.write_nonblock(data)
           rescue IO::WaitWritable
             IO.select(nil, [socket], nil, time_left)
             log_time
@@ -97,7 +97,7 @@ def write(data)
 
           loop do
             result = socket.write_nonblock(data, :exception => false)
-            break unless result == :wait_writable
+            return result unless result == :wait_writable
 
             IO.select(nil, [socket], nil, time_left)
             log_time

lib/http/timeout/null.rb

@@ -45,7 +45,7 @@ def readpartial(size)
 
       # Write to the socket
       def write(data)
-        @socket << data
+        @socket.write(data)
       end
       alias_method :<<, :write
 

lib/http/timeout/per_operation.rb

@@ -75,7 +75,7 @@ def readpartial(size)
         def write(data)
           loop do
             result = socket.write_nonblock(data, :exception => false)
-            break unless result == :wait_writable
+            return result unless result == :wait_writable
 
             unless IO.select(nil, [socket], nil, write_timeout)
               fail TimeoutError, "Read timed out after #{write_timeout} seconds"

spec/lib/http/client_spec.rb

@@ -267,8 +267,8 @@ def simple_response(body, status = 200)
 
         allow(socket_spy).to receive(:close) { nil }
         allow(socket_spy).to receive(:closed?) { true }
-        allow(socket_spy).to receive(:readpartial) { chunks.shift }
-        allow(socket_spy).to receive(:<<) { nil }
+        allow(socket_spy).to receive(:readpartial) { chunks[0] }
+        allow(socket_spy).to receive(:write) { chunks[0].length }
 
         allow(TCPSocket).to receive(:open) { socket_spy }
       end

spec/lib/http_spec.rb

@@ -44,6 +44,32 @@
         expect(response.to_s.include?("json")).to be true
       end
     end
+
+    context "with a large request body" do
+      %w(global null per_operation).each do |timeout|
+        context "with a #{timeout} timeout" do
+          [16_000, 16_500, 17_000, 34_000, 68_000].each do |size|
+            [0, rand(0..100), rand(100..1000)].each do |fuzzer|
+              context "with a #{size} body and #{fuzzer} of fuzzing" do
+                let(:client) { HTTP.timeout(timeout, :read => 2, :write => 2, :connect => 2) }
+
+                let(:characters) { ("A".."Z").to_a }
+                let(:request_body) do
+                  (size + fuzzer).times.map { |i| characters[i % characters.length] }.join
+                end
+
+                it "returns a large body" do
+                  response = client.post("#{dummy.endpoint}/echo-body", :body => request_body)
+
+                  expect(response.body.to_s).to eq(request_body)
+                  expect(response.headers["Content-Length"].to_i).to eq(request_body.length)
+                end
+              end
+            end
+          end
+        end
+      end
+    end
   end
 
   describe ".via" do

spec/support/dummy_server/servlet.rb

@@ -133,5 +133,10 @@ def do_#{method.upcase}(req, res)
       res["Set-Cookie"] = "foo=bar"
       res.body = req.cookies.map { |c| [c.name, c.value].join ": " }.join("\n")
     end
+
+    post "/echo-body" do |req, res|
+      res.status = 200
+      res.body   = req.body
+    end
   end
 end