[PATCH 3/4] response: support auto-chunking for HTTP/1.1

[Eric Wong <e@80x24.org>]

We might as well do it since puma and thin both do(*),
and we can still do writev for now to get some speedups
by avoiding Rack::Chunked overhead.

timing runs of "curl --no-buffer http://127.0.0.1:9292/ >/dev/null"
results in a best case drop from ~260ms to ~205ms on one VM
by disabling Rack::Chunked in the below config.ru

$ ruby -I lib bin/yahns-rackup -E none config.ru

==> config.ru <==
class Body
  STR = ' ' * 1024 * 16
  def each
    10000.times { yield STR }
  end
end

use Rack::Chunked if ENV['RACK_CHUNKED']
run(lambda do |env|
  [ 200, [ %w(Content-Type text/plain) ], Body.new ]
end)

(*) they can do Content-Length, but I don't think it's
    worth the effort at the server level.
---
 lib/yahns/chunk_body.rb    | 27 ++++++++++++++++++++++
 lib/yahns/http_client.rb   |  8 +++----
 lib/yahns/http_response.rb | 38 +++++++++++++++++++++----------
 test/test_auto_chunk.rb    | 56 ++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 113 insertions(+), 16 deletions(-)
 create mode 100644 lib/yahns/chunk_body.rb
 create mode 100644 test/test_auto_chunk.rb

diff --git a/lib/yahns/chunk_body.rb b/lib/yahns/chunk_body.rb
new file mode 100644
index 0000000..6e56a18
--- /dev/null
+++ b/lib/yahns/chunk_body.rb
@@ -0,0 +1,27 @@
+# -*- encoding: binary -*-
+# Copyright (C) 2016 all contributors <yahns-public@yhbt.net>
+# License: GPL-3.0+ <https://www.gnu.org/licenses/gpl-3.0.txt>
+# frozen_string_literal: true
+class Yahns::ChunkBody
+  def initialize(body, vec)
+    @body = body
+    @vec = vec
+  end
+
+  def each
+    vec = @vec
+    vec[2] = "\r\n".freeze
+    @body.each do |chunk|
+      vec[0] = "#{chunk.bytesize.to_s(16)}\r\n"
+      vec[1] = chunk
+      # vec[2] never changes: "\r\n" above
+      yield vec
+    end
+    vec.clear
+    yield "0\r\n\r\n".freeze
+  end
+
+  def close
+    @body.close if @body.respond_to?(:close)
+  end
+end
diff --git a/lib/yahns/http_client.rb b/lib/yahns/http_client.rb
index 7a1bac1..d8154a4 100644
--- a/lib/yahns/http_client.rb
+++ b/lib/yahns/http_client.rb
@@ -190,10 +190,10 @@ def r100_done
       true
     else # :lazy, false
       env = @hs.env
-      hdr_only = env['REQUEST_METHOD'] == 'HEAD'.freeze
+      opt = http_response_prep(env)
       res = k.app.call(env)
       return :ignore if app_hijacked?(env, res)
-      http_response_write(res, hdr_only)
+      http_response_write(res, opt)
     end
   end
 
@@ -222,7 +222,7 @@ def app_call(input)
       env['SERVER_PORT'] = '443'.freeze
     end
 
-    hdr_only = env['REQUEST_METHOD'] == 'HEAD'.freeze
+    opt = http_response_prep(env)
     # run the rack app
     res = k.app.call(env)
     return :ignore if app_hijacked?(env, res)
@@ -232,7 +232,7 @@ def app_call(input)
     end
 
     # this returns :wait_readable, :wait_writable, :ignore, or nil:
-    http_response_write(res, hdr_only)
+    http_response_write(res, opt)
   end
 
   # called automatically by kgio_write
diff --git a/lib/yahns/http_response.rb b/lib/yahns/http_response.rb
index b157ee4..32d1a45 100644
--- a/lib/yahns/http_response.rb
+++ b/lib/yahns/http_response.rb
@@ -4,12 +4,14 @@
 # frozen_string_literal: true
 require_relative 'stream_file'
 require_relative 'wbuf_str'
+require_relative 'chunk_body'
 
 # Writes a Rack response to your client using the HTTP/1.1 specification.
 # You use it by simply doing:
 #
+#   opt = http_response_prep(env)
 #   res = rack_app.call(env)
-#   http_response_write(res, env['REQUEST_METHOD']=='HEAD')
+#   http_response_write(res, opt)
 #
 # Most header correctness (including Content-Length and Content-Type)
 # is the job of Rack, with the exception of the "Date" header.
@@ -120,14 +122,14 @@ def kv_str(buf, key, value)
 
   # writes the rack_response to socket as an HTTP response
   # returns :wait_readable, :wait_writable, :forget, or nil
-  def http_response_write(res, hdr_only)
+  def http_response_write(res, opt)
     status, headers, body = res
     offset = 0
     count = hijack = nil
