einhorn 0.7.4 → 0.8.2

data/lib/einhorn/command/interface.rb

@@ -289,6 +289,17 @@ EOF
       nil
     end
 
+    command 'worker:ping' do |conn, request|
+      if pid = request['pid']
+        Einhorn::Command.register_ping(pid, request['request_id'])
+      else
+        conn.log_error("Invalid request (no pid): #{request.inspect}")
+      end
+      # Throw away this connection in case the application forgets to
+      conn.close
+      nil
+    end
+
     # Used by einhornsh
     command 'ehlo' do |conn, request|
       <<EOF

data/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 @@ module Einhorn
       if expires_at = @@timers.keys.sort[0]
         expires_at - Time.now
       else
-        nil
+        @@default_timeout
       end
     end
 
@@ -165,6 +166,14 @@ module Einhorn
         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
 

data/lib/einhorn/prctl.rb

@@ -0,0 +1,26 @@
+module Einhorn
+  class PrctlAbstract
+    def self.get_pdeathsig
+      raise NotImplementedError
+    end
+
+    def self.set_pdeathsig(signal)
+      raise NotImplementedError
+    end
+  end
+
+  class PrctlUnimplemented < PrctlAbstract
+    # Deliberately empty; NotImplementedError is intended
+  end
+
+  if RUBY_PLATFORM =~ /linux/ then
+    begin
+      require 'einhorn/prctl_linux'
+      Prctl = PrctlLinux
+    rescue LoadError
+      Prctl = PrctlUnimplemented
+    end
+  else
+    Prctl = PrctlUnimplemented
+  end
+end

data/lib/einhorn/prctl_linux.rb

@@ -0,0 +1,49 @@
+require 'fiddle'
+require 'fiddle/import'
+
+module Einhorn
+  module PrctlRaw
+    extend Fiddle::Importer
+    dlload Fiddle.dlopen(nil) # libc
+    extern 'int prctl(int, unsigned long, unsigned long, unsigned long, unsigned long)'
+
+    # From linux/prctl.h
+    SET_PDEATHSIG = 1
+    GET_PDEATHSIG = 2
+  end
+
+  class PrctlLinux < PrctlAbstract
+    # Reading integers is hard with fiddle. :(
+    IntStruct = Fiddle::CStructBuilder.create(Fiddle::CStruct, [Fiddle::TYPE_INT], ['i'])
+
+    def self.get_pdeathsig
+      out = IntStruct.malloc
+      out.i = 0
+      if PrctlRaw.prctl(PrctlRaw::GET_PDEATHSIG, out.to_i, 0, 0, 0) != 0 then
+        raise SystemCallError.new("get_pdeathsig", Fiddle.last_error)
+      end
+
+      signo = out.i
+      if signo == 0 then
+        return nil
+      end
+
+      return Signal.signame(signo)
+    end
+
+    def self.set_pdeathsig(signal)
+      case
+      when signal == nil
+        signo = 0
+      when signal.instance_of?(String)
+        signo = Signal.list.fetch(signal)
+      else
+        signo = signal
+      end
+
+      if PrctlRaw.prctl(PrctlRaw::SET_PDEATHSIG, signo, 0, 0, 0) != 0 then
+        raise SystemCallError.new("set_pdeathsig(#{signal})", Fiddle.last_error)
+      end
+    end
+  end
+end

data/lib/einhorn/version.rb

@@ -1,3 +1,3 @@
 module Einhorn
-  VERSION = '0.7.4'
+  VERSION = '0.8.2'
 end

data/lib/einhorn/worker.rb

@@ -66,33 +66,28 @@ module Einhorn
     # be useful for anything. Maybe if it's always fd 3, because then
     # the user wouldn't have to provide an arg.
     def self.ack!(discovery=:env, arg=nil)
-      ensure_worker!
-      close_after_use = true
-
-      case discovery
-      when :env
-        socket = ENV['EINHORN_SOCK_PATH']
-        client = Einhorn::Client.for_path(socket)
-      when :fd
-        raise "No EINHORN_SOCK_FD provided in environment. Did you run einhorn with the -g flag?" unless fd_str = ENV['EINHORN_SOCK_FD']
-
-        fd = Integer(fd_str)
-        client = Einhorn::Client.for_fd(fd)
-        close_after_use = false if arg
-      when :direct
-        socket = arg
-        client = Einhorn::Client.for_path(socket)
-      else
-        raise "Unrecognized socket discovery mechanism: #{discovery.inspect}. Must be one of :filesystem, :argv, or :direct"
+      handle_command_socket(discovery, arg) do |client|
+        client.send_command('command' => 'worker:ack', 'pid' => $$)
       end
+    end
 
