contribsys · dfradette-stripe · Jun 22, 2018 · Feb 16, 2018 · Apr 18, 2018 · asf-stripe
diff --git a/lib/einhorn.rb b/lib/einhorn.rb
@@ -411,6 +411,14 @@ def self.run
       Einhorn::State.reloading_for_upgrade = false
     end
 
+    # If setting a signal-timeout, timeout the event loop
+    # in the same timeframe, ensuring processes are culled
+    # on a regular basis.
+    if Einhorn::State.signal_timeout
+      Einhorn::Event.default_timeout = Einhorn::Event.default_timeout.nil? ?
+        Einhorn::State.signal_timeout : [Einhorn::State.signal_timeout, Einhorn::Event.default_timeout].min
+    end
+
     while Einhorn::State.respawn || Einhorn::State.children.size > 0
       log_debug("Entering event loop")
 

diff --git a/lib/einhorn/command.rb b/lib/einhorn/command.rb
@@ -11,7 +11,6 @@ def self.reap
       begin
         while true
           Einhorn.log_debug('Going to reap a child process')
-
           pid = Process.wait(-1, Process::WNOHANG)
           return unless pid
           mourn(pid)
@@ -101,8 +100,8 @@ def self.register_ack(pid)
 
     def self.signal_all(signal, children=nil, record=true)
       children ||= Einhorn::WorkerPool.workers
-
       signaled = {}
+
       Einhorn.log_info("Sending #{signal} to #{children.inspect}", :upgrade)
 
       children.each do |child|
@@ -116,11 +115,13 @@ def self.signal_all(signal, children=nil, record=true)
             Einhorn.log_error("Re-sending #{signal} to already-signaled child #{child.inspect}. It may be slow to spin down, or it may be swallowing #{signal}s.", :upgrade)
           end
           spec[:signaled].add(signal)
+          spec[:last_signaled_at] = Time.now
         end
 
         begin
           Process.kill(signal, child)
         rescue Errno::ESRCH
+          Einhorn.log_debug("Attempted to #{signal} child #{child.inspect} but the process does not exist", :upgrade)
         else
           signaled[child] = spec
         end
@@ -146,11 +147,12 @@ def self.signal_all(signal, children=nil, record=true)
             spec[:signaled].add('KILL')
           end
         end
-      end
 
-      "Successfully sent #{signal}s to #{signaled.length} processes: #{signaled.keys}"
+        Einhorn.log_info("Successfully sent #{signal}s to #{signaled.length} processes: #{signaled.keys}")
+      end
     end
 
+
     def self.increment
       Einhorn::Event.break_loop
       old = Einhorn::State.config[:number]
@@ -314,6 +316,7 @@ def self.spinup(cmd=nil)
         :version => Einhorn::State.version,
         :acked => false,
         :signaled => Set.new,
+        :last_signaled_at => nil,
         :index => index,
         :spinup_time => Einhorn::State.last_spinup,
       }
@@ -488,6 +491,41 @@ def self.cull
         Einhorn.log_info("Have too many workers at the current version, so killing off #{excess.length} of them.")
         signal_all("USR2", excess)
       end
+
+      # Ensure all signaled workers that have outlived signal_timeout get killed.
+      kill_expired_signaled_workers if Einhorn::State.signal_timeout
+    end
+
+    def self.kill_expired_signaled_workers
+      now = Time.now
+      children = Einhorn::State.children.select do |_,c|
+        # Only interested in USR2 signaled workers
+        next unless c[:signaled] && c[:signaled].length > 0
+        next unless c[:signaled].include?('USR2')
+
+        # Ignore processes that have received KILL since it can't be trapped.
+        next if c[:signaled].include?('KILL')
+
+        # Filter out those children that have not reached signal_timeout yet.
+        next unless c[:last_signaled_at]
+        expires_at = c[:last_signaled_at] + Einhorn::State.signal_timeout
+        next unless now >= expires_at
+
+        true
+      end
+
+      Einhorn.log_info("#{children.size} expired signaled workers found.") if children.size > 0
+      children.each do |pid, child|
+        Einhorn.log_info("Child #{pid.inspect} was signaled #{(child[:last_signaled_at] - now).abs.to_i}s ago. Sending SIGKILL as it is still active after #{Einhorn::State.signal_timeout}s timeout.", :upgrade)
+        begin
+          Process.kill('KILL', pid)
+        rescue Errno::ESRCH
+          Einhorn.log_debug("Attempted to SIGKILL child #{pid.inspect} but the process does not exist.")
+        end
+
+        child[:signaled].add('KILL')
+        child[:last_signaled_at] = Time.now
+      end
     end
 
     def self.stop_respawning

diff --git a/lib/einhorn/event.rb b/lib/einhorn/event.rb
@@ -4,6 +4,7 @@ module Einhorn
   module Event
     @@loopbreak_reader = nil
     @@loopbreak_writer = nil
+    @@default_timeout = nil
     @@signal_actions = []
     @@readable = {}
     @@writeable = {}
@@ -120,7 +121,7 @@ def self.timeout
       if expires_at = @@timers.keys.sort[0]
         expires_at - Time.now
       else
