einhorn 0.3.0

data/lib/einhorn/event.rb

@@ -0,0 +1,150 @@
+require 'set'
+
+module Einhorn
+  module Event
+    @@loopbreak_reader = nil
+    @@loopbreak_writer = nil
+    @@readable = {}
+    @@writeable = {}
+    @@timers = {}
+
+    def self.init
+      readable, writeable = IO.pipe
+      @@loopbreak_reader = LoopBreaker.open(readable)
+      @@loopbreak_writer = writeable
+    end
+
+    def self.uninit
+      # These don't need to persist across Einhorn reloads, so let's not keep.
+      @@loopbreak_reader.close
+      @@loopbreak_writer.close
+    end
+
+    def self.close_all
+      uninit
+      (@@readable.values + @@writeable.values).each do |descriptors|
+        descriptors.each do |descriptor|
+          descriptor.close
+        end
+      end
+    end
+
+    def self.close_all_for_worker
+      close_all
+    end
+
+    def self.persistent_descriptors
+      descriptor_sets = @@readable.values + @@writeable.values + @@timers.values
+      descriptors = descriptor_sets.inject {|a, b| a | b}
+      descriptors.select {|descriptor| Einhorn::Event::Persistent.persistent?(descriptor)}
+    end
+
+    def self.restore_persistent_descriptors(persistent_descriptors)
+      persistent_descriptors.each do |descriptor_state|
+        Einhorn::Event::Persistent.from_state(descriptor_state)
+      end
+    end
+
+    def self.register_readable(reader)
+      @@readable[reader.to_io] ||= Set.new
+      @@readable[reader.to_io] << reader
+    end
+
+    def self.deregister_readable(reader)
+      readers = @@readable[reader.to_io]
+      readers.delete(reader)
+      @@readable.delete(reader.to_io) if readers.length == 0
+    end
+
+    def self.readable_fds
+      readers = @@readable.keys
+      Einhorn.log_debug("Readable fds are #{readers.inspect}")
+      readers
+    end
+
+    def self.register_writeable(writer)
+      @@writeable[writer.to_io] ||= Set.new
+      @@writeable[writer.to_io] << writer
+    end
+
+    def self.deregister_writeable(writer)
+      writers = @@writeable[writer.to_io]
+      writers.delete(writer)
+      @@readable.delete(writer.to_io) if writers.length == 0
+    end
+
+    def self.writeable_fds
+      writers = @@writeable.select do |io, writers|
+        writers.any? {|writer| writer.write_pending?}
+      end.map {|io, writers| io}
+      Einhorn.log_debug("Writeable fds are #{writers.inspect}")
+      writers
+    end
+
+    def self.register_timer(timer)
+      @@timers[timer.expires_at] ||= Set.new
+      @@timers[timer.expires_at] << timer
+    end
+
+    def self.deregister_timer(timer)
+      timers = @@timers[timer.expires_at]
+      timers.delete(timer)
+      @@timers.delete(timer.expires_at) if timers.length == 0
+    end
+
+    def self.loop_once
+      run_selectables
+      run_timers
+    end
+
+    def self.timeout
+      # (expires_at of the next timer) - now
+      if expires_at = @@timers.keys.sort[0]
+        expires_at - Time.now
+      else
+        nil
+      end
+    end
+
+    def self.run_selectables
+      time = timeout
+      Einhorn.log_debug("Loop timeout is #{time.inspect}")
+      # Time's already up
+      return if time && time < 0
+
+      readable, writeable, _ = IO.select(readable_fds, writeable_fds, nil, time)
+      (readable || []).each do |io|
+        @@readable[io].each {|reader| reader.notify_readable}
+      end
+
+      (writeable || []).each do |io|
+        @@writeable[io].each {|writer| writer.notify_writeable}
+      end
+    end
+
+    def self.run_timers
+      @@timers.select {|expires_at, _| expires_at <= Time.now}.each do |expires_at, timers|
+        # Going to be modifying the set, so let's dup it.
+        timers.dup.each {|timer| timer.ring!}
+      end
+    end
+
+    def self.break_loop
+      Einhorn.log_debug("Breaking the loop")
+      begin
+        @@loopbreak_writer.write_nonblock('a')
+      rescue Errno::EWOULDBLOCK, Errno::EAGAIN
+        Einhorn.log_error("Loop break pipe is full -- probably means that we are quite backlogged")
+      end
+    end
+  end
+end
+
+require 'einhorn/event/persistent'
+require 'einhorn/event/timer'
+
+require 'einhorn/event/abstract_text_descriptor'
+require 'einhorn/event/ack_timer'
+require 'einhorn/event/command_server'
+require 'einhorn/event/connection'
+require 'einhorn/event/loop_breaker'

