From bb566d44a5073dc03f161448f311f49507de1ca6 Mon Sep 17 00:00:00 2001
From: Will Jordan <will@code.org>
Date: Thu, 15 Apr 2021 20:28:50 -0700
Subject: [PATCH] Refactor drain_on_shutdown into main server accept loop Also
 add test coverage for drain_on_shutdown option.

---
 lib/puma/queue_close.rb  | 14 +++++++-------
 lib/puma/server.rb       | 30 ++++++------------------------
 test/test_puma_server.rb | 30 ++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 31 deletions(-)

diff --git a/lib/puma/queue_close.rb b/lib/puma/queue_close.rb
index a1d568ccfc..d9a94dc665 100644
--- a/lib/puma/queue_close.rb
+++ b/lib/puma/queue_close.rb
@@ -5,22 +5,22 @@ module Puma
   # Add a simple implementation for earlier Ruby versions.
   #
   module QueueClose
-    def initialize
-      @closed = false
-      super
-    end
     def close
+      num_waiting.times {push nil}
       @closed = true
     end
     def closed?
-      @closed
+      @closed ||= false
     end
     def push(object)
-      @closed ||= false
-      raise ClosedQueueError if @closed
+      raise ClosedQueueError if closed?
       super
     end
     alias << push
+    def pop(non_block=false)
+      return nil if !non_block && closed? && empty?
+      super
+    end
   end
   ::Queue.prepend QueueClose
 end
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index 1f1ffc2161..3912699cf6 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -311,6 +311,7 @@ def handle_servers
         sockets = [check] + @binder.ios
         pool = @thread_pool
         queue_requests = @queue_requests
+        drain = @options[:drain_on_shutdown] ? 0 : nil
 
         remote_addr_value = nil
         remote_addr_header = nil
@@ -322,9 +323,10 @@ def handle_servers
           remote_addr_header = @options[:remote_address_header]
         end
 
-        while @status == :run
+        while @status == :run || (drain && shutting_down?)
           begin
-            ios = IO.select sockets
+            ios = IO.select sockets, nil, nil, (shutting_down? ? 0 : nil)
+            break unless ios
             ios.first.each do |sock|
               if sock == check
                 break if handle_check
@@ -337,6 +339,7 @@ def handle_servers
                 rescue IO::WaitReadable
                   next
                 end
+                drain += 1 if shutting_down?
                 client = Client.new io, @binder.env(sock)
                 if remote_addr_value
                   client.peerip = remote_addr_value
@@ -351,6 +354,7 @@ def handle_servers
           end
         end
 
+        @events.debug "Drained #{drain} additional connections." if drain
         @events.fire :state, @status
 
         if queue_requests
@@ -553,28 +557,6 @@ def graceful_shutdown
         $stdout.syswrite "#{pid}: === End thread backtrace dump ===\n"
       end
 
-      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
-                client = Client.new io, @binder.env(sock)
-                @thread_pool << client
-              end
-            rescue SystemCallError
-            end
-          end
-        end
-
-        @events.debug "Drained #{count} additional connections."
-      end
-
       if @status != :restart
         @binder.close
       end
diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
index 9c340ee3bb..7ed629e4af 100644
--- a/test/test_puma_server.rb
+++ b/test/test_puma_server.rb
@@ -1282,4 +1282,34 @@ def test_custom_io_selector
 
     assert_equal selector.backend, backend
   end
+
+  def test_drain_on_shutdown(drain=true)
+    num_connections = 10
+
+    wait = Queue.new
+    server_run(drain_on_shutdown: drain, max_threads: 1) do
+      wait.pop
+      [200, {}, ["DONE"]]
+    end
+    connections = Array.new(num_connections) {send_http "GET / HTTP/1.0\r\n\r\n"}
+    @server.stop
+    wait.close
+    bad = 0
+    connections.each do |s|
+      begin
+        assert_match 'DONE', s.read
+      rescue Errno::ECONNRESET
+        bad += 1
+      end
+    end
+    if drain
+      assert_equal 0, bad
+    else
+      refute_equal 0, bad
+    end
+  end
+
+  def test_not_drain_on_shutdown
+    test_drain_on_shutdown false
+  end
 end