-        nil
+        @@default_timeout
       end
     end
 
@@ -165,6 +166,14 @@ def self.break_loop
         Einhorn.log_error("Loop break pipe is full -- probably means that we are quite backlogged")
       end
     end
+
+    def self.default_timeout=(val)
+      @@default_timeout = val.to_i == 0 ? nil : val.to_i
+    end
+
+    def self.default_timeout
+      @@default_timeout
+    end
   end
 end
 

diff --git a/lib/einhorn/version.rb b/lib/einhorn/version.rb
@@ -1,3 +1,3 @@
 module Einhorn
-  VERSION = '0.7.5'
+  VERSION = '0.7.6'
 end
diff --git a/test/integration/_lib/fixtures/signal_timeout/sleepy_server.rb b/test/integration/_lib/fixtures/signal_timeout/sleepy_server.rb
@@ -0,0 +1,22 @@
+require 'bundler/setup'
+require 'socket'
+require 'einhorn/worker'
+
+def einhorn_main
+  serv = Socket.for_fd(Einhorn::Worker.socket!)
+  Einhorn::Worker.ack!
+
+  Signal.trap('USR2') do
+    sleep ENV.fetch("TRAP_SLEEP").to_i
+    exit
+  end
+
+  while true
+    s, _ = serv.accept
+    s.write($$)
+    s.flush
+    s.close
+  end
+end
+
+einhorn_main if $0 == __FILE__
diff --git a/test/integration/_lib/helpers/einhorn_helpers.rb b/test/integration/_lib/helpers/einhorn_helpers.rb
@@ -106,6 +106,11 @@ def find_free_port(host='127.0.0.1')
       open_port.close
     end
 
+    def get_state(client)
+      client.send_command('command' => 'state')
+      YAML.load(client.receive_message['message'])[:state]
+    end
+
     def wait_for_open_port
       max_retries = 50
       begin

diff --git a/test/integration/upgrading.rb b/test/integration/upgrading.rb
@@ -154,4 +154,47 @@ class UpgradeTests < EinhornIntegrationTestCase
       end
     end
   end
+
+  describe 
F438
"with --signal-timeout" do
+    before do
+      @dir = prepare_fixture_directory('signal_timeout')
+      @port = find_free_port
+      @server_program = File.join(@dir, "sleepy_server.rb")
+      @socket_path = File.join(@dir, "einhorn.sock")
+    end
+
+    after { cleanup_fixtured_directories }
+
+    it 'issues a SIGKILL to outdated children when signal-timeout has passed' do
+      signal_timeout = 2
+      sleep_for = 10
+      cmd = %W{
+        einhorn
+        -b 127.0.0.1:#{@port}
+        -d #{@socket_path}
+        --signal-timeout #{signal_timeout}
+        -- ruby #{@server_program}
+      }
+
+      with_running_einhorn(cmd, env: ENV.to_h.merge({'TRAP_SLEEP' => sleep_for.to_s})) do |process|
+        wait_for_open_port
+        client = Einhorn::Client.for_path(@socket_path)
+        einhornsh(%W{-d #{@socket_path} -e upgrade})
+
+        state = get_state(client)
+        assert_equal(2, state[:children].count)
+        signaled_children = state[:children].select{|_,c| c[:signaled].length > 0}
+        assert_equal(1, signaled_children.length)
+
+        sleep(signal_timeout * 2)
+
+        state = get_state(client)
+        assert_equal(1, state[:children].count)
+        signaled_children = state[:children].select{|_,c| c[:signaled].length > 0}
+        assert_equal(0, signaled_children.length)
+
+        process.terminate
+      end
+    end
+  end
 end
diff --git a/test/unit/einhorn/command.rb b/test/unit/einhorn/command.rb
@@ -18,4 +18,36 @@ class CommandTest < EinhornTestCase
       Command.quieter
     end
   end
+
+  describe "resignal_timeout" do
+    it "does not kill any children" do
+      Einhorn::State.stubs(signal_timeout: 5 * 60)
+      Einhorn::State.stubs(children: {
+        12345 => {last_signaled_at: nil},
+        12346 => {signaled: Set.new(["USR1"]), last_signaled_at: Time.now - (2 * 60)},
+      })
+
+      Process.expects(:kill).never
+      Einhorn::Command.kill_expired_signaled_workers
+
+      refute(Einhorn::State.children[12346][:signaled].include?("KILL"), "Process was KILLed when it shouldn't have been")
+    end
+
+    it "KILLs stuck child processes" do
+      Time.stub :now, Time.at(0) do
+        Process.stub(:kill, true) do
+          Einhorn::State.stubs(signal_timeout: 60)
+          Einhorn::State.stubs(children: {
+            12346 => {signaled: Set.new(["USR2"]), last_signaled_at: Time.now - (2 * 60)},
+          })
+
+          Einhorn::Command.kill_expired_signaled_workers
+
+          child = Einhorn::State.children[12346]
+          assert(child[:signaled].include?("KILL"), "Process was not KILLed as expected")
+          assert(child[:last_signaled_at] == Time.now, "The last_processed_at was not updated as expected")
+        end
+      end
+    end
+  end
 end