data/lib/einhorn/version.rb

@@ -0,0 +1,3 @@
+module Einhorn
+  VERSION = '0.3.0'
+end

data/lib/einhorn/worker.rb

@@ -0,0 +1,94 @@
+require 'einhorn/client'
+require 'einhorn/command/interface'
+
+module Einhorn
+  module Worker
+    class WorkerError < RuntimeError; end
+
+    def self.is_worker?
+      begin
+        ensure_worker!
+      rescue WorkerError
+        false
+      else
+        true
+      end
+    end
+
+    def self.ensure_worker!
+      # Make sure that EINHORN_MASTER_PID is my parent
+      if ppid_s = ENV['EINHORN_MASTER_PID']
+        ppid = ppid_s.to_i
+        raise WorkerError.new("EINHORN_MASTER_PID environment variable is #{ppid_s.inspect}, but my parent's pid is #{Process.ppid.inspect}. This probably means that I am a subprocess of an Einhorn worker, but am not one myself.") unless Process.ppid == ppid
+        true
+      else
+        raise WorkerError.new("No EINHORN_MASTER_PID environment variable set. Are you running your process under Einhorn?") unless Process.ppid == ppid
+      end
+    end
+
+    def self.ack(*args)
+      begin
+        ack!(*args)
+      rescue WorkerError
+      end
+    end
+
+    # Call this once your app is up and running in a good state.
+    # Arguments:
+    #
+    # @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_FD'].
+    #                Must run the master with the -b 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.
+    #
+    # TODO: add a :fileno option? Easy to implement; not sure if it'd
+    # 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_FD provided in environment. Did you run einhorn with the -b flag?" unless fd_str = ENV['EINHORN_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"
+      end
+
+      client.command({
+          'command' => 'worker:ack',
+          'pid' => $$
+        })
+
+      client.close if close_after_use
+      true
+    end
+
+    # Call this to handle graceful shutdown requests to your app.
+    def self.graceful_shutdown(&blk)
+      Signal.trap('USR2', &blk)
+    end
+
+    private
+
+    def self.socket_from_filesystem(cmd_name)
+      ppid = Process.ppid
+      socket_path_file = Einhorn::Command::Interface.socket_path_file(ppid)
+      File.read(socket_path_file)
+    end
+  end
+end

data/lib/einhorn/worker_pool.rb

@@ -0,0 +1,56 @@
+module Einhorn
+  module WorkerPool
+    def self.unsignaled_workers
+      Einhorn::State.children.select do |pid, spec|
+        spec[:signaled].length == 0
+      end.map {|pid, _| pid}
+    end
+
+    def self.modern_workers_with_state
+      Einhorn::State.children.select do |pid, spec|
+        spec[:version] == Einhorn::State.version
+      end
+    end
+
+    def self.acked_modern_workers_with_state
+      modern_workers_with_state.select {|pid, spec| spec[:acked]}
+    end
+
+    def self.modern_workers
+      modern_workers_with_state.map {|pid, _| pid}
+    end
+
+    def self.acked_modern_workers
+      acked_modern_workers_with_state.map {|pid, _| pid}
+    end
+
+    def self.acked_unsignaled_modern_workers
+      acked_modern_workers_with_state.select do |_, spec|
+        spec[:signaled].length == 0
+      end.map {|pid, _| pid}
+    end
+
+    # Use the number of modern workers, rather than unsignaled modern
+    # workers. This means if e.g. we do bunch of decs and then incs,
+    # any workers which haven't died yet will count towards our number
+    # of workers. Since workers really should be dying shortly after
+    # they are USR2'd, that indicates a bad state and we shouldn't
+    # make it worse by spinning up more processes. Once they die,
+    # order will be restored.
+    def self.missing_worker_count
+      ack_target - modern_workers.length
+    end
+
+    def self.ack_count
+      acked_unsignaled_modern_workers.length
+    end
+
+    def self.ack_target
+      Einhorn::State.config[:number]
+    end
+
+    def self.old_workers
+      unsignaled_workers - modern_workers
+    end
+  end
+end

