From 9a9bd739e555771cacfd5b3757a251754ef9512b Mon Sep 17 00:00:00 2001
From: Simon Eskildsen <simon.eskildsen@shopify.com>
Date: Mon, 6 Mar 2017 16:32:02 -0500
Subject: check_client_connection: use tcp state on linux

* Use a frozen empty array and a class variable for TCP_Info to avoid
  garbage. As far as I can tell, this shouldn't result in any garbage on
  any requests (other than on the first request).
* Pass listener socket to #read to only check the client connection on
  a TCP server.
* Short circuit CLOSE_WAIT after ESTABLISHED since in my testing it's
  the most common state after ESTABLISHED, it makes the numbers
  un-ordered, though. But comment should make it OK.
* Definition of of `check_client_connection` based on whether
  Raindrops::TCP_Info is defined, instead of the class variable
  approach.
* Changed the unit tests to pass a `nil` listener.

Tested on our staging environment, and still works like a dream.

I should note that I got the idea between this patch into Puma as well!

https://github.com/puma/puma/pull/1227

[ew: squashed in temporary change for oob_gc.rb, but we'll come
 up with a different change to avoid breaking gctools
 <https://github.com/tmm1/gctools>]

Acked-by: Eric Wong <e@80x24.org>
---
 lib/unicorn/http_request.rb | 45 +++++++++++++++++++++++++++++++++++++++------
 lib/unicorn/http_server.rb  |  6 +++---
 lib/unicorn/oob_gc.rb       |  4 ++--
 3 files changed, 44 insertions(+), 11 deletions(-)

(limited to 'lib/unicorn')

diff --git a/lib/unicorn/http_request.rb b/lib/unicorn/http_request.rb
index c176083..9acde50 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,6 +26,7 @@ class Unicorn::HttpParser
 
   # :stopdoc:
   HTTP_RESPONSE_START = [ 'HTTP'.freeze, '/1.1 '.freeze ]
+  EMPTY_ARRAY = [].freeze
   @@input_class = Unicorn::TeeInput
   @@check_client_connection = false
 
@@ -59,7 +61,7 @@ class Unicorn::HttpParser
   # returns an environment hash suitable for Rack if successful
   # This does minimal exception trapping and it is up to the caller
   # to handle any socket errors (e.g. user aborted upload).
-  def read(socket)
+  def read(socket, listener)
     clear
     e = env
 
@@ -80,11 +82,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, listener) if @@check_client_connection
 
     e['rack.input'] = 0 == content_length ?
                       NULL_IO : @@input_class.new(socket, self)
@@ -105,4 +103,39 @@ class Unicorn::HttpParser
   def hijacked?
     env.include?('rack.hijack_io'.freeze)
   end
+
+  if defined?(Raindrops::TCP_Info)
+    def check_client_connection(socket, listener) # :nodoc:
+      if Kgio::TCPServer === listener
+        @@tcp_info ||= Raindrops::TCP_Info.new(socket)
+        @@tcp_info.get!(socket)
+        raise Errno::EPIPE, "client closed connection".freeze,
+              EMPTY_ARRAY if closed_state?(@@tcp_info.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
+    def check_client_connection(socket, listener) # :nodoc:
+      write_http_header(socket)
+    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/http_server.rb b/lib/unicorn/http_server.rb
index c2086cb..2aa1072 100644
--- a/lib/unicorn/http_server.rb
+++ b/lib/unicorn/http_server.rb
@@ -558,8 +558,8 @@ class Unicorn::HttpServer
 
   # once a client is accepted, it is processed in its entirety here
   # in 3 easy steps: read request, call app, write app response
-  def process_client(client)
-    status, headers, body = @app.call(env = @request.read(client))
+  def process_client(client, listener)
+    status, headers, body = @app.call(env = @request.read(client, listener))
 
     begin
       return if @request.hijacked?
@@ -655,7 +655,7 @@ class Unicorn::HttpServer
         # Unicorn::Worker#kgio_tryaccept is not like accept(2) at all,
         # but that will return false
         if client = sock.kgio_tryaccept
-          process_client(client)
+          process_client(client, sock)
           nr += 1
           worker.tick = time_now.to_i
         end
diff --git a/lib/unicorn/oob_gc.rb b/lib/unicorn/oob_gc.rb
index 5572e59..74a1d51 100644
--- a/lib/unicorn/oob_gc.rb
+++ b/lib/unicorn/oob_gc.rb
@@ -67,8 +67,8 @@ module Unicorn::OobGC
 
   #:stopdoc:
   PATH_INFO = "PATH_INFO"
-  def process_client(client)
-    super(client) # Unicorn::HttpServer#process_client
+  def process_client(client, listener)
+    super(client, listener) # Unicorn::HttpServer#process_client
     if OOBGC_PATH =~ OOBGC_ENV[PATH_INFO] && ((@@nr -= 1) <= 0)
       @@nr = OOBGC_INTERVAL
       OOBGC_ENV.clear
-- 
cgit v1.2.3-24-ge0c7