-    k = self.class
-    alive = @hs.next? && k.persistent_connections
+    alive = @hs.next? && self.class.persistent_connections
     flags = MSG_DONTWAIT
     term = false
+    hdr_only, chunk_ok = opt
 
     if @hs.headers?
       code = status.to_i
@@ -161,6 +163,11 @@ def http_response_write(res, hdr_only)
           kv_str(buf, key, value)
         end
       end
+      if !term && chunk_ok
+        term = true
+        body = Yahns::ChunkBody.new(body, opt)
+        buf << "Transfer-Encoding: chunked\r\n".freeze
+      end
       alive &&= term
       buf << (alive ? "Connection: keep-alive\r\n\r\n".freeze
                     : "Connection: close\r\n\r\n".freeze)
@@ -173,7 +180,7 @@ def http_response_write(res, hdr_only)
         flags = MSG_DONTWAIT
         buf = rv # unlikely, hope the skb grows
       when :wait_writable, :wait_readable # unlikely
-        if k.output_buffering
+        if self.class.output_buffering
           alive = hijack ? hijack : alive
           rv = response_header_blocked(buf, body, alive, offset, count)
           body = nil # ensure we do not close body in ensure
@@ -193,19 +200,19 @@ def http_response_write(res, hdr_only)
     end
 
     wbuf = rv = nil
-    body.each do |chunk|
+    body.each do |x|
       if wbuf
-        rv = wbuf.wbuf_write(self, chunk)
+        rv = wbuf.wbuf_write(self, x)
       else
-        case rv = kgio_trywrite(chunk)
+        case rv = String === x ? kgio_trywrite(x) : kgio_trywritev(x)
         when nil # all done, likely and good!
           break
-        when String
-          chunk = rv # hope the skb grows when we loop into the trywrite
+        when String, Array
+          x = rv # hope the skb grows when we loop into the trywrite
         when :wait_writable, :wait_readable
-          if k.output_buffering
+          if self.class.output_buffering
             wbuf = Yahns::Wbuf.new(body, alive)
-            rv = wbuf.wbuf_write(self, chunk)
+            rv = wbuf.wbuf_write(self, x)
             break
           else
             response_wait_write(rv) or return :close
@@ -278,4 +285,11 @@ def http_100_response(env)
       return rv
     end while true
   end
+
+  # must be called before app dispatch, since the app can
+  # do all sorts of nasty things to env
+  def http_response_prep(env)
+    [ env['REQUEST_METHOD'] == 'HEAD'.freeze, # hdr_only
+      env['HTTP_VERSION'] == 'HTTP/1.1'.freeze ] # chunk_ok
+  end
 end
diff --git a/test/test_auto_chunk.rb b/test/test_auto_chunk.rb
new file mode 100644
index 0000000..a97fe26
--- /dev/null
+++ b/test/test_auto_chunk.rb
@@ -0,0 +1,56 @@
+# Copyright (C) 2013-2016 all contributors <yahns-public@yhbt.net>
+# License: GPL-3.0+ <https://www.gnu.org/licenses/gpl-3.0.txt>
+# frozen_string_literal: true
+require_relative 'server_helper'
+
+class TestAutoChunk < Testcase
+  ENV["N"].to_i > 1 and parallelize_me!
+  include ServerHelper
+  alias setup server_helper_setup
+  alias teardown server_helper_teardown
+
+  def test_auto_head
+    err = @err
+    cfg = Yahns::Config.new
+    host, port = @srv.addr[3], @srv.addr[1]
+    cfg.instance_eval do
+      GTL.synchronize do
+        app = Rack::Builder.new do
+          use Rack::ContentType, "text/plain"
+          run(lambda do |env|
+            [ 200, {}, %w(a b c) ]
+          end)
+        end
+        app(:rack, app) { listen "#{host}:#{port}" }
+      end
+      logger(Logger.new(err.path))
+    end
+    pid = mkserver(cfg)
+    s = TCPSocket.new(host, port)
+    s.write("GET / HTTP/1.0\r\n\r\n")
+    assert s.wait(30), "IO wait failed"
+    buf = s.read
+    assert_match %r{\r\n\r\nabc\z}, buf
+    s.close
+
+    s = TCPSocket.new(host, port)
+    s.write("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
+    buf = ''.dup
+    Timeout.timeout(30) do
+      until buf =~ /\r\n\r\n1\r\na\r\n1\r\nb\r\n1\r\nc\r\n0\r\n\r\n\z/
+        buf << s.readpartial(16384)
+      end
+    end
+    assert_match(%r{^Transfer-Encoding: chunked\r\n}, buf)
+    s.close
+
+    Net::HTTP.start(host, port) do |http|
+      req = Net::HTTP::Get.new("/")
+      res = http.request(req)
+      assert_equal 200, res.code.to_i
+      assert_equal 'abc', res.body
+    end
+  ensure
+    quit_wait(pid)
+  end
+end
-- 
EW