einhorn 0.3.0

data/bin/einhornsh

@@ -0,0 +1,120 @@
+#!/usr/bin/env ruby
+require 'logger'
+require 'optparse'
+
+require 'readline'
+
+require 'rubygems'
+require 'einhorn'
+
+module Einhorn
+  class EinhornSH
+    def initialize(path_to_socket)
+      @path_to_socket = path_to_socket
+      reconnect
+    end
+
+    def run
+      puts "Enter 'help' if you're not sure what to do."
+      puts
+      puts 'Type "quit" or "exit" to quit at any time'
+
+      while line = Readline.readline('> ', true)
+        if ['quit', 'exit'].include?(line)
+          puts "Goodbye!"
+          return
+        end
+
+        begin
+          response = @client.command({'command' => line})
+        rescue Errno::EPIPE => e
+          puts "einhornsh: Error communicating with Einhorn: #{e} (#{e.class})"
+          puts "einhornsh: Attempting to reconnect..."
+          reconnect
+
+          retry
+        end
+        puts response['message']
+      end
+    end
+
+    def reconnect
+      begin
+        @client = Einhorn::Client.for_path(@path_to_socket)
+      rescue Errno::ENOENT => e
+        # TODO: The exit here is a biit of a layering violation.
+        puts <<EOF
+Could not connect to Einhorn master process:
+
+  #{e}
+
+HINT: Are you sure you are running an Einhorn master?  If so, you
+should pass einhornsh the cmd_name (-c argument) provided to Einhorn.
+EOF
+        exit(1)
+      end
+      ehlo
+    end
+
+    def ehlo
+      response = @client.command('command' => 'ehlo', 'user' => ENV['USER'])
+      puts response['message']
+    end
+  end
+end
+
+def main
+  options = {}
+  optparse = OptionParser.new do |opts|
+    opts.banner = "Usage: #{$0} [options] [cmd_name]
+
+Welcome to Einhornsh: the Einhorn shell.
+
+Pass the cmd_name of the Einhorn master you are connecting to either
+as a positional argument or using `-c`. If you're running your Einhorn
+with a `-d`, provide the same argument here."
+
+    opts.on('-h', '--help', 'Display this message') do
+      puts opts
+      exit(1)
+    end
+
+    opts.on('-c CMD_NAME', '--command-name CMD_NAME', 'Connect to the Einhorn master with this cmd_name') do |cmd_name|
+      options[:cmd_name] = cmd_name
+    end
+
+    opts.on('-d PATH', '--socket-path PATH', 'Path to the Einhorn command socket') do |path|
+      options[:socket_path] = path
+    end
+  end
+  optparse.parse!
+
+  if ARGV.length > 1
+    puts optparse
+    return 1
+  end
+
+  Signal.trap("INT") {puts; exit(0)}
+
+  path_to_socket = options[:socket_path]
+
+  unless path_to_socket
+    cmd_name = options[:cmd_name] || ARGV[0]
+    path_to_socket = Einhorn::Command::Interface.default_socket_path(cmd_name)
+  end
+
+  sh = Einhorn::EinhornSH.new(path_to_socket)
+  sh.run
+  return 0
+end
+
+# Would be nice if this could be loadable rather than always
+# executing, but when run under gem it's a bit hard to do so.
+if true # $0 == __FILE__
+  ret = main
+  begin
+    exit(ret)
+  rescue TypeError
+    exit(0)
+  end
+end

data/einhorn.gemspec

