royw-drbman 0.0.1

data/lib/drbman/drb_pool.rb

@@ -0,0 +1,111 @@
+# == Synopsis
+# A pool of drb objects
+class DrbPool
+  # Create the pool of drb objects.
+  # @param [Array<HostMachine>] the host_machine instances to use to populate the pool of drb objects
+  # @param [Logger] the logger
+  # @yield [self]
+  # @example Without using a block
+  #   pool = DrbPool.new(hosts, logger)
+  #   pool.getObject {|obj| obj.do_something}
+  #   pool.shutdown
+  # @example Using a block
+  #  DrbPool.new(hosts, logger) do |pool|
+  #    pool.getObject {|obj| obj.do_something}
+  #  end
+  def initialize(hosts, logger, &block)
+    @logger = logger
+    @objects = []
+
+    threads = []
+    mutex = Mutex.new
+    hosts.each do |host_name, host_machine|
+      threads << Thread.new(host_machine) do |host|
+        obj = get_drb_object(host.machine, host.port)
+        mutex.synchronize do
+          @objects << obj
+        end
+      end
+    end
+    threads.each {|thrd| thrd.join}
+
+    unless block.nil?
+      block.call(self)
+      shutdown
+    end
+  end
+  
+  # Use an object from the pool
+  # @yield [DRbObject]
+  # @example Usage
+  #   pool.get_object {|obj| obj.do_something}
+  def get_object(&block)
+    mutex = Mutex.new
+    while((object = next_object(mutex)).nil?)
+      sleep 0.1 
+    end
+    raise ArgumentError.new('a block is required') if block.nil?
+    block.call(object)
+    object.in_use = false
+  end
+  
+  # Shut the pool down
+  # Only necessary if not using a block with DrbPool.new
+  # @example
+  #   pool = DrbPool.new(hosts, logger)
+  #   pool.getObject {|obj| obj.do_something}
+  #   pool.shutdown
+  def shutdown
+    @objects.each do |obj|
+      obj.stop_service
+    end
+  end
+  
+  private
+
+  # find the next available object
+  # @return [DRbObject, nil] returns nil if no objects are available
+  def next_object(mutex)
+    object = nil
+    mutex.synchronize do
+      object = @objects.select {|obj| !obj.in_use?}.first
+      object.in_use = true unless object.nil?
+    end
+    object
+  end
+  
+  # Get a DRbObject to place into the pool
+  # The object is extended with an in_use attribute
+  # @param machine [String] the host machine name, for example: 'foo.example.com'
+  # @param port [Integer] the port the drb server on the host machine is using
+  # @return [DRbObject,nil] returns nil if unable to get the DRbObject
+  def get_drb_object(machine, port)
+    obj = nil
+    retry_cnt = 0
+    DRb.start_service
+    begin
+      obj = DRbObject.new(nil, "druby://#{machine}:#{port}")
+      obj.extend(InUse)
+      name = obj.name
+      obj.in_use = false
+    rescue Exception => e
+      retry_cnt += 1
+      raise e if retry_cnt > 10
+      sleep 0.2
+      @logger.debug {"retrying (#{retry_cnt})"}
+      retry
+    end
+    obj
+  end
+  
+  # Adds an in_use attribute
+  module InUse
+    def in_use=(flag)
+      @in_use = (flag ? true : false)
+    end
+    def in_use?
+      @in_use
+    end
+  end
+  
+end

data/lib/drbman/drbman.rb

