io_buffer.rb, request.rb - improve handling with a body array/enumeration

lib/puma/io_buffer.rb

@@ -1,11 +1,37 @@
 # frozen_string_literal: true
 
+require 'stringio'
+
 module Puma
-  class IOBuffer < String
-    def append(*args)
-      args.each { |a| concat(a) }
+  class IOBuffer < StringIO
+    def initialize
+      super.binmode
+    end
+
+    def empty?
+      length.zero?
+    end
+
+    def reset
+      truncate 0
+      rewind
     end
 
-    alias reset clear
+    def read
+      rewind
+      super.tap { |s| truncate 0; rewind }
+    end
+
+    # don't use, added just for existing CI tests
+    alias_method :to_s, :string
+
+    # before Ruby 2.5, `write` would only take one argument
+    if RUBY_VERSION >= '2.5' && RUBY_ENGINE != 'truffleruby'
+      alias_method :append, :write
+    else
+      def append(*strs)
+        strs.each { |str| write str }
+      end
+    end
   end
 end

lib/puma/request.rb

@@ -12,6 +12,8 @@ module Puma
   #
   module Request
 
+    BUFFER_LENGTH = 128 * 1024
+
     include Puma::Const
 
     # Takes the request contained in +client+, invokes the Rack application to construct
@@ -48,8 +50,6 @@ def handle_request(client, lines, requests)
 
       body = client.body
 
-      head = env[REQUEST_METHOD] == HEAD
-
       env[RACK_INPUT] = body
       env[RACK_URL_SCHEME] ||= default_server_port(env) == PORT_443 ? HTTPS : HTTP
 
@@ -99,19 +99,10 @@ def handle_request(client, lines, requests)
           status, headers, res_body = lowlevel_error(e, env, 500)
         end
 
-        res_info = {}
-        res_info[:content_length] = nil
-        res_info[:no_body] = head
-
-        res_info[:content_length] = if res_body.kind_of? Array and res_body.size == 1
-          res_body[0].bytesize
-        else
-          nil
-        end
-
-        cork_socket io
+        content_length = res_body.kind_of?(::Array) && res_body.size == 1 ?
+          res_body[0].bytesize : nil
 
-        str_headers(env, status, headers, res_info, lines, requests, client)
+        res_info = str_headers(env, status, headers, content_length, lines, requests, client)
 
         line_ending = LINE_END
 
@@ -124,7 +115,7 @@ def handle_request(client, lines, requests)
           end
 
           lines << LINE_END
-          fast_write io, lines.to_s
+          fast_write io, lines.read
           return res_info[:keep_alive]
         end
 
@@ -138,45 +129,29 @@ def handle_request(client, lines, requests)
 
         lines << line_ending
 
-        fast_write io, lines.to_s
-
         if response_hijack
+          fast_write io, lines.read
           response_hijack.call io
           return :async
         end
 
-        begin
-          res_body.each do |part|
-            next if part.bytesize.zero?
-            if chunked
-               fast_write io, (part.bytesize.to_s(16) << line_ending)
-               fast_write io, part            # part may have different encoding
-               fast_write io, line_ending
-            else
-              fast_write io, part
-            end
-            io.flush
-          end
-
-          if chunked
-            fast_write io, CLOSE_CHUNKED
-            io.flush
-          end
-        rescue SystemCallError, IOError
-          raise ConnectionError, "Connection error detected during write"
-        end
+        fast_write_body io, res_body, lines, chunked
 
       ensure
-        uncork_socket io
-
+        io.flush
+        lines.reset
         body.close
+
+        keep_alive = res_info[:keep_alive]
+        res_info = nil
+
         client.tempfile.unlink if client.tempfile
         res_body.close if res_body.respond_to? :close
 
         after_reply.each { |o| o.call }
       end
 
-      res_info[:keep_alive]
+      keep_alive
     end
 
     # @param env [Hash] see Puma::Client#env, from request
@@ -190,39 +165,73 @@ def default_server_port(env)
       end
     end
 
-    # Writes to an io (normally Client#io) using #syswrite
-    # @param io [#syswrite] the io to write to
+    # Used to write 'early hints', 'no body' responses, 'hijacked' responses,
+    # and body segments (called by `fast_write_ary`).
+    # Writes a string to an io (normally `Client#io`) using `write_nonblock`.
+    # Large strings may not be written in one pass, especially if `io` is a
+    # `MiniSSL::Socket`.
+    # @param io [#write_nonblock] the io to write to
     # @param str [String] the string written to the io
     # @raise [ConnectionError]
     #
     def fast_write(io, str)
       n = 0
-      while true
+      byte_size = str.bytesize
+      while n < byte_size
         begin