-      client.send_command({
-          'command' => 'worker:ack',
-          'pid' => $$
-        })
-
-      client.close if close_after_use
-      true
+    # Call this to indicate your child process is up and in a healthy state.
+    # Arguments:
+    #
+    # @request_id: Identifies the request ID of the worker, can be used to debug wedged workers.
+    #
+    # @discovery: How to discover the master process's command socket.
+    #   :env:        Discover the path from ENV['EINHORN_SOCK_PATH']
+    #   :fd:         Just use the file descriptor in ENV['EINHORN_SOCK_FD'].
+    #                Must run the master with the -g flag. This is mostly
+    #                useful if you don't have a nice library like Einhorn::Worker.
+    #                Then @arg being true causes the FD to be left open after ACK;
+    #                otherwise it is closed.
+    #   :direct:     Provide the path to the command socket in @arg.
+    def self.ping!(request_id, discovery=:env, arg=nil)
+      handle_command_socket(discovery, arg) do |client|
+        client.send_command('command' => 'worker:ping', 'pid' => $$, 'request_id' => request_id)
+      end
     end
 
     def self.socket(number=nil)
@@ -139,6 +134,33 @@ module Einhorn
 
     private
 
+    def self.handle_command_socket(discovery, contextual_arg)
+      ensure_worker!
+      close_after_use = true
+
+      case discovery
+      when :env
+        socket = ENV['EINHORN_SOCK_PATH']
+        client = Einhorn::Client.for_path(socket)
+      when :fd
+        raise "No EINHORN_SOCK_FD provided in environment. Did you run einhorn with the -g flag?" unless fd_str = ENV['EINHORN_SOCK_FD']
+
+        fd = Integer(fd_str)
+        client = Einhorn::Client.for_fd(fd)
+        close_after_use = false if contextual_arg
+      when :direct
+        socket = contextual_arg
+        client = Einhorn::Client.for_path(socket)
+      else
+        raise "Unrecognized socket discovery mechanism: #{discovery.inspect}. Must be one of :filesystem, :argv, or :direct"
+      end
+
+      yield client
+      client.close if close_after_use
+
+      true
+    end
+
     def self.socket_from_filesystem(cmd_name)
       ppid = Process.ppid
       socket_path_file = Einhorn::Command::Interface.socket_path_file(ppid)

data/test/integration/_lib/fixtures/exit_during_upgrade/exiting_server.rb

@@ -5,6 +5,7 @@ require 'einhorn/worker'
 def einhorn_main
   serv = Socket.for_fd(Einhorn::Worker.socket!)
   Einhorn::Worker.ack!
+  Einhorn::Worker.ping!("id-1")
 
   Signal.trap('USR2') do
     sleep 3

data/test/integration/_lib/fixtures/pdeathsig_printer/pdeathsig_printer.rb

@@ -0,0 +1,29 @@
+require 'bundler/setup'
+require 'socket'
+require 'einhorn/worker'
+require 'einhorn/prctl'
+
+def einhorn_main
+  serv = Socket.for_fd(Einhorn::Worker.socket!)
+
+  Signal.trap("USR2") { exit }
+
+  begin
+    output = Einhorn::Prctl.get_pdeathsig
+    if output == nil then
+      output = "nil"
+    end
+  rescue NotImplementedError
+    output = "not implemented"
+  end
+
+  Einhorn::Worker.ack!
+  while true
+    s, _ = serv.accept
+    s.write(output)
+    s.flush
+    s.close
+  end
+end
+
+einhorn_main if $0 == __FILE__

data/test/integration/_lib/fixtures/signal_timeout/sleepy_server.rb

@@ -0,0 +1,23 @@
+require 'bundler/setup'
+require 'socket'
+require 'einhorn/worker'
+
+def einhorn_main
+  serv = Socket.for_fd(Einhorn::Worker.socket!)
+  Einhorn::Worker.ack!
+  Einhorn::Worker.ping!("id-1")
+
+  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__

data/test/integration/_lib/fixtures/upgrade_project/upgrading_server.rb

@@ -9,6 +9,8 @@ def einhorn_main
   $stderr.puts "Worker has a socket"
   Einhorn::Worker.ack!
   $stderr.puts "Worker sent ack to einhorn"
+  Einhorn::Worker.ping!("id-1")
+  $stderr.puts "Worker has sent a ping to einhorn"
   while true
     s, addrinfo = serv.accept
     $stderr.puts "Worker got a socket!"

data/test/integration/_lib/helpers/einhorn_helpers.rb

@@ -106,6 +106,11 @@ module Helpers
       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

data/test/integration/pdeathsig.rb