@@ -0,0 +1,229 @@
+# Drbman is the drb manager
+#
+# == Synopsis
+# Drbman will create a project directory on a host machine,
+# then copy a set of files to the host machine, make sure
+# a given set of gems is installed on the host machine, then
+# run the drb server on the host machine.
+# Drbman also supports issuing a termination command to
+# the drb server on the remote machine and cleaning up
+# the project by removing the files installed onto the
+# host machine.
+#
+# == Notes
+# Uses the Command design pattern
+class Drbman
+  # @param [Logger] logger the logger
+  # @param [UserChoices,Hash] choices
+  # @option choices [Array<String>] :dirs array of local directories to copy to the host machines.  REQUIRED
+  # @option choices [String] :run the name of the file to run on the host machine.  REQUIRED
+  #  This file should start the drb server.  Note, this file will be daemonized before running.
+  # @option choices [Array<String>] :hosts array of host machine descriptions "{user{:password}@}machine{:port}"
+  #  This defaults to ['localhost']
+  # @option choices [Integer] :port default port number used to assign to hosts without a port number.
+  #  The port number is incremented for each host.  This defaults to 9000
+  # @option choices [Array<String>] :gems array of gem names to verify are installed on the host machine.
+  #  Note, 'daemons' is always added to this array.
+  # @yield [Drbman]
+  # @example Usage
+  # Drbman.new(logger, choices) do |drbman|
+  #   drbman.get_object do |obj|
+  #     obj.do_something
+  #   end
+  # end
+  def initialize(logger, choices, &block)
+    @logger = logger
+    @user_choices = choices
+    
+    # @hosts[machine_description] = HostMachine instance
+    @hosts = {}
+
+    choices[:port] ||= 9000
+    choices[:hosts] ||= ['localhost']
+    choices[:gems] ||= []
+    choices[:gems] = (choices[:gems] + ['daemons']).uniq.compact
+    
+    raise ArgumentError.new('Missing choices[:run]') if choices[:run].blank?
+    raise ArgumentError.new('Missing choices[:hosts]') if choices[:hosts].blank?
+    raise ArgumentError.new('Missing choices[:dirs]') if choices[:dirs].blank?
+    
+    # populate the @hosts hash.  key => host machine description, value => HostMachine instance
+    port = choices[:port]
+    @user_choices[:hosts].each do |host|
+      host = "#{host}:#{port}" unless host =~ /\:\d+\s*$/
+      @hosts[host] = HostMachine.new(host, @logger)
+      port += 1
+    end
+    
+    unless block.nil?
+      begin
+        setup
+        @pool = DrbPool.new(@hosts, @logger)
+        block.call(self)
+      rescue Exception => e
+        @logger.error { e }
+        @logger.debug { e.backtrace.join("\n") }
+      ensure
+        @pool.shutdown unless @pool.nil?
+        shutdown
+      end
+    end
+  end
+  
+  # Use an object from the pool
+  # @yield [DRbObject]
+  # @example Usage
+  #   drbman.get_object {|obj| obj.do_something}
+  def get_object(&block)
+    @pool.get_object(&block)
+  end
+  
+  private
+
+  # setup the host machine drb servers
+  # @raise [Exception] when a component is not installed on the host
+  def setup
+    threads = []
+    @hosts.each do |name, machine|
+      threads << Thread.new(machine) do |host_machine|
+        host_machine.session do |host|
+          # begin
+            startup(host)
+          # rescue Exception => e
+          #   @logger.error { e }
+          #   @logger.debug { e.backtrace.join("\n") }
+          # end
+        end
+      end
+    end
+    threads.each {|thrd| thrd.join}
+    sleep 1  # give the drb servers a little time to get running
+  end
+  
+  # stop and remove the host machine drb servers
+  def shutdown
+    threads = []
+    @hosts.each do |name, machine|
+      threads << Thread.new(machine) do |host_machine|
+        host_machine.session do |host|
+          begin
+            cleanup(host)
+          rescue Exception => e
+            @logger.error { e }
+            @logger.debug { e.backtrace.join("\n") }
+          end
+        end
+      end
+    end
+    threads.each {|thrd| thrd.join}
+  end
+  
+  # Setup the drb server on the given host then start it
+  # @param [HostMachine] host the host machine
+  # @raise [Exception] when a component is not installed on the host
+  def startup(host)
+    @logger.info { "Setting up: #{host.name}" }
+    check_gems(host)
+    create_directory(host)
+    upload_dirs(host)
+    create_controller(host)
+    run_drb_server(host)
+  end
+  
+  # Stop the drb server on the given host then remove the drb files from the host
+  # @param [HostMachine] host the host machine
+  def cleanup(host)
+    @logger.info { "Cleaning up: #{host.name}" }
+    stop_drb_server(host)
+    cleanup_files(host)
+  end
+  
+  # remove the host directory and any files in it from
+  # the host machine
+  # @param [HostMachine] host the host machine
+  def cleanup_files(host)
+    unless host.dir.blank? || (host.dir =~ /[\*\?]/)
+      @logger.debug { "#{host.name}: rm -rf #{host.dir}"}
+      host.sh("rm -rf #{host.dir}")
+    end
+  end
+  
+  # run the drb server stop command on the host machine
+  # @param [HostMachine] host the host machine
+  # @raise [Exception] when ruby is not installed on the host
+  def stop_drb_server(host)
+    case host.sh("cd #{host.dir};ruby #{host.controller} stop")
+    when /command not found/
+      raise Exception.new "Ruby is not installed on #{host.name}"
+    end
+  end
+  
+  # run the drb server start command on the host machine
+  # @param [HostMachine] host the host machine
+  # @raise [Exception] when ruby is not installed on the host
+  def run_drb_server(host)
+    unless host.controller.blank?
+      case host.sh("cd #{host.dir};ruby #{host.controller} start -- #{host.machine} #{host.port}")
+      when /command not found/
+        raise Exception.new "Ruby is not installed on #{host.name}"
+      end
+    end
+  end
+
+  # Create the daemon controller on the host machine
+  # @param [HostMachine] host the host machine
+  def create_controller(host)
+    unless @user_choices[:run].blank?
+      host.controller = File.basename(@user_choices[:run], '.*') + '_controller.rb'
+      controller = "require 'rubygems' ; require 'daemons' ; Daemons.run('#{@user_choices[:run]}')"
+      host.sh("cd #{host.dir};echo \"#{controller}\" > #{host.controller}")
+    end
+  end
+  
+  # copy files from local directories to the created 
+  # directory on the host machine
+  # @param [HostMachine] host the host machine
+  def upload_dirs(host)
+    unless @user_choices[:dirs].blank?
+      @user_choices[:dirs].each do |name|
+        if File.directory?(name)
+          host.upload(name, "#{host.dir}/#{File.basename(name)}")
+        else
+          @logger.error { "\"#{name}\" is not a directory" }
+        end
+      end
+    end
+  end
+  
+  # check if the required gems are installed on the host
+  # @param [HostMachine] host the host machine
+  # @raise [Exception] when rubygems is not installed on the host
+  def check_gems(host)
+    missing_gems = []
+    @user_choices[:gems].each do |gem_name|
+      case host.sh("gem list -l -i #{gem_name}")
+      when /false/i
+        missing_gems << gem_name
+      when /command not found/
+        raise Exception.new "Rubygems is not installed on #{host.name}"
+      end
+    end
+    unless missing_gems.empty?
+      raise Exception.new "The following gems are not installed on #{host.name}: #{missing_gems.join(', ')}"
+    end
+  end
+
+  # Create the directory that will hold the drb files.
+  # The directory created is ~/.drbman/{uuid}
+  # Note, requires 'uuidgen' in the path on the host machine
+  # @todo maybe generate a random uuid locally instead
+  # @param [HostMachine] host the host machine
+  def create_directory(host)
+    host.uuid = UUIDTools::UUID.random_create
+    host.dir = "~/.drbman/#{host.uuid}".strip
+    host.sh("mkdir -p #{host.dir}")
+    @logger.debug { "host directory: #{host.dir}" }
+  end
+  
+end
+

