Add ability to drain accept socket on shutdown

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

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)

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

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