set :source_machine, "{{active_master}}"
set :install_dir,  "/opt/hypertable"
set :hypertable_version, "0.9.4.3"
set :default_dfs, "hadoop"
set :default_config, "/root/mesos-ec2/hypertable/hypertable.cfg"
set :default_additional_args, ""
set :hbase_home, "/opt/hbase/current"
set :default_client_multiplier, 1
set :default_test_driver, "hypertable"
set :default_test_args, ""

role :source, "{{active_master}}"
role :master, "{{active_master}}"
role :hyperspace, "{{active_master}}"
open("/root/mesos-ec2/slaves").each do |slave|
  role :slave, slave
end
role :thriftbroker
role :spare
role :localhost, "{{active_master}}"
role :test_client
role :test_dispatcher

######################### END OF USER CONFIGURATION ############################

def install_machines
  (roles[:master].servers | \
   roles[:hyperspace].servers | \
   roles[:slave].servers | \
   roles[:thriftbroker].servers | \
   roles[:spare].servers | \
   roles[:test_client].servers | \
   roles[:test_dispatcher].servers) - roles[:source].servers
end

role(:install) { install_machines }

set(:dfs) do
  "#{default_dfs}"
end unless exists?(:dfs)

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

set(:additional_args) do
  "#{default_additional_args}"
end unless exists?(:additional_args)

set(:test_driver) do
  "#{default_test_driver}"
end unless exists?(:test_driver)

set(:test_args) do
  "#{default_test_args}"
end unless exists?(:test_args)

set(:client_multiplier) do
  "#{default_client_multiplier}".to_i