data/lib/drbman/host_machine.rb

@@ -0,0 +1,216 @@
+# HostMachine is used to interface with a host machine
+# 
+# == Notes
+# A host machine may be either another machine or the localhost.
+# By supporting localhost, it is likely that the process will
+# be on a different core than the current processes.
+#
+# Once HostMachine opens an ssh connection, it does not close
+# the connection until a disconnect() is invoked.
+class HostMachine
+  attr_accessor :uuid, :dir, :controller
+  attr_reader :name, :machine, :user, :port
+  
+  # @param [String] host_string describes the host to connect to.
+  #  The format is "{user{:password}@}machine{:port}"
+  # @param [Logger] logger the logger to use
+  def initialize(host_string, logger)
+    @logger = logger
+    @machine = 'localhost'
+    @user = ENV['USER']
+    @port = 9000
+    @password = {:keys => ['~/.ssh/id_dsa']}
+    case host_string
+    when /^(\S+)\:(\S+)\@(\S+)\:(\d+)$/
+      @user = $1
+      @password = {:password => $2}
+      @machine = $3
+      @port = $4.to_i
+    when /^(\S+)\:(\S+)\@(\S+)$/
+      @user = $1
+      @password = {:password => $2}
+      @machine = $3
+    when /^(\S+)\@(\S+)\:(\d+)$/
+      @user = $1
+      @machine = $2
+      @port = $3.to_i
+    when /^(\S+)\@(\S+)$/
+      @user = $1
+      @machine = $2
+    when /^(\S+)\:(\d+)$/
+      @machine = $1
+      @port = $2.to_i
+    when /^(\S+)$/
+      @machine = $1
+    end
+    @name = "#{user}@#{machine}:#{port}"
+    @ssh = nil
+    @logger.debug { self.pretty_inspect }
+  end
+  
+  # Connect to the host, execute the given block, then disconnect from the host
+  # @yield [HostMachine]
+  # @example
+  # host_machine = HostMachine.new('localhost', @logger)
+  # host_machine.session do |host|
+  #   host.upload(local_dir, "#{host.dir}/#{File.basename(local_dir)}")
+  #   @logger.debug { host.sh("ls -lR #{host.dir}") }
+  # end
+  def session(&block)
+    connect
+    yield self
+    disconnect
+  end
+  
+  # upload a directory structure to the host machine.
+  # @param [String] local_src the source directory on the local machine
+  # @param [String] remote_dest the destination directory on the host machine
+  # @raise [Exception] if the files are not copied
+  def upload(local_src, remote_dest)
+    @logger.debug { "upload(\"#{local_src}\", \"#{remote_dest}\")" }
+    connect
+    result = nil
+    unless @ssh.nil?
+      begin
+        @ssh.scp.upload!(local_src, remote_dest, :recursive => true) do |ch, name, sent, total|
+          @logger.debug { "#{name}: #{sent}/#{total}" }
+        end
+        @ssh.loop
+      rescue Exception => e
+        # only raise the exception if the files differ
+        raise e unless same_files?(local_src, remote_dest)
+      end
+    end
+  end
+  
+  # download a directory structure from the host machine.
+  # @param [String] remote_src the source directory on the host machine
+  # @param [String] local_dest the destination directory on the local machine
+  # @raise [Exception] if the files are not copied
+  def download(remote_src, local_dest)
+    connect
+    result = nil
+    unless @ssh.nil?
+      begin
+        @ssh.scp.download!(local_src, remote_dest, :recursive => true) do |ch, name, sent, total|
+          @logger.debug { "#{name}: #{sent}/#{total}" }
+        end
+        @ssh.loop
+      rescue Exception => e
+        # only raise the exception if the files differ
+        raise e unless same_files?(local_dest, remote_src)
+      end
+    end
+  end
+
+  # run a command on the host machine
+  # Note that the environment on the host machine is the default environment instead
+  # of the user's environment.  So by default we try to source ~/.profile and ~/.bashrc
+  #
+  # @param [String] command the command to run
+  # @param [Hash] opts
+  # @options opts [Array<String>] :source array of files to source. defaults to ['~/.profile', '~/.bashrc']
+  def sh(command, opts={})
+    @logger.debug { "sh \"#{command}\""}
+    # if opts[:source].blank?
+    #   opts[:source] = ['~/.profile', '~/.bashrc']
+    # end
+    connect
+    result = nil
+    unless @ssh.nil?
+      if @pre_commands.nil?
+        @pre_commands = []
+        opts[:source] ||= []
+        opts[:source].each do |name|
+          ls_out = @ssh.exec!("ls #{name}")
+          @pre_commands << "source #{name}" if ls_out =~ /^\s*\S+\/#{File.basename(name)}\s*$/
+        end
+      end
+      commands = @pre_commands.clone
+      commands << command
+      command_line = commands.join(' && ')
+      result = @ssh.exec!(command_line)
+    end
+    result
+  end
+  
+  # run a command as the superuser on the host machine
+  # Note that the environment on the host machine is the default environment instead
+  # of the user's environment.  So by default we try to source ~/.profile and ~/.bashrc
+  #
+  # @param [String] command the command to run
+  # @param [Hash] opts
+  # @options opts [Array<String>] :source array of files to source. defaults to ['~/.profile', '~/.bashrc']
+  def sudo(command, opts={})
+    @logger.debug { "sudo \"#{command}\""}
+    # if opts[:source].blank?
+    #   opts[:source] = ['~/.profile', '~/.bashrc']
+    # end
+    connect
+    result = nil
+    unless @ssh.nil?
+      buf = []
+      @ssh.open_channel do |channel|
+        if @pre_commands.nil?
+          @pre_commands = []
+          opts[:source] ||= []
+          opts[:source].each do |name|
+            ls_out = @ssh.exec!("ls #{name}")
+            @pre_commands << "source #{name}" if ls_out =~ /^\s*\S+\/#{File.basename(name)}\s*$/
+          end
+        end
+        commands = @pre_commands.clone
+        commands << "sudo -p 'sudo password: ' #{command}"
+        command_line = commands.join(' && ')
+        channel.exec(command_line) do |ch, success|
+          ch.on_data do |ch, data|
+            if data =~ /sudo password: /
+              ch.send_data("#{@password[:password]}\n")
+            else
+              buf << data
+            end
+          end
+        end
+      end
+      @ssh.loop
+      result = buf.compact.join('')
+    end
+    result
+  end
+
+  # connect to the host machine
+  def connect
+    if @ssh.nil?
+      @ssh = Net::SSH.start(@machine, @user, @password)
+      # @ssh.forward.local(@port, @machine, @port)
+    end
+  end
+  
+  # disconnect from the host machine
+  def disconnect
+    if @ssh
+      @ssh.close
+      @ssh = nil
+    end
+  end
+
+  private
+  
+  # Does the local directory tree and the remote directory tree contain the same files?
+  # Calculates a MD5 hash for each file then compares the hashes
+  # @param [String] local_path local directory
+  # @param [String] remote_path remote directory
+  def same_files?(local_path, remote_path)
+    remote_md5 = @ssh.exec!(md5_command_line(remote_path))
+    local_md5 = `#{md5_command_line(local_path)}`
+    @logger.debug { "same_files? local_md5 => #{local_md5}, remote_md5 => #{remote_md5}"}
+    remote_md5 == local_md5
+  end
+  
+  def md5_command_line(dirname)
+    line = "cat \`find #{dirname} -type f | sort\` | ruby -e \"require 'digest/md5';puts Digest::MD5.hexdigest(STDIN.read)\""
+    @logger.debug { line }
+    line
+  end
+  
+end

