set :install_dir,  "/opt/hypertable"
set :hypertable_version, "0.9.8.4"
set :default_pkg, "/tmp/hypertable-0.9.8.4-linux-x86_64.deb"
set :default_fs, "hadoop"
set :default_distro, "cdh4"
set :default_config, "/opt/hypertable/dev-hypertable.cfg"

role :source, "master"
role :master, "master"
role :hyperspace, "master", "hyperspace01", "hyperspace02"
role :slave,  "slave000", "slave001", "slave002", "slave003", "slave004",
              "slave005", "slave006", "slave007", "slave008", "slave009"
role :thriftbroker_additional, "thrift001", "thrift002", "thrift003", "thrift004", "thrift005"
role :spare, "spare000", "spare001"

######################### END OF USER CONFIGURATION ############################
set :prompt_stop, 0
set :prompt_clean, 1

#ssh_options[:forward_agent] = true

def supported_pkgs
  {"bz2"=>1, "deb"=>1, "rpm"=>1, "gz"=>1}
end

def pkg_regex
  '.*\.(bz2|deb|rpm|gz)$'
end

set(:pkg) do
  "#{default_pkg}"
end unless exists?(:pkg)

set(:fs) do
  "#{default_fs}"
end unless exists?(:fs)

set(:distro) do
  "#{default_distro}"
end unless exists?(:distro)

set(:config) do
  "#{default_config}"
end unless exists?(:config)

set :config_file, "#{config}".split('/')[-1]
set :config_option, \
    "--config=#{install_dir}/#{hypertable_version}/conf/#{config_file}"

set(:thriftbroker_args) do
  ""
