Spread ec-restore across Front Ends

parallel-restore.rb

#!/usr/bin/env ruby

require 'fileutils'
require 'optparse'

ARGV << '-h' if ARGV.length != 6

$options = {}

OptionParser.new do |opts|
  opts.banner = "Usage: #{$PROGRAM_NAME} [options]"
  opts.on('-d', '--backup-dir PATH', '/path/to/backup/directory') { |v| $options[:backup_dir] = v }
  opts.on('-s', '--batch-size SIZE', 'Number of FrontEnds to generate scripts for') { |v| $options[:batch_size] = v }
  opts.on('-i', '--batch-id ID', 'A number from 1 to Batch Size') { |v| $options[:batch_id] = v }
  opts.on_tail('-h', '--help', 'Show this message') do
    puts opts
    exit
  end
end.parse!

def org_path(org)
  File.expand_path(File.join($options[:backup_dir], 'organizations', org))
end

def tree_size(org_path)
  Dir[File.join(org_path, '**/*')].size
end

def organizations_path
  File.expand_path(File.join($options[:backup_dir], 'organizations'))
end

def populate_org_hash
  orgs = {}
  Dir[File.join(organizations_path, '**')].each do |org_path|
    org = File.basename(org_path)
    orgs[org] = {}
    orgs[org]['path'] = org_path
    orgs[org]['size'] = tree_size(org_path)
  end
  orgs
end

def batch_orgs_by_size(orgs)
  sorted_list = Hash[orgs.sort_by { |_k, v| v['size'] }.reverse]
  index = 0
  h = {}
  sorted_list.each do |s|
    batch = (index % $options[:batch_size].to_i) + 1
    h[batch] ||= []
    h[batch] << s
    index += 1
  end
  h
end

def write_script(batches)
  full_org_list_by_size = []
  batches.each_key do |key|
    full_org_list_by_size << batches[key].map {|k, v| k}.join(' ')
  end

  first_org = full_org_list_by_size.join(' ').split(' ')[0]
  first_time = ''
  this_org_list = batches[$options[:batch_id].to_i].map { |k, _v| k }

  if $options[:batch_id].to_i == 1
    this_org_list.shift
    first_time = "( time knife ec restore #{$options[:backup_dir]} --with-user-sql -c /etc/opscode/pivotal.rb --only-org #{first_org} ) 2>&1 | tee -a restore.log"
  end

  fname = "chef_restore-#{$options[:batch_id].to_i}of#{$options[:batch_size].to_i}.sh"

  script = <<-EOS
#!/bin/bash

# run this script with sudo -E ie. sudo -E ./#{fname}

SECONDS=0

export PATH=/opt/opscode/embedded/bin:$PATH

# FULL ORG LIST
# #{full_org_list_by_size.join(' ')}
#
# THIS WAS GENERATED FOR SERVER #{$options[:batch_id].to_i} OF #{$options[:batch_size].to_i}
ORGS='#{this_org_list.join(" ")}'

check_for_org () {
 `knife raw -m GET "/organizations" -c /etc/opscode/pivotal.rb | grep -q "\\"#{first_org}\\""`
}

#{first_time}

for org in $ORGS; do
  check_for_org
  while [ $? -ne 0 ]; do
    echo "Checking for existence of Org #{first_org} before starting.."
    sleep 10
    check_for_org
  done
  ( time knife ec restore #{$options[:backup_dir]} --with-user-sql -c /etc/opscode/pivotal.rb --only-org $org --skip-users --skip-useracl ) 2>&1 | tee -a restore.log
done

duration=$SECONDS
echo
echo "A total of $(($duration / 60)) minutes and $(($duration % 60)) seconds elapsed."
EOS
  
  puts "Writing: #{fname}"
  File.open(fname, 'w') do |f|
    f.write(script)
  end
  FileUtils.chmod "u=rwx", "#{fname}"
end

orgs = populate_org_hash
batches = batch_orgs_by_size(orgs)
write_script(batches)

~/ (FE1)$ ./parallel-restore.rb -h
Usage: ./parallel-restore.rb [options]
    -d, --backup-dir PATH            /path/to/backup/directory
    -s, --batch-size SIZE            Number of FrontEnds to generate scripts for
    -i, --batch-id ID                A number from 1 to Batch Size
    -h, --help                       Show this message
~/ (FE1)$ ./parallel-restore.rb -d chef-server-vagrant/backups/ -s 3 -i 1
Writing: chef_restore-1of1.sh
~/ (FE2)$ ./parallel-restore.rb -d chef-server-vagrant/backups/ -s 3 -i 2
Writing: chef_restore-1of2.sh
~/ (FE3)$ ./parallel-restore.rb -d chef-server-vagrant/backups/ -s 3 -i 3
Writing: chef_restore-1of3.sh