@@ -0,0 +1,26 @@
+require(File.expand_path('_lib', File.dirname(__FILE__)))
+
+class PdeathsigTest < EinhornIntegrationTestCase
+  include Helpers::EinhornHelpers
+
+  describe 'when run with -k' do
+    before do
+      @dir = prepare_fixture_directory('pdeathsig_printer')
+      @port = find_free_port
+      @server_program = File.join(@dir, 'pdeathsig_printer.rb')
+      @socket_path = File.join(@dir, 'einhorn.sock')
+    end
+    after { cleanup_fixtured_directories }
+
+    it 'sets pdeathsig to USR2 in the child process' do
+      with_running_einhorn(%W{einhorn -m manual -b 127.0.0.1:#{@port} -d #{@socket_path} -k -- ruby #{@server_program}}) do |process|
+        wait_for_open_port
+        output = read_from_port
+        if output != "not implemented" then
+          assert_equal("USR2", output)
+        end
+        process.terminate
+      end
+    end
+  end
+end

data/test/integration/upgrading.rb

@@ -78,6 +78,10 @@ class UpgradeTests < EinhornIntegrationTestCase
           state = YAML.load(resp['message'])
           assert_equal(1, state[:state][:children].count)
 
+          child = state[:state][:children].first.last
+          assert_in_delta(child[:pinged_at], Time.now, 60)
+          assert_equal("id-1", child[:pinged_request_id])
+
           process.terminate
         end
       end
@@ -154,4 +158,47 @@ class UpgradeTests < EinhornIntegrationTestCase
       end
     end
   end
+
+  describe "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

data/test/unit/_lib/bad_worker.rb

@@ -0,0 +1,7 @@
+#!/usr/bin/env ruby
+
+raise 'This worker has an error during preload'
+
+def einhorn_main
+
+end

data/test/unit/_lib/sleep_worker.rb

@@ -0,0 +1,5 @@
+#!/usr/bin/env ruby
+
+def einhorn_main
+  sleep 1
+end

data/test/unit/einhorn.rb

@@ -10,7 +10,7 @@ class EinhornTest < EinhornTestCase
 
     it "correctly parses srv: arguments" do
       cmd = ['foo', 'srv:1.2.3.4:123,llama,test', 'bar']
-      Einhorn.expects(:bind).once.with('1.2.3.4', '123', ['llama', 'test']).returns(4)
+      Einhorn.expects(:bind).once.with('1.2.3.4', '123', ['llama', 'test']).returns([4, 10087])
 
       Einhorn.socketify!(cmd)
 
@@ -19,7 +19,7 @@ class EinhornTest < EinhornTestCase
 
     it "correctly parses --opt=srv: arguments" do
       cmd = ['foo', '--opt=srv:1.2.3.4:456', 'baz']
-      Einhorn.expects(:bind).once.with('1.2.3.4', '456', []).returns(5)
+      Einhorn.expects(:bind).once.with('1.2.3.4', '456', []).returns([5, 10088])
 
       Einhorn.socketify!(cmd)
 
@@ -28,7 +28,7 @@ class EinhornTest < EinhornTestCase
 
     it "uses the same fd number for the same server spec" do
       cmd = ['foo', '--opt=srv:1.2.3.4:8910', 'srv:1.2.3.4:8910']
-      Einhorn.expects(:bind).once.with('1.2.3.4', '8910', []).returns(10)
+      Einhorn.expects(:bind).once.with('1.2.3.4', '8910', []).returns([10, 10089])
 
       Einhorn.socketify!(cmd)
 
@@ -55,4 +55,42 @@ class EinhornTest < EinhornTestCase
       assert(message.nil?)
     end
   end
+
+  describe ".preload" do
+    before do
+      Einhorn::State.preloaded = false
+    end
+
+    it "updates preload on success" do
+      Einhorn.stubs(:set_argv).returns
+      # preloads the sleep worker since it has einhorn main
+      Einhorn::State.path = "#{__dir__}/_lib/sleep_worker.rb"
+      assert_equal(false, Einhorn::State.preloaded)
+      Einhorn.preload
+      assert_equal(true, Einhorn::State.preloaded)
+      # Attempt another preload 
+      Einhorn.preload
+      assert_equal(true, Einhorn::State.preloaded)
+    end
+
+    it "updates preload to failed with previous success" do
+      Einhorn.stubs(:set_argv).returns
+      Einhorn::State.path = "#{__dir__}/_lib/sleep_worker.rb"
+      assert_equal(false, Einhorn::State.preloaded)
+      Einhorn.preload
+      assert_equal(true, Einhorn::State.preloaded)
+      # Change path to bad worker and preload again, should be false
+      Einhorn::State.path = "#{__dir__}/_lib/bad_worker.rb"
+      Einhorn.preload
+      assert_equal(false, Einhorn::State.preloaded)
+    end
+
+    it "preload is false after failing" do
+      Einhorn.stubs(:set_argv).returns
+      Einhorn::State.path = "#{__dir__}/bad_worker.rb"
+      assert_equal(false, Einhorn::State.preloaded)
+      Einhorn.preload
+      assert_equal(false, Einhorn::State.preloaded)
+    end
+  end
 end