close client socket after closing response body

Response bodies may capture the block passed to each
and save it for body.close, so don't close the socket
before we have a chance to call body.close

lib/unicorn/http_response.rb

@@ -38,7 +38,6 @@ def http_response_write(socket, status, headers, body)
     end
 
     body.each { |chunk| socket.write(chunk) }
-    socket.close # flushes and uncorks the socket immediately
     ensure
       body.respond_to?(:close) and body.close
   end

lib/unicorn/http_server.rb

@@ -538,6 +538,7 @@ def process_client(client)
     end
     @request.headers? or headers = nil
     http_response_write(client, status, headers, body)
+    client.close # flush and uncork socket immediately, no keepalive
   rescue => e
     handle_error(client, e)
   end

t/t0018-write-on-close.sh

@@ -0,0 +1,23 @@
+#!/bin/sh
+. ./test-lib.sh
+t_plan 4 "write-on-close tests for funky response-bodies"
+
+t_begin "setup and start" && {
+	unicorn_setup
+	unicorn -D -c $unicorn_config write-on-close.ru
+	unicorn_wait_start
+}
+
+t_begin "write-on-close response body succeeds" && {
+	test xGoodbye = x"$(curl -sSf http://$listen/)"
+}
+
+t_begin "killing succeeds" && {
+	kill $unicorn_pid
+}
+
+t_begin "check stderr" && {
+	check_stderr
+}
+
+t_done

t/write-on-close.ru

@@ -0,0 +1,11 @@
+class WriteOnClose
+  def each(&block)
+    @callback = block
+  end
+
+  def close
+    @callback.call "7\r\nGoodbye\r\n0\r\n\r\n"
+  end
+end
+use Rack::ContentType, "text/plain"
+run(lambda { |_| [ 200, [%w(Transfer-Encoding chunked)], WriteOnClose.new ] })

test/unit/test_response.rb

@@ -27,31 +27,31 @@ def test_httpdate
   def test_response_headers
     out = StringIO.new
     http_response_write(out, 200, {"X-Whatever" => "stuff"}, ["cool"])
-    assert out.closed?
+    assert ! out.closed?
 
     assert out.length > 0, "output didn't have data"
   end
 
   def test_response_string_status
     out = StringIO.new
     http_response_write(out,'200', {}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert out.length > 0, "output didn't have data"
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/).size
   end
 
   def test_response_200
     io = StringIO.new
     http_response_write(io, 200, {}, [])
-    assert io.closed?
+    assert ! io.closed?
     assert io.length > 0, "output didn't have data"
   end
 
   def test_response_with_default_reason
     code = 400
     io = StringIO.new
     http_response_write(io, code, {}, [])
-    assert io.closed?
+    assert ! io.closed?
     lines = io.string.split(/\r\n/)
     assert_match(/.* Bad Request$/, lines.first,
                  "wrong default reason phrase")
@@ -60,7 +60,7 @@ def test_response_with_default_reason
   def test_rack_multivalue_headers
     out = StringIO.new
     http_response_write(out,200, {"X-Whatever" => "stuff\nbleh"}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_match(/^X-Whatever: stuff\r\nX-Whatever: bleh\r\n/, out.string)
   end
 
@@ -69,7 +69,7 @@ def test_rack_multivalue_headers
   def test_status_header_added
     out = StringIO.new
     http_response_write(out,200, {"X-Whatever" => "stuff"}, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/i).size
   end
 
@@ -80,7 +80,7 @@ def test_status_header_ignores_app_hash
     out = StringIO.new
     header_hash = {"X-Whatever" => "stuff", 'StaTus' => "666" }
     http_response_write(out,200, header_hash, [])
-    assert out.closed?
+    assert ! out.closed?
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status: 200 OK/i).size
     assert_equal 1, out.string.split(/\r\n/).grep(/^Status:/i).size
   end
@@ -91,15 +91,15 @@ def test_body_closed
     body.rewind
     out = StringIO.new
     http_response_write(out,200, {}, body)
-    assert out.closed?
+    assert ! out.closed?
     assert body.closed?
     assert_match(expect_body, out.string.split(/\r\n/).last)
   end
 
   def test_unknown_status_pass_through
     out = StringIO.new
     http_response_write(out,"666 I AM THE BEAST", {}, [] )
-    assert out.closed?
+    assert ! out.closed?
     headers = out.string.split(/\r\n\r\n/).first.split(/\r\n/)
     assert %r{\AHTTP/\d\.\d 666 I AM THE BEAST\z}.match(headers[0])
     status = headers.grep(/\AStatus:/i).first