@@ -0,0 +1,21 @@
+# -*- encoding: utf-8 -*-
+require File.expand_path('../lib/einhorn/version', __FILE__)
+
+Gem::Specification.new do |gem|
+  gem.authors       = ["Greg Brockman"]
+  gem.email         = ["gdb@stripe.com"]
+  gem.summary       = "Einhorn: the language-independent shared socket manager"
+  gem.description   = "Einhorn makes it easy to run multiple instances of an application server, all listening on the same port. You can also seamlessly restart your workers without dropping any requests. Einhorn requires minimal application-level support, making it easy to use with an existing project."
+  gem.homepage      = "https://github.com/stripe/einhorn"
+
+  gem.files         = `git ls-files`.split($\)
+  gem.executables   = gem.files.grep(%r{^bin/}).map{ |f| File.basename(f) }
+  gem.test_files    = gem.files.grep(%r{^(test|spec|features)/})
+  gem.name          = "einhorn"
+  gem.require_paths = ["lib"]
+  # maybe swap out for YAML? Then don't need any gems.
+  gem.add_dependency('json')
+  gem.add_development_dependency('shoulda')
+  gem.add_development_dependency('mocha')
+  gem.version       = Einhorn::VERSION
+end

data/example/pool_worker.rb

@@ -0,0 +1,19 @@
+#!/usr/bin/env ruby
+#
+# Used as an example of preloading in Einhorn blog post
+# (https://stripe.com/blog/meet-einhorn). Program name ends in .rb in
+# order to make explicit that it's written in Ruby, though this isn't
+# actually necessary for preloading to work.
+#
+# Run as
+#
+# einhorn -p ./pool_worker.rb ./pool_worker.rb
+
+puts "From PID #{$$}: loading #{__FILE__}"
+
+def einhorn_main
+  while true
+    puts "From PID #{$$}: Doing some work"
+    sleep 1
+  end
+end

data/example/thin_example

@@ -0,0 +1,52 @@
+#!/usr/bin/env ruby
+#
+# An example application using our patched Thin and EventMachine. You
+# can obtain these from:
+#
+#   https://github.com/stripe/thin.git, and
+#   https://github.com/stripe/eventmachine.git
+
+require 'rubygems'
+require 'einhorn'
+
+# Make sure we're using the patched versions.
+gem 'thin', '1.3.2.stripe.0'
+gem 'eventmachine', '1.0.0.beta.4.stripe.0'
+
+require 'thin'
+
+class App
+  def initialize(id)
+    @id = id
+  end
+
+  def call(env)
+    return [200, {}, "[#{$$}] From server instance #{@id}: Got your request!\n"]
+  end
+end
+
+def einhorn_main
+  puts "Called with #{ARGV.inspect}"
+
+  if ARGV.length == 0
+    raise "Need to call with at least one argument. Try running 'einhorn #{$0} srv:127.0.0.1:5000,r,n srv:127.0.0.1:5001,r,n' and then running 'curl 127.0.0.1:5000' or 'curl 127.0.0.1:5001'"
+  end
+
+  Einhorn::Worker.graceful_shutdown do
+    puts "#{$$} is now exiting..."
+    exit(0)
+  end
+  Einhorn::Worker.ack!
+
+  EventMachine.run do
+    ARGV.each_with_index do |arg, i|
+      sock = Integer(arg)
+      srv = Thin::Server.new(sock, App.new(i), :reuse => true)
+      srv.start
+    end
+  end
+end
+
+if $0 == __FILE__
+  einhorn_main
+end

data/example/time_server

@@ -0,0 +1,48 @@
+#!/usr/bin/env ruby
+#
+# A simple example showing how to use Einhorn's shared-socket
+# features. Einhorn translates the srv:(addr:port[,flags...]) spec in
+# the arg string into a file descriptor number.
+#
+# Invoke through Einhorn as
+#
+#     einhorn ./time_server srv:127.0.0.1:2345,r
+#
+# or, if you want to try out preloading:
+#
+#     einhorn -p ./time_server ./time_server srv:127.0.0.1:2345,r
+
+require 'rubygems'
+require 'einhorn/worker'
+
+def einhorn_main
+  puts "Called with #{ARGV.inspect}"
+
+  if ARGV.length != 1
+    raise "Need to call with a port spec as the first argument. Try running 'einhorn #{$0} srv:127.0.0.1:2345,r' and then running 'nc 127.0.0.1 2345'"
+  end
+
+  socket = Socket.for_fd(Integer(ARGV[0]))
+
+  # Came up successfully, so let's set up graceful handler and ACK the
+  # master.
+  Einhorn::Worker.graceful_shutdown do
+    puts "Goodbye from #{$$}"
+    exit(0)
+  end
+  Einhorn::Worker.ack!
+
+  # Real work happens here.
+  begin
+    while true
+      accepted, _ = socket.accept
+      accepted.write("[#{$$}] The current time is: #{Time.now}!\n")
+      accepted.close
+    end
+  rescue Exception
+  end
+end
+
+if $0 == __FILE__
+  einhorn_main
+end

