1 files changed, 69 insertions, 7 deletions

diff --git a/lib/yahns/wbuf_lite.rb b/lib/yahns/wbuf_lite.rb
index 25daf21..2da5349 100644
--- a/lib/yahns/wbuf_lite.rb
+++ b/lib/yahns/wbuf_lite.rb
@@ -2,30 +2,88 @@
 # 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
-require_relative 'wbuf'
+require 'stringio'
+require_relative 'wbuf_common'
 
 # This is only used for "proxy_buffering: false"
-class Yahns::WbufLite < Yahns::Wbuf # :nodoc:
+class Yahns::WbufLite # :nodoc:
+  include Yahns::WbufCommon
   attr_reader :busy
   attr_writer :req_res
 
   def initialize(req_res)
-    alive = req_res.alive
-    super(nil, alive ? :ignore : false)
+    @tmpio = nil
+    @sf_offset = @sf_count = 0
+    @wbuf_persist = :ignore
+    @busy = false
     @req_res = req_res
   end
 
-  def wbuf_write(client, buf)
-    super
+  def wbuf_write(c, buf)
+    buf = buf.join if Array === buf
+    # try to bypass the VFS layer and write directly to the socket
+    # if we're all caught up
+    case rv = c.kgio_trywrite(buf)
+    when String
+      buf = rv # retry in loop
+    when nil
+      return # yay! hopefully we don't have to buffer again
+    when :wait_writable, :wait_readable
+      @busy = rv
+    end until @busy
+
+    @tmpio ||= StringIO.new(''.dup) # relies on encoding: binary above
+    @sf_count += @tmpio.write(buf)
+
+    # we spent some time copying to the FS, try to write to
+    # the socket again in case some space opened up...
+    case rv = c.trysendio(@tmpio, @sf_offset, @sf_count)
+    when Integer
+      @sf_count -= rv
+      @sf_offset += rv
+    when :wait_writable, :wait_readable
+      @busy = rv
+      return rv
+    else
+      raise "BUG: #{rv.nil? ? "EOF" : rv.inspect} on tmpio " \
+            "sf_offset=#@sf_offset sf_count=#@sf_count"
+    end while @sf_count > 0
+
+    # we're all caught up, try to prevent dirty data from getting flushed
+    # to disk if we can help it.
+    wbuf_abort
+    @sf_offset = 0
+    @busy = false
+    nil
   rescue
     @req_res = @req_res.close if @req_res
     raise
   end
 
   def wbuf_flush(client)
-    super
+    case rv = client.trysendio(@tmpio, @sf_offset, @sf_count)
+    when Integer
+      return wbuf_close(client) if (@sf_count -= rv) == 0 # all sent!
+      @sf_offset += rv # keep going otherwise
+    when :wait_writable, :wait_readable
+      return rv
+    when nil
+      # response got truncated, drop the connection
+      # this may happens when using Rack::File or similar, we can't
+      # keep the connection alive because we already sent our Content-Length
+      # header the client would be confused.
+      @wbuf_persist = false
+      return wbuf_close(client)
+    else
+      raise "BUG: rv=#{rv.inspect} " \
+            "on tmpio=#{@tmpio.inspect} " \
+            "sf_offset=#@sf_offset sf_count=#@sf_count"
+    end while @sf_count > 0
+    wbuf_close(client)
   rescue
+    @wbuf_persist = false # ensure a hijack response is not called
     @req_res = @req_res.close if @req_res
+    wbuf_close(client)
     raise
   end
 
@@ -48,4 +106,8 @@ class Yahns::WbufLite < Yahns::Wbuf # :nodoc:
     @req_res = @req_res.close if @req_res
     raise
   end
+
+  def wbuf_abort
+    @tmpio = @tmpio.close if @tmpio
+  end
 end