end unless exists?(:thriftbroker_args)

 desc <<-DESC
    Copies config file into installation.
    This task runs on the source machine (#{roles[:source].servers[0]}) and \
    copies the config file specified by the variable 'config' (#{config}) \
    to the installation directory specified by the variable 'install_dir' \
    (#{install_dir})
 DESC
task :copy_config_local, :roles => :source do
  run <<-CMD
    echo #{distro} > #{install_dir}/#{hypertable_version}/conf/hadoop-distro &&
    rsync -e \"ssh -o StrictHostKeyChecking=no\" #{config} #{install_dir}/#{hypertable_version}/conf
  CMD
end
 
 desc <<-DESC
    Copies config file to installation on all servers in cluster.
    This task copies the dir\
    #{roles[:source].servers[0]}:#{install_dir}/{#hypertable_version}/conf
    to all machines in the cluster    
 DESC
task :push_config_all do
   run <<-CMD
     rsync -av -e "ssh -o StrictHostKeyChecking=no" --exclude=log --exclude=run --exclude=demo --exclude=fs --exclude=hyperspace #{roles[:source].servers[0]}:#{install_dir}/#{hypertable_version}/conf/ #{install_dir}/#{hypertable_version}/conf 
  CMD
end

 desc <<-DESC
    Pushes config file to all macines.
    This task copies the config file (#{config}), specified by the 'config' \
    variable, into the conf/ directory of the installation on the source \
    machine (#{roles[:source].servers[0]}) and then rsyncs it to all machines.
 DESC
task :push_config do
  copy_config_local
  push_config_all 
end

desc <<-DESC
    Set Hadoop distro (#{distro})
 DESC
task :set_distro do
  run "echo #{distro} > #{install_dir}/#{hypertable_version}/conf/hadoop-distro"
end

 desc <<-DESC
    rsyncs installation directory to cluster.  For each machine in the \
    cluster, his commannd  rsyncs the installation from the source \
    installation machine specified by the role 'source' \
    (default=#{roles[:source].servers[0]})
 DESC
task :rsync do
  run <<-CMD
     rsync -av -e "ssh -o StrictHostKeyChecking=no" --exclude=log --exclude=run --exclude=demo --exclude=fs --exclude=conf --exclude=hyperspace #{roles[:source].servers[0]}:#{install_dir}/#{hypertable_version} #{install_dir} &&
     rsync -av -e "ssh -o StrictHostKeyChecking=no" --exclude=log --exclude=run --exclude=demo --exclude=fs --exclude=hyperspace #{roles[:source].servers[0]}:#{install_dir}/#{hypertable_version}/conf/ #{install_dir}/#{hypertable_version}/conf 
  CMD
end

 desc <<-DESC
    sets up the symbolic link 'current' in the installation area \
    to point to the directory of the current version
    (default=#{hypertable_version})
 DESC
task :set_current, :roles => [:master, :hyperspace, :slave, :thriftbroker_additional, :spare] do
  run <<-CMD
   cd #{install_dir} &&
   rm -f current &&
   ln -s #{hypertable_version} current
  CMD
end

 desc <<-DESC
   Distributes installation.  This task rsyncs everything under\
   #{roles[:source].servers[0]}:#{install_dir}/#{hypertable_version} to #{install_dir}\
   on all machines in the cluster
 DESC
task :dist, :roles => :source do
  run "echo #{distro} > #{install_dir}/#{hypertable_version}/conf/hadoop-distro"
  transaction do
    rsync
  end
end

 desc <<-DESC
   Alias for install_package command 
 DESC
task :install_pkg do
  install_package
end

 desc <<-DESC
    rsyncs binary packages and installs on each machine in the cluster
 DESC
task :install_package, :roles => [:master, :hyperspace, :slave, :thriftbroker_additional, :spare] do
  pkg_basename = File.basename(pkg)
  pkg_basename =~ /#{pkg_regex}/
  pkg_type = $1
  
  if (!supported_pkgs.has_key?(pkg_type))
    raise "Package file #{pkg} is of unsupported type. Expected one of #{supported_pkgs.keys.inspect}"
  end
  if (/-#{hypertable_version}-/ =~ pkg_basename).nil?
    raise "Package #{pkg} doesn't match version #{hypertable_version}"
  end

  run("rsync -e \"ssh -o StrictHostKeyChecking=no\" #{roles[:source].servers[0]}:#{pkg} #{install_dir}/")
   
  if (pkg_type == "deb")
    run("dpkg -i #{install_dir}/#{pkg_basename} && rm #{install_dir}/#{pkg_basename}") 
  elsif (pkg_type == "rpm")
    run("rpm -ivh --replacepkgs --nomd5 --nodeps --oldpackage #{install_dir}/#{pkg_basename} && rm #{install_dir}/#{pkg_basename}")
  elsif (pkg_type == "bz2")
    ext_index = pkg_basename.index(".tar.bz2")
    if ext_index.nil?
      raise "Package #{pkg} is not a .tar.bz2 archive"
    end
    dir_basename = pkg_basename[0..ext_index-1]
    run("tar xjv -f #{install_dir}/#{pkg_basename} -C #{install_dir}")
    run("mv #{install_dir}/#{dir_basename}/opt/hypertable/#{hypertable_version} #{install_dir}");
    run("/bin/rm -rf #{install_dir}/#{dir_basename}");
  elsif (pkg_type == "gz")
    ext_index = pkg_basename.index(".tar.gz")
    if ext_index.nil?
      raise "Package #{pkg} is not a .tar.gz archive"
    end
    dir_basename = pkg_basename[0..ext_index-1]
    run("tar xzv -f #{install_dir}/#{pkg_basename} -C #{install_dir}")
    run("mv #{install_dir}/#{dir_basename}/opt/hypertable/#{hypertable_version} #{install_dir}");
    run("/bin/rm -rf #{install_dir}/#{dir_basename}");
  else
    raise "Unrecognized package type - #{pkg}"
  end
  run "echo #{distro} > #{install_dir}/#{hypertable_version}/conf/hadoop-distro"
end

 desc <<-DESC
    fhsize's the installations
 DESC
task :fhsize do
  transaction do 
    run <<-CMD
      #{install_dir}/#{hypertable_version}/bin/fhsize.sh &&
      echo #{distro} > #{install_dir}/#{hypertable_version}/conf/hadoop-distro
    CMD
  end
end

 desc <<-DESC
    Upgrades installation.  Checks upgrade, fhsizes if needed 
    then copies hyperspace and the rangeserver
    state in the run/ directory to new installation
 DESC
task :upgrade do 
  transaction do
    qualify_upgrade
    upgrade_all
    set_current
  end
end

 desc <<-DESC
   Verify that upgrade is OK.
 DESC
task :qualify_upgrade, :roles => :source do
  run <<-CMD
    #{install_dir}/#{hypertable_version}/bin/upgrade-ok.sh \
        #{install_dir}/current #{hypertable_version}
  CMD
end

 desc <<-DESC
    Upgrades (copies or uses previous symlink) for "hyperspace", "conf", "run", "log" 
    and "fs" dirs from the current installation to
    installation specified by the hypertable_version
    (#{hypertable_version})
DESC
task :upgrade_all, :roles => [:master, :hyperspace, :slave, :thriftbroker_additional, :spare] do
  run <<-CMD
    #{install_dir}/#{hypertable_version}/bin/upgrade.sh \
        #{install_dir}/current #{hypertable_version}
  CMD
end

desc "Starts all processes."
task :start do
  transaction do
    start_servers
    start_thriftbrokers
  end
end

desc "Stops all processes."
task :stop do
  transaction do
    stop_thriftbrokers
    stop_servers
  end
end

desc "Starts Hyperspace, Master, and Slave processes."
task :start_servers do
  transaction do
    start_hyperspace
    start_master
    start_slaves
  end
end

desc "Starts hyperspace processes."
task :start_hyperspace, :roles => :hyperspace do
  run <<-CMD
   #{install_dir}/current/bin/start-hyperspace.sh \
      #{config_option}
  CMD
end

desc "Starts master processes."
task :start_master, :roles => :master do
  run <<-CMD
   #{install_dir}/current/bin/start-fsbroker.sh #{fs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-master.sh #{config_option} &&
   #{install_dir}/current/bin/start-monitoring.sh
  CMD
end

desc "Starts monitoring server."
task :start_monitoring, :roles => :master do
  #{install_dir}/current/bin/start-monitoring.sh
end

desc "Starts slave processes."
task :start_slaves, :roles => :slave do
  run <<-CMD
   #{install_dir}/current/bin/random-wait.sh 5 &&
   #{install_dir}/current/bin/start-fsbroker.sh #{fs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-rangeserver.sh \
      #{config_option}
  CMD
end

desc "Starts ThriftBrokers."
task :start_thriftbrokers do
  transaction do
    start_thriftbrokers_primary
    if !roles[:thriftbroker_additional].empty?
      start_thriftbrokers_additional
    end
  end
end

desc "Stops ThriftBrokers."
task :stop_thriftbrokers do
  transaction do
    stop_thriftbrokers_primary
    if !roles[:thriftbroker_additional].empty?
      stop_thriftbrokers_additional
    end
  end
end

desc "Starts ThriftBroker on master and slave machines."
task :start_thriftbrokers_primary, :roles => [:master, :slave] do
  run <<-CMD
   #{install_dir}/current/bin/random-wait.sh 5 &&
   #{install_dir}/current/bin/start-thriftbroker.sh \
      #{config_option} #{thriftbroker_args}
  CMD
end

desc "Stops ThriftBroker on master and slave machines."
task :stop_thriftbrokers_primary, :roles => [:master, :slave] do
  run <<-CMD
   #{install_dir}/current/bin/stop-servers.sh thriftbroker
  CMD
end

desc "Starts ThriftBroker on thriftbroker_additional machines."
task :start_thriftbrokers_additional, :roles => :thriftbroker_additional do
  run <<-CMD
   #{install_dir}/current/bin/random-wait.sh 5 &&
   #{install_dir}/current/bin/start-fsbroker.sh #{fs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-thriftbroker.sh \
      #{config_option} #{thriftbroker_args}
  CMD
end

desc "Stops ThriftBroker on thriftbroker_additional machines."
task :stop_thriftbrokers_additional, :roles => :thriftbroker_additional do
  run <<-CMD
   #{install_dir}/current/bin/stop-servers.sh thriftbroker fsbroker
  CMD
end

desc "Starts FS brokers."
task :start_fsbrokers, :roles => [:master, :slave] do
  run "#{install_dir}/current/bin/start-fsbroker.sh #{fs} \
      #{config_option}"
end

desc "Stops Slave, Master, and Hyperspace processes."
task :stop_servers do
   
  if (prompt_stop == 1)
    puts "This will STOP ALL SERVERS in the Hypertable instance with MASTER='#{roles[:master].first}'. Proceed? ('yes' to proceed)"
    value = STDIN.gets.chomp rescue nil
    exit unless value == "yes"
  end
 
  transaction do
    stop_master
    stop_slaves
    stop_hyperspace
    stop_fsbrokers
  end
end

desc "Stops FS brokers."
task :stop_fsbrokers, :roles => [:master, :slave] do
  run <<-CMD
  #{install_dir}/current/bin/stop-servers.sh
  CMD
end

desc "Stops slave processes."
task :stop_slaves, :roles => :slave do
  run <<-CMD
  #{install_dir}/current/bin/stop-servers.sh --no-hyperspace --no-master --no-fsbroker
  CMD
end

desc "Stops master processes."
task :stop_master, :roles => :master do
  run <<-CMD
  #{install_dir}/current/bin/stop-servers.sh --no-hyperspace --no-rangeserver --no-fsbroker &&
  #{install_dir}/current/bin/stop-monitoring.sh
  CMD
end

desc "Stops monitoring server."
task :stop_monitoring, :roles => :master do
  #{install_dir}/current/bin/stop-monitoring.sh
end

desc "Stops hyperspace processes."
task :stop_hyperspace, :roles => :hyperspace do
  run <<-CMD 
  #{install_dir}/current/bin/stop-hyperspace.sh
  CMD
end

desc "Destroys the database, removing all tables."
task :cleandb do
  
  if prompt_clean == 1
    puts "This will DELETE ALL DATA stored in the Hypertable instance with MASTER='#{roles[:master].first}'. ARE YOU SURE you want to proceed? ('Yes' to proceed)"
    value = STDIN.gets.chomp rescue nil
    exit unless value == "Yes"
  end

  transaction do
    clean_master
    clean_hyperspace
    clean_slaves
  end  
end

desc "Cleans master state but not hyperspace."
task :clean_master, :roles => :master do
  run <<-CMD
   #{install_dir}/current/bin/start-fsbroker.sh #{fs} \
      #{config_option} && \
   #{install_dir}/current/bin/clean-database.sh #{config_option} && \ 
   #{install_dir}/current/bin/stop-monitoring.sh ;
  CMD
end

desc "Cleans hyperspace."
task :clean_hyperspace, :roles => :hyperspace do
  run <<-CMD
   #{install_dir}/current/bin/clean-hyperspace.sh
  CMD
end

desc "Cleans rangeservers and master state but not hyperspace."
task :clean_slaves, :roles => :slave do
  run <<-CMD
   #{install_dir}/current/bin/stop-servers.sh --no-hyperspace --no-master &&
   rm -rf #{install_dir}/current/run/*
  CMD
end

desc "Reports status for all processes."
task :status do
  transaction do
    fs_status
    master_status
    hyperspace_status
    rangeserver_status
  end
end

desc "Get status for fs processes."
task :fs_status, :roles => [:master, :slave] do
  run <<-CMD
   #{install_dir}/current/bin/ht serverup fsbroker
  CMD
end

desc "Get status for Hypertable.Master process."
task :master_status, :roles => [:master] do
  run <<-CMD
   #{install_dir}/current/bin/ht serverup master
  CMD
end

desc "Get status for Hyperspace.Master process."
task :hyperspace_status, :roles => [:hyperspace] do
  srand
  sleep(rand(4))
  run <<-CMD
   #{install_dir}/current/bin/ht serverup hyperspace
  CMD
end

desc "Get status for rangeserver processes."
task :rangeserver_status, :roles => [:slave] do
  run <<-CMD
   #{install_dir}/current/bin/ht serverup rangeserver
  CMD
end

set :default_dumpfile, "/tmp/rsdump.txt"

set(:dumpfile) do
  "#{default_dumpfile}"
end unless exists?(:dumpfile)

desc "Run dump command on each rangeserver"
task :rangeserver_dump, :roles => [:slave] do
  run <<-CMD
   echo "dump NOKEYS '#{dumpfile}';" | #{install_dir}/current/bin/ht ht_rsclient --batch #{config_option}
  CMD
end