data/lib/drbman.rb

@@ -0,0 +1,28 @@
+$:.unshift(File.dirname(__FILE__)) unless
+  $:.include?(File.dirname(__FILE__)) || $:.include?(File.expand_path(File.dirname(__FILE__)))
+
+require 'rubygems'
+
+# gem 'ruby-debug',       '>=0.10.3'
+gem 'log4r',            '>=1.0.5'
+gem 'user-choices',     '>=1.1.6'
+gem 'extlib',           '>=0.9.12'
+gem 'net-ssh',          '>=2.0.11'
+gem 'net-scp',          '>=1.0.2'
+gem 'daemons',          '>=1.0.10'
+gem 'uuidtools',        '>=2.0.0'
+
+require 'extlib'
+# require 'ruby-debug'
+require 'log4r'
+require 'user-choices'
+require 'net/ssh'
+require 'net/scp'
+require 'daemons'
+require 'uuidtools'
+
+# require all of the .rb files in the drbman/ subdirectory
+Dir.glob(File.join(File.dirname(__FILE__), 'drbman/**/*.rb')).each do |name| 
+  require name
+end
+

data/spec/drbman_spec.rb

@@ -0,0 +1,7 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe "Drbman" do
+  it "fails" do
+    fail "hey buddy, you should probably rename this file and start specing for real"
+  end
+end

