Merge remote-tracking branch 'origin/ccc-tcp-v3'

* origin/ccc-tcp-v3:
  http_request: reduce insn size for check_client_connection
  support "struct tcp_info" on non-Linux and Ruby 2.2+
  revert signature change to HttpServer#process_client
  new test for check_client_connection
  check_client_connection: use tcp state on linux

3 files changed, 164 insertions, 7 deletions

diff --git a/lib/unicorn/http_request.rb b/lib/unicorn/http_request.rb
index c176083..9010007 100644
--- a/lib/unicorn/http_request.rb
+++ b/lib/unicorn/http_request.rb
@@ -2,6 +2,7 @@
 # :enddoc:
 # no stable API here
 require 'unicorn_http'
+require 'raindrops'
 
 # TODO: remove redundant names
 Unicorn.const_set(:HttpRequest, Unicorn::HttpParser)
@@ -25,8 +26,10 @@ class Unicorn::HttpParser
 
   # :stopdoc:
   HTTP_RESPONSE_START = [ 'HTTP'.freeze, '/1.1 '.freeze ]
+  EMPTY_ARRAY = [].freeze
   @@input_class = Unicorn::TeeInput
   @@check_client_connection = false
+  @@tcpi_inspect_ok = true
 
   def self.input_class
     @@input_class
@@ -80,11 +83,7 @@ class Unicorn::HttpParser
       false until add_parse(socket.kgio_read!(16384))
     end
 
-    # detect if the socket is valid by writing a partial response:
-    if @@check_client_connection && headers?
-      self.response_start_sent = true
-      HTTP_RESPONSE_START.each { |c| socket.write(c) }
-    end
+    check_client_connection(socket) if @@check_client_connection
 
     e['rack.input'] = 0 == content_length ?
                       NULL_IO : @@input_class.new(socket, self)
@@ -105,4 +104,69 @@ class Unicorn::HttpParser
   def hijacked?
     env.include?('rack.hijack_io'.freeze)
   end
+
+  if defined?(Raindrops::TCP_Info)
+    TCPI = Raindrops::TCP_Info.allocate
+
+    def check_client_connection(socket) # :nodoc:
+      if Unicorn::TCPClient === socket
+        # Raindrops::TCP_Info#get!, #state (reads struct tcp_info#tcpi_state)
+        raise Errno::EPIPE, "client closed connection".freeze,
+              EMPTY_ARRAY if closed_state?(TCPI.get!(socket).state)
+      else
+        write_http_header(socket)
+      end
+    end
+
+    def closed_state?(state) # :nodoc:
+      case state
+      when 1 # ESTABLISHED
+        false
+      when 8, 6, 7, 9, 11 # CLOSE_WAIT, TIME_WAIT, CLOSE, LAST_ACK, CLOSING
+        true
+      else
+        false
+      end
+    end
+  else
+
+    # Ruby 2.2+ can show struct tcp_info as a string Socket::Option#inspect.
+    # Not that efficient, but probably still better than doing unnecessary
+    # work after a client gives up.
+    def check_client_connection(socket) # :nodoc:
+      if Unicorn::TCPClient === socket && @@tcpi_inspect_ok
+        opt = socket.getsockopt(:IPPROTO_TCP, :TCP_INFO).inspect
+        if opt =~ /\bstate=(\S+)/
+          @@tcpi_inspect_ok = true
+          raise Errno::EPIPE, "client closed connection".freeze,
+                EMPTY_ARRAY if closed_state_str?($1)
+        else
+          @@tcpi_inspect_ok = false
+          write_http_header(socket)
+        end
+        opt.clear
+      else
+        write_http_header(socket)
+      end
+    end
+
+    def closed_state_str?(state)
+      case state
+      when 'ESTABLISHED'
+        false
+      # not a typo, ruby maps TCP_CLOSE (no 'D') to state=CLOSED (w/ 'D')
+      when 'CLOSE_WAIT', 'TIME_WAIT', 'CLOSED', 'LAST_ACK', 'CLOSING'
+        true
+      else
+        false
+      end
+    end
+  end
+
+  def write_http_header(socket) # :nodoc:
+    if headers?
+      self.response_start_sent = true
+      HTTP_RESPONSE_START.each { |c| socket.write(c) }
+    end
+  end
 end
diff --git a/lib/unicorn/socket_helper.rb b/lib/unicorn/socket_helper.rb
index 7aa2bb0..f52dde2 100644
--- a/lib/unicorn/socket_helper.rb
+++ b/lib/unicorn/socket_helper.rb
@@ -3,6 +3,18 @@
 require 'socket'
 
 module Unicorn
