From 816c67d2e92732407d5ed5998cc57c52bf106892 Mon Sep 17 00:00:00 2001
From: Evan Phoenix <evan@fallingsnow.net>
Date: Mon, 15 Jul 2013 14:29:10 -0700
Subject: [PATCH] Add ability to drain accept socket on shutdown

---
 lib/puma/configuration.rb |  8 ++++++++
 lib/puma/events.rb        |  6 ++++++
 lib/puma/runner.rb        |  2 +-
 lib/puma/server.rb        | 26 +++++++++++++++++++++++++-
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
index 1c1538bb..0648b316 100644
--- a/lib/puma/configuration.rb
+++ b/lib/puma/configuration.rb
@@ -194,6 +194,14 @@ module Puma
         @options[:daemon] = which
       end
 
+      # When shutting down, drain the accept socket of pending
+      # connections and proces them. This loops over the accept
+      # socket until there are no more read events and then stops
+      # looking and waits for the requests to finish.
+      def drain_on_shutdown(which=true)
+        @options[:drain_on_shutdown] = which
+      end
+
       # Set the environment in which the Rack's app will run.
       def environment(environment)
         @options[:environment] = environment
diff --git a/lib/puma/events.rb b/lib/puma/events.rb
index 734084f9..994084ef 100644
--- a/lib/puma/events.rb
+++ b/lib/puma/events.rb
@@ -20,6 +20,8 @@ module Puma
       @stdout.sync = true
       @stderr.sync = true
 
+      @debug = ENV.key? 'PUMA_DEBUG'
+
       @on_booted = []
     end
 
@@ -35,6 +37,10 @@ module Puma
       @stdout.write str
     end
 
+    def debug(str)
+      log("% #{str}") if @debug
+    end
+
     # Write +str+ to +@stderr+
     #
     def error(str)
diff --git a/lib/puma/runner.rb b/lib/puma/runner.rb
index 29553cb5..a7b03b0a 100644
--- a/lib/puma/runner.rb
+++ b/lib/puma/runner.rb
@@ -115,7 +115,7 @@ module Puma
       min_t = @options[:min_threads]
       max_t = @options[:max_threads]
 
-      server = Puma::Server.new app, @cli.events
+      server = Puma::Server.new app, @cli.events, @options
       server.min_threads = min_t
       server.max_threads = max_t
       server.inherit_binder @cli.binder
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index 9e5698a9..a29bb414 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -46,7 +46,7 @@ module Puma
     # Server#run returns a thread that you can join on to wait for the server
     # to do it's work.
     #
-    def initialize(app, events=Events.stdio)
+    def initialize(app, events=Events.stdio, options={})
       @app = app
       @events = events
 
@@ -70,6 +70,8 @@ module Puma
 
       @leak_stack_on_error = true
 
+      @options = options
+
       ENV['RACK_ENV'] ||= "development"
     end
 
@@ -600,6 +602,28 @@ module Puma
     # Wait for all outstanding requests to finish.
     #
     def graceful_shutdown
+      if @options[:drain_on_shutdown]
+        count = 0
+
+        while true
+          ios = IO.select @binder.ios, nil, nil, 0
+          break unless ios
+
+          ios.first.each do |sock|
+            begin
+              if io = sock.accept_nonblock
+                count += 1
+                c = Client.new io, @binder.env(sock)
+                @thread_pool << c
+              end
+            rescue SystemCallError
+            end
+          end
+        end
+
+        @events.debug "Drained #{count} additional connections."
+      end
+
       @thread_pool.shutdown if @thread_pool
     end