end unless exists?(:client_multiplier)

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

 desc <<-DESC
    Copies config file to installation on localhost.
    This task runs on localhost and copies the config file specified \
    by the variable 'config' (default=#{config}) \
    to the installation directory specified by the variable 'install_dir' \
    (default=#{install_dir})
 DESC
task :copy_config, :roles => :localhost do
  run("rsync -e \"ssh -o StrictHostKeyChecking=no\" #{config} #{install_dir}/#{hypertable_version}/conf/")
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 variable 'source_machine' \
    (default=#{source_machine})
 DESC
task :rsync, :roles => :install do
  run <<-CMD
     rsync -av -e "ssh -o StrictHostKeyChecking=no" --exclude=log --exclude=run --exclude=demo --exclude=fs --exclude=hyperspace/ #{source_machine}:#{install_dir}/#{hypertable_version} #{install_dir} &&
     rsync -av -e "ssh -o StrictHostKeyChecking=no" --exclude=log --exclude=run --exclude=demo --exclude=fs --exclude=hyperspace/ #{source_machine}:#{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 => [:install, :source] do
  run <<-CMD
   cd #{install_dir} &&
   rm -f current &&
   ln -s #{hypertable_version} current
  CMD
end

 desc <<-DESC
    Distributes installation.  This task copiles the config file and \
    then rsyncs the installation to each machine in the cluster
 DESC
task :dist do
  transaction do
    copy_config
    rsync
  end
end

 desc <<-DESC
    Distributes and fhsizes the installation, then copies
    config and rsyncs
 DESC
task :fhsize do
  transaction do
    rsync
    fhsize_install
    copy_config
    rsync
  end
end


 desc <<-DESC
    fhsize's the installations
 DESC
task :fhsize_install, :roles => [:install, :source] do
  run <<-CMD
     #{install_dir}/#{hypertable_version}/bin/fhsize.sh
  CMD
end

desc "Verify that upgrade is OK."
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 installation.  Stops servers, copies config, rsyncs
    the installation, then copies hyperspace and the rangeserver
    state in the run/ directory to new installation
 DESC
task :upgrade do
  transaction do
    qualify_upgrade
    stop
    copy_config
    rsync
    upgrade_hyperspace
    upgrade_rangeservers
    set_current
  end
end

 desc <<-DESC
    Upgrades (copies) the Hyperspace database from the current
    installation to the new installation specified by the
    hypertable_version (#{hypertable_version})
 DESC
task :upgrade_hyperspace, :roles => :hyperspace do
  run <<-CMD
    cp -dpR #{install_dir}/current/hyperspace \
       #{install_dir}/#{hypertable_version}
  CMD
end

 desc <<-DESC
    Upgrades (copies) the RangeServers by copying the contents
    of the run directory from the current installation to
    installation specified by the hypertable_version
    (#{hypertable_version})
 DESC
task :upgrade_rangeservers, :roles => :slave do
  run <<-CMD
    cp -dpR #{install_dir}/current/run \
       #{install_dir}/#{hypertable_version}
  CMD
end

desc "Starts all processes."
task :start do
  transaction do
    start_hyperspace
    start_master
    start_slaves
    start_master_thriftbroker
  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-dfsbroker.sh #{dfs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-master.sh #{config_option} &&
   #{install_dir}/current/bin/start-monitoring.sh
  CMD
end

desc "Starts ThriftBroker on master."
task :start_master_thriftbroker, :roles => :master do
  run <<-CMD
   #{install_dir}/current/bin/start-thriftbroker.sh \
      #{config_option}
  CMD
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-dfsbroker.sh #{dfs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-rangeserver.sh \
      #{config_option} &&
   #{install_dir}/current/bin/start-thriftbroker.sh \
      #{config_option}
  CMD
end

desc "Starts ThriftBroker processes."
task :start_thriftbrokers, :roles => :thriftbroker do
  run <<-CMD
   #{install_dir}/current/bin/random-wait.sh 5 &&
   #{install_dir}/current/bin/start-dfsbroker.sh #{dfs} \
      #{config_option} &&
   #{install_dir}/current/bin/start-thriftbroker.sh \
      #{config_option}
  CMD
end


desc "Starts DFS brokers."
task :start_dfsbrokers, :roles => [:master, :slave] do
  run "#{install_dir}/current/bin/start-dfsbroker.sh #{dfs} \
      #{config_option}"
end

desc "Stops all servers."
task :stop do
  transaction do
    stop_slaves
    stop_master
    stop_hyperspace
  end
end

desc "Stops slave processes."
task :stop_slaves, :roles => :slave do
  run <<-CMD
  #{install_dir}/current/bin/stop-servers.sh --no-hyperspace --no-master #{additional_args}
  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 #{additional_args} &&
  #{install_dir}/current/bin/stop-monitoring.sh
  CMD
end

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

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

desc "Cleans hyperspace & rangeservers, removing all tables."
task :cleandb do
  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-dfsbroker.sh #{dfs} \
      #{config_option} && \
   #{install_dir}/current/bin/clean-database.sh #{config_option} ;
  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
    dfs_status
    master_status
    hyperspace_status
    rangeserver_status
  end
end

desc "Get status for dfs processes."
task :dfs_status, :roles => [:master, :slave] do
  run <<-CMD
   #{install_dir}/current/bin/ht serverup dfsbroker
  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
  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


if "#{test_driver}" == "hypertable"
  set :thrift_broker_command, "#{install_dir}/current/bin/start-thriftbroker.sh #{config_option}"
  set :start_test_client_command, "#{install_dir}/current/bin/start-test-client.sh --count #{client_multiplier} #{roles[:test_dispatcher].servers[0]}"
  set :run_test_dispatcher_command, "#{install_dir}/current/bin/jrun --pidfile #{install_dir}/#{hypertable_version}/run/Hypertable.TestDispatcher.pid org.hypertable.examples.PerformanceTest.Dispatcher --driver=#{test_driver} --clients=#{roles[:test_client].servers.length*client_multiplier} #{test_args}"
  set :stop_test_args, ""
elsif "#{test_driver}" == "hbase"
  set :thrift_broker_command, "echo -n"
  set :start_test_client_command, "#{install_dir}/current/bin/start-test-client.sh --jrun-opts \"--add-to-classpath #{hbase_home}/conf\" --count #{client_multiplier} #{roles[:test_dispatcher].servers[0]}"
  set :run_test_dispatcher_command, "#{install_dir}/current/bin/jrun --pidfile #{install_dir}/#{hypertable_version}/run/Hypertable.TestDispatcher.pid --add-to-classpath #{hbase_home}/conf org.hypertable.examples.PerformanceTest.Dispatcher --driver=#{test_driver} --clients=#{roles[:test_client].servers.length*client_multiplier} #{test_args}"
  set :stop_test_args, "--no-thriftbroker --no-dfsbroker"
else
  set :thrift_broker_command, "echo Invalid test driver - #{test_driver}"
  set :start_test_client_command, "echo Invalid test driver - #{test_driver}"
  set :run_test_dispatcher_command, "echo Invalid test driver - #{test_driver}"
  set :stop_test_args, "--no-thriftbroker --no-dfsbroker"
end

desc "Starts test clients."
task :start_test_clients, :roles => :test_client do
  run <<-CMD
   #{install_dir}/current/bin/random-wait.sh 5 &&
   #{thrift_broker_command} &&
   #{start_test_client_command}
  CMD
end

desc "Run test dispatcher."
task :run_test_dispatcher, :roles => :test_dispatcher do
  run <<-CMD
   #{thrift_broker_command} &&
   #{run_test_dispatcher_command}
  CMD
end

desc "Stops test."
task :stop_test, :roles => [:test_client, :test_dispatcher] do
  run <<-CMD
  #{install_dir}/current/bin/stop-servers.sh --no-hyperspace --no-master --no-rangeserver #{stop_test_args}
  CMD
end

desc "Run test"
task :run_test do
  transaction do
    stop_test
    start_test_clients
    run_test_dispatcher
  end
end