Merge pull request #639 from ebeigarts/fix-phased-restarts

Fix phased restart with worker shutdown timeout

.gitignore

@@ -13,3 +13,6 @@ t/
 .rbx/
 Gemfile.lock
 .idea/
+/test/test_puma.state
+/test/test_server.sock
+/test/test_control.sock

lib/puma/cli.rb

@@ -491,7 +491,7 @@ def run
       set_rack_environment
 
       if clustered?
-        @events = PidEvents.new STDOUT, STDERR
+        @events.formatter = Events::PidFormatter.new
         @options[:logger] = @events
 
         @runner = Cluster.new(self)

lib/puma/cluster.rb

@@ -43,12 +43,14 @@ def redirect_io
     end
 
     class Worker
-      def initialize(idx, pid, phase)
+      def initialize(idx, pid, phase, options)
         @index = idx
         @pid = pid
         @phase = phase
         @stage = :started
         @signal = "TERM"
+        @options = options
+        @first_term_sent = nil
         @last_checkin = Time.now
       end
 
@@ -105,7 +107,7 @@ def spawn_workers
 
         pid = fork { worker(idx, master) }
         @cli.debug "Spawned worker: #{pid}"
-        @workers << Worker.new(idx, pid, @phase)
+        @workers << Worker.new(idx, pid, @phase, @options)
         @options[:after_worker_boot].each { |h| h.call }
       end
 

lib/puma/events.rb

@@ -8,12 +8,24 @@ module Puma
   # The methods available are the events that the Server fires.
   #
   class Events
+    class DefaultFormatter
+      def call(str)
+        str
+      end
+    end
+
+    class PidFormatter
+      def call(str)
+        "[#{$$}] #{str}"
+      end
+    end
 
     include Const
 
     # Create an Events object that prints to +stdout+ and +stderr+.
     #
     def initialize(stdout, stderr)
+      @formatter = DefaultFormatter.new
       @stdout = stdout
       @stderr = stderr
 
@@ -28,6 +40,7 @@ def initialize(stdout, stderr)
     end
 
     attr_reader :stdout, :stderr
+    attr_accessor :formatter
 
     # Fire callbacks for the named hook
     #
@@ -52,11 +65,11 @@ def register(hook, obj=nil, &blk)
     # Write +str+ to +@stdout+
     #
     def log(str)
-      @stdout.puts str
+      @stdout.puts format(str)
     end
 
     def write(str)
-      @stdout.write str
+      @stdout.write format(str)
     end
 
     def debug(str)
@@ -66,10 +79,14 @@ def debug(str)
     # Write +str+ to +@stderr+
     #
     def error(str)
-      @stderr.puts "ERROR: #{str}"
+      @stderr.puts format("ERROR: #{str}")
       exit 1
     end
 
+    def format(str)
+      formatter.call(str)
+    end
+
     # An HTTP parse error has occured.
     # +server+ is the Server object, +env+ the request, and +error+ a
     # parsing exception.
@@ -113,18 +130,4 @@ def self.stdio
       Events.new $stdout, $stderr
     end
   end
-
-  class PidEvents < Events
-    def log(str)
-      super "[#{$$}] #{str}"
-    end
-
-    def write(str)
-      super "[#{$$}] #{str}"
-    end
-
-    def error(str)
-      super "[#{$$}] #{str}"
-    end
-  end
 end

test/hello-stuck.ru

@@ -0,0 +1 @@
+run lambda { |env| sleep 60; [200, {"Content-Type" => "text/plain"}, ["Hello World"]] }

test/test_integration.rb

@@ -102,6 +102,47 @@ def test_stop_via_pumactl
     assert_kind_of Thread, t.join(1), "server didn't stop"
   end
 
+  def test_phased_restart_via_pumactl
+    if defined?(JRUBY_VERSION) || RbConfig::CONFIG["host_os"] =~ /mingw|mswin/
+      assert true
+      return
+    end
+
+    cli = Puma::CLI.new %W!-q -S #{@state_path} -b unix://#{@bind_path} --control unix://#{@control_path} -w 2 test/hello-stuck.ru!, @events
+    cli.options[:worker_shutdown_timeout] = 1
+
+    t = Thread.new do
+      cli.run
+    end
+
+    wait_booted
+
+    # Make both workers stuck
+    s1 = UNIXSocket.new @bind_path
+    s1 << "GET / HTTP/1.0\r\n\r\n"
+    s2 = UNIXSocket.new @bind_path
+    s2 << "GET / HTTP/1.0\r\n\r\n"
+
+    sout = StringIO.new
+
+    # Phased restart
+    ccli = Puma::ControlCLI.new %W!-S #{@state_path} phased-restart!, sout
+    ccli.run
+    sleep 20
+    @events.stdout.rewind
+    log = @events.stdout.readlines.join("")
+    assert_match(/TERM sent/, log)
+    assert_match(/KILL sent/, log)
+    assert_match(/Worker 0 \(pid: \d+\) booted, phase: 1/, log)
+    assert_match(/Worker 1 \(pid: \d+\) booted, phase: 1/, log)
+
+    # Stop
+    ccli = Puma::ControlCLI.new %W!-S #{@state_path} stop!, sout
+    ccli.run
+
+    assert_kind_of Thread, t.join(5), "server didn't stop"
+  end
+
   def notest_restart_closes_keepalive_sockets
     server("-q test/hello.ru")