data/lib/einhorn/client.rb

@@ -0,0 +1,48 @@
+require 'json'
+require 'set'
+
+module Einhorn
+  class Client
+    @@responseless_commands = Set.new(['worker:ack'])
+
+    def self.for_path(path_to_socket)
+      socket = UNIXSocket.open(path_to_socket)
+      self.new(socket)
+    end
+
+    def self.for_fd(fileno)
+      socket = UNIXSocket.for_fd(fileno)
+      self.new(socket)
+    end
+
+    def initialize(socket)
+      @socket = socket
+    end
+
+    def command(command_hash)
+      command = JSON.generate(command_hash) + "\n"
+      write(command)
+      recvmessage if expect_response?(command_hash)
+    end
+
+    def expect_response?(command_hash)
+      !@@responseless_commands.include?(command_hash['command'])
+    end
+
+    def close
+      @socket.close
+    end
+
+    private
+
+    def write(bytes)
+      @socket.write(bytes)
+    end
+
+    # TODO: use a streaming JSON parser instead?
+    def recvmessage
+      line = @socket.readline
+      JSON.parse(line)
+    end
+  end
+end

data/lib/einhorn/command/interface.rb

@@ -0,0 +1,336 @@
+require 'tmpdir'
+require 'socket'
+
+module Einhorn::Command
+  module Interface
+    @@commands = {}
+    @@command_server = nil
+
+    def self.command_server=(server)
+      raise "Command server already set" if @@command_server && server
+      @@command_server = server
+    end
+
+    def self.command_server
+      @@command_server
+    end
+
+    def self.init
+      install_handlers
+      at_exit do
+        if Einhorn::TransientState.whatami == :master
+          to_remove = [pidfile]
+          # Don't nuke socket_path if we never successfully acquired it
+          to_remove << socket_path if @@command_server
+          to_remove.each do |file|
+            begin
+              File.unlink(file)
+            rescue Errno::ENOENT
+            end
+          end
+        end
+      end
+    end
+
+    def self.persistent_init
+      socket = open_command_socket
+      Einhorn::Event::CommandServer.open(socket)
+
+      # Could also rewrite this on reload. Might be useful in case
+      # someone goes and accidentally clobbers/deletes. Should make
+      # sure that the clobber is atomic if we we were do do that.
+      write_pidfile
+    end
+
+    def self.open_command_socket
+      path = socket_path
+
+      with_file_lock do
+        # Need to avoid time-of-check to time-of-use bugs in blowing
+        # away and recreating the old socketfile.
+        destroy_old_command_socket(path)
+        UNIXServer.new(path)
+      end
+    end
+
+    # Lock against other Einhorn workers. Unfortunately, have to leave
+    # this lockfile lying around forever.
+    def self.with_file_lock(&blk)
+      path = lockfile
+      File.open(path, 'w', 0600) do |f|
+        unless f.flock(File::LOCK_EX|File::LOCK_NB)
+          raise "File lock already acquired by another Einhorn process. This likely indicates you tried to run Einhorn masters with the same cmd_name at the same time. This is a pretty rare race condition."
+        end
+
+        blk.call
+      end
+    end
+
+    def self.destroy_old_command_socket(path)
+      # Socket isn't actually owned by anyone
+      begin
+        sock = UNIXSocket.new(path)
+      rescue Errno::ECONNREFUSED
+        # This happens with non-socket files and when the listening
+        # end of a socket has exited.
+      rescue Errno::ENOENT
+        # Socket doesn't exist
+        return
+      else
+        # Rats, it's still active
+        sock.close
+        raise Errno::EADDRINUSE.new("Another process (probably another Einhorn) is listening on the Einhorn command socket at #{path}. If you'd like to run this Einhorn as well, pass a `-d PATH_TO_SOCKET` to change the command socket location.")
+      end
+
+      # Socket should still exist, so don't need to handle error.
+      stat = File.stat(path)
+      unless stat.socket?
+        raise Errno::EADDRINUSE.new("Non-socket file present at Einhorn command socket path #{path}. Either remove that file and restart Einhorn, or pass a `-d PATH_TO_SOCKET` to change the command socket location.")
+      end
+
+      Einhorn.log_info("Blowing away old Einhorn command socket at #{path}. This likely indicates a previous Einhorn worker which exited uncleanly.")
+      # Whee, blow it away.
+      File.unlink(path)
+    end
+
+    def self.write_pidfile
+      file = pidfile
+      Einhorn.log_info("Writing PID to #{file}")
+      File.open(file, 'w') {|f| f.write($$)}
+    end
+
+    def self.uninit
+      remove_handlers
+    end
+
+    def self.socket_path
+      Einhorn::State.socket_path || default_socket_path
+    end
+
+    def self.default_socket_path(cmd_name=nil)
+      cmd_name ||= Einhorn::State.cmd_name
+      if cmd_name
+        filename = "einhorn-#{cmd_name}.sock"
+      else
+        filename = "einhorn.sock"
+      end
+      File.join(Dir.tmpdir, filename)
+    end
+
+    def self.lockfile
+      Einhorn::State.lockfile || default_lockfile_path
+    end
+
+    def self.default_lockfile_path(cmd_name=nil)
+      cmd_name ||= Einhorn::State.cmd_name
+      if cmd_name
+        filename = "einhorn-#{cmd_name}.lock"
+      else
+        filename = "einhorn.lock"
+      end
+      File.join(Dir.tmpdir, filename)
+    end
+
+    def self.pidfile
+      Einhorn::State.pidfile || default_pidfile
+    end
+
+    def self.default_pidfile(cmd_name=nil)
+      cmd_name ||= Einhorn::State.cmd_name
+      if cmd_name
+        filename = "einhorn-#{cmd_name}.pid"
+      else
+        filename = "einhorn.pid"
+      end
+      File.join(Dir.tmpdir, filename)
+    end
+
+    ## Signals
+    def self.install_handlers
+      Signal.trap("INT") do
+        Einhorn::Command.signal_all("USR2", Einhorn::State.children.keys)
+        Einhorn::State.respawn = false
+      end
+      Signal.trap("TERM") do
+        Einhorn::Command.signal_all("TERM", Einhorn::State.children.keys)
+        Einhorn::State.respawn = false
+      end
+      # Note that quit is a bit different, in that it will actually
+      # make Einhorn quit without waiting for children to exit.
+      Signal.trap("QUIT") do
+        Einhorn::Command.signal_all("QUIT", Einhorn::State.children.keys)
+        Einhorn::State.respawn = false
+        exit(1)
+      end
+      Signal.trap("HUP") {Einhorn::Command.reload}
+      Signal.trap("ALRM") {Einhorn::Command.full_upgrade}
+      Signal.trap("CHLD") {Einhorn::Event.break_loop}
+      Signal.trap("USR2") do
+        Einhorn::Command.signal_all("USR2", Einhorn::State.children.keys)
+        Einhorn::State.respawn = false
+      end
+      at_exit do
+        if Einhorn::State.kill_children_on_exit && Einhorn::TransientState.whatami == :master
+          Einhorn::Command.signal_all("USR2", Einhorn::State.children.keys)
+          Einhorn::State.respawn = false
+        end
+      end
+    end
+
+    def self.remove_handlers
+      %w{INT TERM QUIT HUP ALRM CHLD USR2}.each do |signal|
+        Signal.trap(signal, "DEFAULT")
+      end
+    end
+
+    ## Commands
+    def self.command(name, description=nil, &code)
+      @@commands[name] = {:description => description, :code => code}
+    end
+
+    def self.process_command(conn, command)
+      response = generate_response(conn, command)
+      if !response.nil?
+        send_message(conn, response)
+      else
+        conn.log_debug("Got back nil response, so not responding to command.")
+      end
+    end
+
+    def self.send_message(conn, response)
+      if response.kind_of?(String)
+        response = {'message' => response}
+      end
+      message = pack_message(response)
+      conn.write(message)
+    end
+
+    def self.generate_response(conn, command)
+      begin
+        request = JSON.parse(command)
+      rescue JSON::ParserError => e
+        return {
+          'message' => "Could not parse command: #{e}"
+        }
+      end
+
+      unless command_name = request['command']
+        return {
+          'message' => 'No "command" parameter provided; not sure what you want me to do.'
+        }
+      end
+
+      if command_spec = @@commands[command_name]
+        conn.log_debug("Received command: #{command.inspect}")
+        begin
+          return command_spec[:code].call(conn, request)
+        rescue StandardError => e
+          msg = "Error while processing command #{command_name.inspect}: #{e} (#{e.class})\n  #{e.backtrace.join("\n  ")}"
+          conn.log_error(msg)
+          return msg
+        end
+      else
+        conn.log_debug("Received unrecognized command: #{command.inspect}")
+        return unrecognized_command(conn, request)
+      end
+    end
+
+    def self.pack_message(message_struct)
+      begin
+        JSON.generate(message_struct) + "\n"
+      rescue JSON::GeneratorError => e
+        response = {
+          'message' => "Error generating JSON message for #{message_struct.inspect} (this indicates a bug): #{e}"
+        }
+        JSON.generate(response) + "\n"
+      end
+    end
+
+    def self.command_descriptions
+      command_specs = @@commands.select do |_, spec|
+        spec[:description]
+      end.sort_by {|name, _| name}
+
+      command_specs.map do |name, spec|
+        "#{name}: #{spec[:description]}"
+      end.join("\n")
+    end
+
+    def self.unrecognized_command(conn, request)
+      <<EOF
+Unrecognized command: #{request['command'].inspect}
+
+#{command_descriptions}
+EOF
+    end
+
+    # Used by workers
+    command 'worker:ack' do |conn, request|
+      if pid = request['pid']
+        Einhorn::Command.register_manual_ack(pid)
+      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
+Welcome #{request['user']}! You are speaking to Einhorn Master Process #{$$}#{Einhorn::State.cmd_name ? " (#{Einhorn::State.cmd_name})" : ''}
+EOF
+    end
+
+    command 'help', 'Print out available commands' do
+"You are speaking to the Einhorn command socket. You can run the following commands:
+
+#{command_descriptions}
+"
+    end
+
+    command 'state', "Get a dump of Einhorn's current state" do
+      Einhorn::Command.dumpable_state.pretty_inspect
+    end
+
+    command 'reload', 'Reload Einhorn' do |conn, _|
+      # TODO: make reload actually work (command socket reopening is
+      # an issue). Would also be nice if user got a confirmation that
+      # the reload completed, though that's not strictly necessary.
+
+      # In the normal case, this will do a write
+      # synchronously. Otherwise, the bytes will be stuck into the
+      # buffer and lost upon reload.
+      send_message(conn, 'Reloading, as commanded')
+      Einhorn::Command.reload
+
+      # Reload should not return
+      raise "Not reachable"
+    end
+
+    command 'inc', 'Increment the number of Einhorn child processes' do
+      Einhorn::Command.increment
+    end
+
+    command 'dec', 'Decrement the number of Einhorn child processes' do
+      Einhorn::Command.decrement
+    end
+
+    command 'quieter', 'Decrease verbosity' do
+      Einhorn::Command.quieter
+    end
+
+    command 'louder', 'Increase verbosity' do
+      Einhorn::Command.louder
+    end
+
+    command 'upgrade', 'Upgrade all Einhorn workers. This may result in Einhorn reloading its own code as well.' do |conn, _|
+      # TODO: send confirmation when this is done
+      send_message(conn, 'Upgrading, as commanded')
+      # This or may not return
+      Einhorn::Command.full_upgrade
+      nil
+    end
+  end
+end