+
+  # Instead of using a generic Kgio::Socket for everything,
+  # tag TCP sockets so we can use TCP_INFO under Linux without
+  # incurring extra syscalls for Unix domain sockets.
+  # TODO: remove these when we remove kgio
+  TCPClient = Class.new(Kgio::Socket) # :nodoc:
+  class TCPSrv < Kgio::TCPServer # :nodoc:
+    def kgio_tryaccept # :nodoc:
+      super(TCPClient)
+    end
+  end
+
   module SocketHelper
 
     # internal interface
@@ -151,7 +163,7 @@ module Unicorn
       end
       sock.bind(Socket.pack_sockaddr_in(port, addr))
       sock.autoclose = false
-      Kgio::TCPServer.for_fd(sock.fileno)
+      TCPSrv.for_fd(sock.fileno)
     end
 
     # returns rfc2732-style (e.g. "[::1]:666") addresses for IPv6
@@ -188,7 +200,7 @@ module Unicorn
     def server_cast(sock)
       begin
         Socket.unpack_sockaddr_in(sock.getsockname)
-        Kgio::TCPServer.for_fd(sock.fileno)
+        TCPSrv.for_fd(sock.fileno)
       rescue ArgumentError
         Kgio::UNIXServer.for_fd(sock.fileno)
       end
diff --git a/test/unit/test_ccc.rb b/test/unit/test_ccc.rb
new file mode 100644
index 0000000..22b1a9c
--- /dev/null
+++ b/test/unit/test_ccc.rb
@@ -0,0 +1,81 @@
+require 'socket'
+require 'unicorn'
+require 'io/wait'
+require 'tempfile'
+require 'test/unit'
+
+class TestCccTCPI < Test::Unit::TestCase
+  def test_ccc_tcpi
+    start_pid = $$
+    host = '127.0.0.1'
+    srv = TCPServer.new(host, 0)
+    port = srv.addr[1]
+    err = Tempfile.new('unicorn_ccc')
+    rd, wr = IO.pipe
+    pid = fork do
+      reqs = 0
+      rd.close
+      worker_pid = nil
+      app = lambda do |env|
+        worker_pid ||= begin
+          at_exit { wr.write(reqs.to_s) if worker_pid == $$ }
+          $$
+        end
+        reqs += 1
+        sleep(1) if env['PATH_INFO'] == '/sleep'
+        [ 200, [ %w(Content-Length 0),  %w(Content-Type text/plain) ], [] ]
+      end
+      ENV['UNICORN_FD'] = srv.fileno.to_s
+      opts = {
+        listeners: [ "#{host}:#{port}" ],
+        stderr_path: err.path,
+        check_client_connection: true,
+      }
+      uni = Unicorn::HttpServer.new(app, opts)
+      uni.start.join
+    end
+    wr.close
+
+    # make sure the server is running, at least
+    client = TCPSocket.new(host, port)
+    client.write("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
+    assert client.wait_readable(10), 'never got response from server'
+    res = client.read
+    assert_match %r{\AHTTP/1\.1 200}, res, 'got part of first response'
+    assert_match %r{\r\n\r\n\z}, res, 'got end of response, server is ready'
+    client.close
+
+    # start a slow request...
+    sleeper = TCPSocket.new(host, port)
+    sleeper.write("GET /sleep HTTP/1.1\r\nHost: example.com\r\n\r\n")
+
+    # and a bunch of aborted ones
+    nr = 100
+    nr.times do |i|
+      client = TCPSocket.new(host, port)
+      client.write("GET /collections/#{rand(10000)} HTTP/1.1\r\n" \
+                   "Host: example.com\r\n\r\n")
+      client.close
+    end
+    sleeper.close
+    kpid = pid
+    pid = nil
+    Process.kill(:QUIT, kpid)
+    _, status = Process.waitpid2(kpid)
+    assert status.success?
+    reqs = rd.read.to_i
+    warn "server got #{reqs} requests with #{nr} CCC aborted\n" if $DEBUG
+    assert_operator reqs, :<, nr
+    assert_operator reqs, :>=, 2, 'first 2 requests got through, at least'
+  ensure
+    return if start_pid != $$
+    srv.close if srv
+    if pid
+      Process.kill(:QUIT, pid)
+      _, status = Process.waitpid2(pid)
+      assert status.success?
+    end
+    err.close! if err
+    rd.close if rd
+  end
+end