-          n = io.syswrite str
+          n += io.syswrite(n == 0 ? str : str.byteslice(n..-1))
         rescue Errno::EAGAIN, Errno::EWOULDBLOCK
           unless io.wait_writable WRITE_TIMEOUT
             raise ConnectionError, "Socket timeout writing data"
           end
-
           retry
         rescue  Errno::EPIPE, SystemCallError, IOError
           raise ConnectionError, "Socket timeout writing data"
         end
-
-        return if n == str.bytesize
-        str = str.byteslice(n..-1)
       end
     end
-    private :fast_write
 
-    # @param status [Integer] status from the app
-    # @return [String] the text description from Puma::HTTP_STATUS_CODES
+    # Used to write headers and body.
+    # Writes to an io (normally `Client#io`) using `#fast_write`.
+    # Accumulates `body` items into `strm`, then writes anytime `strm` is 128kB
+    # or larger.
+    # @param io [#write] the io to write to
+    # @param body [Enumerable, File] the body object
+    # @param strm [Puma::IOBuffer] strm to write body body into
+    # @param chunk [Boolean]
+    # @raise [ConnectionError]
     #
-    def fetch_status_code(status)
-      HTTP_STATUS_CODES.fetch(status) { 'CUSTOM' }
+    def fast_write_body(io, body, strm, chunked)
+      running_len = 0
+      if body.is_a? ::File
+        while part = body.read(BUFFER_LENGTH)
+          if chunked
+            strm.append part.bytesize.to_s(16), LINE_END, part, LINE_END
+            fast_write io, strm.read
+          else
+            fast_write io, part
+          end
+        end
+        fast_write(io, CLOSE_CHUNKED) if chunked
+      else
+        body.each do |part|
+          next if (byte_size = part.bytesize).zero?
+          running_len += byte_size
+          if running_len > BUFFER_LENGTH && byte_size != running_len
+            fast_write io, strm.read
+            running_len = 0
+          end
+          if chunked
+            strm.append byte_size.to_s(16), LINE_END, part, LINE_END
+          else
+            strm.append part
+          end
+        end
+        fast_write io, (chunked ? (strm << CLOSE_CHUNKED).read : strm.read)
+      end
     end
-    private :fetch_status_code
+
+    private :fast_write, :fast_write_body
 
     # Given a Hash +env+ for the request read from +client+, add
     # and fixup keys to comply with Rack's env guidelines.
@@ -362,6 +371,14 @@ def str_early_hints(headers)
     end
     private :str_early_hints
 
+    # @param status [Integer] status from the app
+    # @return [String] the text description from Puma::HTTP_STATUS_CODES
+    #
+    def fetch_status_code(status)
+      HTTP_STATUS_CODES.fetch status, 'CUSTOM'
+    end
+    private :fetch_status_code
+
     # Processes and write headers to the IOBuffer.
     # @param env [Hash] see Puma::Client#env, from request
     # @param status [Integer] the status returned by the Rack application
@@ -372,14 +389,18 @@ def str_early_hints(headers)
     # @param client [Puma::Client]
     # @version 5.0.3
     #
-    def str_headers(env, status, headers, res_info, lines, requests, client)
+    def str_headers(env, status, headers, content_length, lines, requests, client)
       line_ending = LINE_END
       colon = COLON
 
+      res_info = {}
+      res_info[:content_length] = content_length
+      res_info[:no_body] = env[REQUEST_METHOD] == HEAD
+
       http_11 = env[HTTP_VERSION] == HTTP_11
       if http_11
         res_info[:allow_chunked] = true
-        res_info[:keep_alive] = env.fetch(HTTP_CONNECTION, "").downcase != CLOSE
+        res_info[:keep_alive] = env.fetch(HTTP_CONNECTION, '').downcase != CLOSE
 
         # An optimization. The most common response is 200, so we can
         # reply with the proper 200 status without having to compute
@@ -395,7 +416,7 @@ def str_headers(env, status, headers, res_info, lines, requests, client)
         end
       else
         res_info[:allow_chunked] = false
-        res_info[:keep_alive] = env.fetch(HTTP_CONNECTION, "").downcase == KEEP_ALIVE
+        res_info[:keep_alive] = env.fetch(HTTP_CONNECTION, '').downcase == KEEP_ALIVE
 
         # Same optimization as above for HTTP/1.1
         #
@@ -461,6 +482,7 @@ def str_headers(env, status, headers, res_info, lines, requests, client)
       else
         lines << CONNECTION_KEEP_ALIVE if res_info[:keep_alive]
       end
+      res_info
     end
     private :str_headers
   end