data/lib/einhorn.rb

@@ -0,0 +1,282 @@
+require 'fcntl'
+require 'optparse'
+require 'pp'
+require 'set'
+require 'socket'
+require 'tmpdir'
+require 'yaml'
+
+require 'rubygems'
+
+module Einhorn
+  module AbstractState
+    def default_state; raise NotImplementedError.new('Override in extended modules'); end
+    def state; @state ||= default_state; end
+    def state=(v); @state = v; end
+
+    def method_missing(name, *args)
+      if (name.to_s =~ /(.*)=$/) && state.has_key?($1.to_sym)
+        state.send(:[]=, $1.to_sym, *args)
+      elsif state.has_key?(name)
+        state[name]
+      else
+        ds = default_state
+        if ds.has_key?(name)
+          ds[name]
+        else
+          super
+        end
+      end
+    end
+  end
+
+  module State
+    extend AbstractState
+    def self.default_state
+      {
+        :children => {},
+        :config => {:number => 1, :backlog => 100, :seconds => 1},
+        :versions => {},
+        :version => 0,
+        :sockets => {},
+        :orig_cmd => nil,
+        :cmd => nil,
+        :script_name => nil,
+        :respawn => true,
+        :upgrading => false,
+        :reloading_for_preload_upgrade => false,
+        :path => nil,
+        :cmd_name => nil,
+        :verbosity => 1,
+        :generation => 0,
+        :last_spinup => nil,
+        :ack_mode => {:type => :timer, :timeout => 1},
+        :kill_children_on_exit => false,
+        :command_socket_as_fd => false,
+        :socket_path => nil,
+        :pidfile => nil,
+        :lockfile => nil
+      }
+    end
+  end
+
+  module TransientState
+    extend AbstractState
+    def self.default_state
+      {
+        :whatami => :master,
+        :preloaded => false,
+        :script_name => nil,
+        :argv => [],
+        :has_outstanding_spinup_timer => false,
+        :stateful => nil,
+        # Holds references so that the GC doesn't go and close your sockets.
+        :socket_handles => Set.new
+      }
+    end
+  end
+
+  def self.restore_state(state)
+    parsed = YAML.load(state)
+    Einhorn::State.state = parsed[:state]
+    Einhorn::Event.restore_persistent_descriptors(parsed[:persistent_descriptors])
+    # Do this after setting state so verbosity is right9
+    Einhorn.log_info("Using loaded state: #{parsed.inspect}")
+  end
+
+  def self.print_state
+    log_info(Einhorn::State.state.pretty_inspect)
+  end
+
+  def self.bind(addr, port, flags)
+    log_info("Binding to #{addr}:#{port} with flags #{flags.inspect}")
+    sd = Socket.new(Socket::AF_INET, Socket::SOCK_STREAM, 0)
+
+    if flags.include?('r') || flags.include?('so_reuseaddr')
+      sd.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, 1)
+    end
+
+    sd.bind(Socket.pack_sockaddr_in(port, addr))
+    sd.listen(Einhorn::State.config[:backlog])
+
+    if flags.include?('n') || flags.include?('o_nonblock')
+      fl = sd.fcntl(Fcntl::F_GETFL)
+      sd.fcntl(Fcntl::F_SETFL, fl | Fcntl::O_NONBLOCK)
+    end
+
+    Einhorn::TransientState.socket_handles << sd
+    sd.fileno
+  end
+
+  # Implement these ourselves so it plays nicely with state persistence
+  def self.log_debug(msg)
+    $stderr.puts("#{log_tag} DEBUG: #{msg}") if Einhorn::State.verbosity <= 0
+  end
+  def self.log_info(msg)
+    $stderr.puts("#{log_tag} INFO: #{msg}") if Einhorn::State.verbosity <= 1
+  end
+  def self.log_error(msg)
+    $stderr.puts("#{log_tag} ERROR: #{msg}") if Einhorn::State.verbosity <= 2
+  end
+
+  private
+
+  def self.log_tag
+    case whatami = Einhorn::TransientState.whatami
+    when :master
+      "[MASTER #{$$}]"
+    when :worker
+      "[WORKER #{$$}]"
+    when :state_passer
+      "[STATE_PASSER #{$$}]"
+    else
+      "[UNKNOWN (#{whatami.inspect}) #{$$}]"
+    end
+  end
+
+  public
+
+  def self.which(cmd)
+    if cmd.include?('/')
+      return cmd if File.exists?(cmd)
+      raise "Could not find #{cmd}"
+    else
+      ENV['PATH'].split(':').each do |f|
+        abs = File.join(f, cmd)
+        return abs if File.exists?(abs)
+      end
+      raise "Could not find #{cmd} in PATH"
+    end
+  end
+
+  # Not really a thing, but whatever.
+  def self.is_script(file)
+    File.open(file) do |f|
+      bytes = f.read(2)
+      bytes == '#!'
+    end
+  end
+
+  def self.preload
+    if path = Einhorn::State.path
+      set_argv(Einhorn::State.cmd, false)
+
+      begin
+        # If it's not going to be requireable, then load it.
+        if !path.end_with?('.rb') && File.exists?(path)
+          log_info("Loading #{path} (if this hangs, make sure your code can be properly loaded as a library)")
+          load path
+        else
+          log_info("Requiring #{path} (if this hangs, make sure your code can be properly loaded as a library)")
+          require path
+        end
+      rescue Exception => e
+        log_info("Proceeding with postload -- could not load #{path}: #{e} (#{e.class})\n  #{e.backtrace.join("\n  ")}")
+      else
+        if defined?(einhorn_main)
+          log_info("Successfully loaded #{path}")
+          Einhorn::TransientState.preloaded = true
+        else
+          log_info("Proceeding with postload -- loaded #{path}, but no einhorn_main method was defined")
+        end
+      end
+    end
+  end
+
+  def self.set_argv(cmd, set_ps_name)
+    # TODO: clean up this hack
+    idx = 0
+    if cmd[0] =~ /(^|\/)ruby$/
+      idx = 1
+    elsif !is_script(cmd[0])
+      log_info("WARNING: Going to set $0 to #{cmd[idx]}, but it doesn't look like a script")
+    end
+
+    if set_ps_name
+      # Note this will mess up $0 if we try using it in our code, but
+      # we don't so that's basically ok. It's a bit annoying that this
+      # is how Ruby exposes changing the output of ps. Note that Ruby
+      # doesn't seem to shrink your cmdline buffer, so ps just ends up
+      # having lots of trailing spaces if we set $0 to something
+      # short. In the future, we could try to not pass einhorn's
+      # state in ARGV.
+      $0 = worker_ps_name
+    end
+
+    ARGV[0..-1] = cmd[idx+1..-1]
+    log_info("Set#{set_ps_name ? " $0 = #{$0.inspect}, " : nil} ARGV = #{ARGV.inspect}")
+  end
+
+  def self.set_master_ps_name
+    $0 = master_ps_name
+  end
+
+  def self.master_ps_name
+    "einhorn: #{worker_ps_name}"
+  end
+
+  def self.worker_ps_name
+    Einhorn::State.cmd_name ? "ruby #{Einhorn::State.cmd_name}" : Einhorn::State.orig_cmd.join(' ')
+  end
+
+  def self.socketify!(cmd)
+    cmd.map! do |arg|
+      if arg =~ /^(.*=|)srv:([^:]+):(\d+)((?:,\w+)*)$/
+        opt = $1
+        host = $2
+        port = $3
+        flags = $4.split(',').select {|flag| flag.length > 0}.map {|flag| flag.downcase}
+        fd = (Einhorn::State.sockets[[host, port]] ||= bind(host, port, flags))
+        "#{opt}#{fd}"
+      else
+        arg
+      end
+    end
+  end
+
+  def self.run
+    Einhorn::Command::Interface.init
+    Einhorn::Event.init
+
+    unless Einhorn::TransientState.stateful
+      if Einhorn::State.config[:number] < 1
+        log_error("You need to spin up at least at least 1 copy of the process")
+        return
+      end
+      Einhorn::Command::Interface.persistent_init
+
+      Einhorn::State.orig_cmd = ARGV.dup
+      Einhorn::State.cmd = ARGV.dup
+      # TODO: don't actually alter ARGV[0]?
+      Einhorn::State.cmd[0] = which(Einhorn::State.cmd[0])
+      socketify!(Einhorn::State.cmd)
+    end
+
+    set_master_ps_name
+    preload
+
+    # In the middle of upgrading
+    if Einhorn::State.reloading_for_preload_upgrade
+      Einhorn::Command.upgrade_workers
+      Einhorn::State.reloading_for_preload_upgrade = false
+    end
+
+    while Einhorn::State.respawn || Einhorn::State.children.size > 0
+      log_debug("Entering event loop")
+      # All of these are non-blocking
+      Einhorn::Command.reap
+      Einhorn::Command.replenish
+      Einhorn::Command.cull
+
+      # Make sure to do this last, as it's blocking.
+      Einhorn::Event.loop_once
+    end
+  end
+end
+
+require 'einhorn/command'
+require 'einhorn/client'
+require 'einhorn/event'
+require 'einhorn/worker'
+require 'einhorn/worker_pool'
+require 'einhorn/version'

data/test/test_helper.rb

@@ -0,0 +1,7 @@
+require 'rubygems'
+
+require 'test/unit'
+require 'mocha'
+require 'shoulda'
+
+$:.unshift(File.join(File.dirname(__FILE__), '../lib'))

data/test/unit/einhorn/command/interface.rb

@@ -0,0 +1,47 @@
+require File.expand_path(File.join(File.dirname(__FILE__), '../../../test_helper'))
+
+require 'einhorn'
+
+class InterfaceTest < Test::Unit::TestCase
+  include Einhorn::Command
+
+  context "when a command is received" do
+    should "call that command" do
+      conn = stub(:log_debug => nil)
+      conn.expects(:write).once.with do |message|
+        parsed = JSON.parse(message)
+        parsed['message'] =~ /Welcome gdb/
+      end
+      request = {
+        'command' => 'ehlo',
+        'user' => 'gdb'
+      }
+      Interface.process_command(conn, JSON.generate(request))
+    end
+  end
+
+  context "when an unrecognized command is received" do
+    should "call the unrecognized_command method" do
+      conn = stub(:log_debug => nil)
+      Interface.expects(:unrecognized_command).once
+      request = {
+        'command' => 'made-up',
+      }
+      Interface.process_command(conn, JSON.generate(request))
+    end
+  end
+
+  context "when a worker ack is received" do
+    should "register ack and close the connection" do
+      conn = stub(:log_debug => nil)
+      conn.expects(:close).once
+      conn.expects(:write).never
+      request = {
+        'command' => 'worker:ack',
+        'pid' => 1234
+      }
+      Einhorn::Command.expects(:register_manual_ack).once.with(1234)
+      Interface.process_command(conn, JSON.generate(request))
+    end
+  end
+end

data/test/unit/einhorn/command.rb

@@ -0,0 +1,21 @@
+require File.expand_path(File.join(File.dirname(__FILE__), '../../test_helper'))
+
+require 'einhorn'
+
+class CommandTest < Test::Unit::TestCase
+  include Einhorn
+
+  context "when running quieter" do
+    should "increase the verbosity threshold" do
+      Einhorn::State.stubs(:verbosity => 1)
+      Einhorn::State.expects(:verbosity=).once.with(2).returns(2)
+      Command.quieter
+    end
+
+    should "max out at 2" do
+      Einhorn::State.stubs(:verbosity => 2)
+      Einhorn::State.expects(:verbosity=).never
+      Command.quieter
+    end
+  end
+end