data/spec/host_machine_spec.rb

@@ -0,0 +1,45 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe('HostMachine') do
+  before(:each) do
+    @logger = Log4r::Logger.new('host_machine_spec')
+  end
+  
+  it "should parse: localhost" do
+    host_machine = HostMachine.new('localhost', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'localhost'
+    host_machine.port.should == 9000
+  end
+
+  it "should parse: localhost:1234" do
+    host_machine = HostMachine.new('localhost:1234', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'localhost:1234'
+    host_machine.port.should == 1234
+  end
+  it "should parse: me@localhost" do
+    host_machine = HostMachine.new('me@localhost', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'me@localhost'
+    host_machine.port.should == 9000
+  end
+  it "should parse: me@localhost:1234" do
+    host_machine = HostMachine.new('me@localhost:1234', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'me@localhost:1234'
+    host_machine.port.should == 1234
+  end
+  it "should parse: me:sekret@localhost" do
+    host_machine = HostMachine.new('me:sekret@localhost', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'me:sekret@localhost'
+    host_machine.port.should == 9000
+  end
+  it "should parse: me:sekret@localhost:1234" do
+    host_machine = HostMachine.new('me:sekret@localhost:1234', @logger)
+    host_machine.machine.should == 'localhost'
+    host_machine.name.should == 'me:sekret@localhost:1234'
+    host_machine.port.should == 1234
+  end
+end

data/spec/spec_helper.rb

@@ -0,0 +1,9 @@
+require 'spec'
+
+$LOAD_PATH.unshift(File.dirname(__FILE__))
+$LOAD_PATH.unshift(File.join(File.dirname(__FILE__), '..', 'lib'))
+require 'drbman'
+
+Spec::Runner.configure do |config|
+  
+end