patched supply_drop gem and vendored it

11 files changed, 410 insertions, 0 deletions

diff --git a/vendor/supply_drop/lib/supply_drop.rb b/vendor/supply_drop/lib/supply_drop.rb
new file mode 100644
index 0000000..9c2c050
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop.rb
@@ -0,0 +1,12 @@
+require 'supply_drop/rsync'
+require 'supply_drop/async_enumerable'
+require 'supply_drop/plugin'
+require 'supply_drop/syntax_checker'
+require 'supply_drop/thread_pool'
+require 'supply_drop/util'
+require 'supply_drop/writer/batched'
+require 'supply_drop/writer/file'
+require 'supply_drop/writer/streaming'
+require 'supply_drop/tasks'
+
+Capistrano.plugin :supply_drop, SupplyDrop::Plugin
diff --git a/vendor/supply_drop/lib/supply_drop/async_enumerable.rb b/vendor/supply_drop/lib/supply_drop/async_enumerable.rb
new file mode 100644
index 0000000..1d5a116
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/async_enumerable.rb
@@ -0,0 +1,19 @@
+module SupplyDrop
+  module AsyncEnumerable
+    def each(&block)
+      pool = SupplyDrop::ThreadPool.new(SupplyDrop::Util.thread_pool_size)
+      super do |item|
+        pool.schedule(item, &block)
+      end
+      pool.shutdown
+    end
+
+    def map(&block)
+      pool = SupplyDrop::ThreadPool.new(SupplyDrop::Util.thread_pool_size)
+      super do |item|
+        pool.schedule(item, &block)
+      end
+      pool.shutdown
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/plugin.rb b/vendor/supply_drop/lib/supply_drop/plugin.rb
new file mode 100644
index 0000000..348b22b
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/plugin.rb
@@ -0,0 +1,94 @@
+module SupplyDrop
+  module Plugin
+
+    def rsync
+      SupplyDrop::Util.thread_pool_size = puppet_parallel_rsync_pool_size
+      servers = SupplyDrop::Util.optionally_async(find_servers_for_task(current_task), puppet_parallel_rsync)
+      failed_servers = servers.map do |server|
+        rsync_cmd = SupplyDrop::Rsync.command(
+          puppet_source,
+          SupplyDrop::Rsync.remote_address(server.user || fetch(:user, ENV['USER']), server.host, puppet_destination),
+          :delete => true,
+          :excludes => puppet_excludes,
+          :ssh => ssh_options.merge(server.options[:ssh_options]||{})
+        )
+        logger.debug rsync_cmd
+        server.host unless system rsync_cmd
+      end.compact
+
+      raise "rsync failed on #{failed_servers.join(',')}" if failed_servers.any?
+    end
+
+    def prepare
+      run "mkdir -p #{puppet_destination}"
+      run "#{sudo} chown -R $USER: #{puppet_destination}"
+    end
+
+    def noop
+      puppet(:noop)
+    end
+
+    def apply
+      puppet(:apply)
+    end
+
+    def lock
+      if should_lock?
+        run <<-GETLOCK
+if [ ! -f #{puppet_lock_file} ]; then
+    touch #{puppet_lock_file};
+else
+    stat -c "#{red_text("Puppet in progress, #{puppet_lock_file} owned by %U since %x")}" #{puppet_lock_file} >&2;
+    exit 1;
+fi
+        GETLOCK
+      end
+    end
+
+    def unlock
+      run "#{sudo} rm -f #{puppet_lock_file}; true" if should_lock?
+    end
+
+    private
+
+    def should_lock?
+      puppet_lock_file && !ENV['NO_PUPPET_LOCK']
+    end
+
+    def puppet(command = :noop)
+      #puppet_cmd = "cd #{puppet_destination} && #{sudo_cmd} #{puppet_command} --modulepath=#{puppet_lib} #{puppet_parameters}"
+      puppet_cmd = "cd #{puppet_destination} && #{sudo_cmd} #{puppet_command} #{puppet_parameters}"
+      flag = command == :noop ? '--noop' : ''
+
+      writer = if puppet_stream_output
+                 SupplyDrop::Writer::Streaming.new(logger)
+               else
+                 SupplyDrop::Writer::Batched.new(logger)
+               end
+
+      writer = SupplyDrop::Writer::File.new(writer, puppet_write_to_file) unless puppet_write_to_file.nil?
+
+      begin
+        run "#{puppet_cmd} #{flag}" do |channel, stream, data|
+          writer.collect_output(channel[:host], data)
+        end
+        logger.debug "Puppet #{command} complete."
+      ensure
+        writer.all_output_collected
+      end
+    end
+
+    def red_text(text)
+      "\033[0;31m#{text}\033[0m"
+    end
+
+    def sudo_cmd
+      if fetch(:use_sudo, true)
+        sudo(:as => puppet_runner)
+      else
+        logger.info "NOTICE: puppet_runner configuration invalid when use_sudo is false, ignoring..." unless puppet_runner.nil?
+        ''
+      end
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/rsync.rb b/vendor/supply_drop/lib/supply_drop/rsync.rb
new file mode 100644
index 0000000..7356653
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/rsync.rb
@@ -0,0 +1,45 @@
+module SupplyDrop
+  class Rsync
+    class << self
+      def command(from, to, options={})
+        flags = ['-az']
+        flags << '--delete' if options[:delete]
+        flags << excludes(options[:excludes]) if options.has_key?(:excludes)
+        flags << ssh_options(options[:ssh]) if options.has_key?(:ssh)
+
+        "rsync #{flags.compact.join(' ')} #{from} #{to}"
+      end
+
+      def remote_address(user, host, path)
+        user_with_host = [user, host].compact.join('@')
+        [user_with_host, path].join(':')
+      end
+
+      def excludes(patterns)
+        [patterns].flatten.map { |p| "--exclude=#{p}" }
+      end
+
+      def ssh_options(options)
+        mapped_options = options.map do |key, value|
+          next unless value
+
+          #
+          # for a list of the options normally support by Net::SSH (and thus Capistrano), see
+          # http://net-ssh.github.com/net-ssh/classes/Net/SSH.html#method-c-start
+          #
+          case key
+          when :keys then [value].flatten.select { |k| File.exist?(k) }.map { |k| "-i #{k}" }
+          when :config then "-F #{value}"
+          when :port then "-p #{value}"
+          when :user_known_hosts_file then "-o 'UserKnownHostsFile=#{value}'"
+          when :host_key_alias then "-o 'HostKeyAlias=#{value}'"
+          when :paranoid then "-o 'StrictHostKeyChecking=yes'"
+          when :host_name then "-o 'HostName=#{value}'"
+          end
+        end.compact
+
+        %[-e "ssh #{mapped_options.join(' ')}"] unless mapped_options.empty?
+      end
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/syntax_checker.rb b/vendor/supply_drop/lib/supply_drop/syntax_checker.rb
new file mode 100644
index 0000000..fe9cda4
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/syntax_checker.rb
@@ -0,0 +1,21 @@
+module SupplyDrop
+  class SyntaxChecker
+    def initialize(path)
+      @path = path
+    end
+
+    def validate_puppet_files
+      Dir.glob("#{@path}/**/*.pp").map do |puppet_file|
+        output = `puppet parser validate #{puppet_file}`
+        $?.to_i == 0 ? nil : [puppet_file, output]
+      end.compact
+    end
+
+    def validate_templates
+      Dir.glob("#{@path}/**/*.erb").map do |template_file|
+        output = `erb -x -T '-' #{template_file} | ruby -c 2>&1`
+        $?.to_i == 0 ? nil : [template_file, output]
+      end.compact
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/tasks.rb b/vendor/supply_drop/lib/supply_drop/tasks.rb
new file mode 100644
index 0000000..e73b280
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/tasks.rb
@@ -0,0 +1,96 @@
+Capistrano::Configuration.instance.load do
+  namespace :puppet do
+    set :puppet_source, '.'
+    set :puppet_destination, '/tmp/supply_drop'
+    set :puppet_command, 'puppet apply'
+    set :puppet_lib, lambda { "#{puppet_destination}/modules" }
+    set :puppet_parameters, lambda { puppet_verbose ? '--debug --trace puppet.pp' : 'puppet.pp' }
+    set :puppet_verbose, false
+    set :puppet_excludes, %w(.git .svn)
+    set :puppet_stream_output, false
+    set :puppet_parallel_rsync, true
+    set :puppet_parallel_rsync_pool_size, 10
+    set :puppet_syntax_check, false
+    set :puppet_write_to_file, nil
+    set :puppet_runner, nil
+    set :puppet_lock_file, '/tmp/puppet.lock'
+
+    namespace :bootstrap do
+      desc "installs puppet via rubygems on an osx host"
+      task :osx do
+        if fetch(:use_sudo, true)
+          run "#{sudo} gem install puppet --no-ri --no-rdoc"
+        else
+          run "gem install puppet --no-ri --no-rdoc"
+        end
+      end
+
+      desc "installs puppet via apt on an ubuntu host"
+      task :ubuntu do
+        run "mkdir -p #{puppet_destination}"
+        run "#{sudo} apt-get update"
+        run "#{sudo} apt-get install -y puppet rsync"
+      end
+
+      desc "installs puppet via yum on a centos/red hat host"
+      task :redhat do
+        run "mkdir -p #{puppet_destination}"
+        run "#{sudo} yum -y install puppet rsync"
+      end
+    end
+
+    desc "checks the syntax of all *.pp and *.erb files"
+    task :syntax_check do
+      checker = SupplyDrop::SyntaxChecker.new(puppet_source)
+      logger.info "Sytax Checking..."
+      errors = false
+      checker.validate_puppet_files.each do |file, error|
+        logger.important "Puppet error: #{file}"
+        logger.important error
+        errors = true
+      end
+      checker.validate_templates.each do |file, error|
+        logger.important "Template error: #{file}"
+        logger.important error
+        errors = true
+      end
+      raise "syntax errors" if errors
+    end
+
+    desc "pushes the current puppet configuration to the server"
+    task :update_code, :except => { :nopuppet => true } do
+      syntax_check if puppet_syntax_check
+      supply_drop.rsync
+    end
+
+    desc "runs puppet with --noop flag to show changes"
+    task :noop, :except => { :nopuppet => true } do
+      transaction do
+        on_rollback { supply_drop.unlock }
+        supply_drop.prepare
+        supply_drop.lock
+        update_code
+        supply_drop.noop
+        supply_drop.unlock
+      end
+    end
+
+    desc "applies the current puppet config to the server"
+    task :apply, :except => { :nopuppet => true } do
+      transaction do
+        on_rollback { supply_drop.unlock }
+        supply_drop.prepare
+        supply_drop.lock
+        update_code
+        supply_drop.apply
+        supply_drop.unlock
+      end
+    end
+
+    desc "clears the puppet lockfile on the server."
+    task :remove_lock, :except => { :nopuppet => true} do
+      supply_drop.lock
+    end
+  end
+end
+
diff --git a/vendor/supply_drop/lib/supply_drop/thread_pool.rb b/vendor/supply_drop/lib/supply_drop/thread_pool.rb
new file mode 100644
index 0000000..082cf4a
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/thread_pool.rb
@@ -0,0 +1,39 @@
+require 'thread'
+
+module SupplyDrop
+  class ThreadPool
+    def initialize(size)
+      @size = size
+      @jobs = Queue.new
+      @retvals = []
+
+      @pool = Array.new(@size) do |i|
+        Thread.new do
+          Thread.current[:id] = i
+
+          catch(:exit) do
+            loop do
+              job, args = @jobs.pop
+              @retvals << job.call(*args)
+            end
+          end
+        end
+      end
+    end
+
+
+    def schedule(*args, &block)
+      @jobs << [block, args]
+    end
+
+
+    def shutdown
+      @size.times do
+        schedule { throw :exit }
+      end
+
+      @pool.map(&:join)
+      @retvals
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/util.rb b/vendor/supply_drop/lib/supply_drop/util.rb
new file mode 100644
index 0000000..5f4f9f4
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/util.rb
@@ -0,0 +1,23 @@
+module SupplyDrop
+  module Util
+    DEFAULT_THREAD_POOL_SIZE = 10
+
+    def self.thread_pool_size
+      @thread_pool_size ||= DEFAULT_THREAD_POOL_SIZE
+    end
+
+    def self.thread_pool_size=(size)
+      @thread_pool_size = size
+    end
+
+    def self.optionally_async(collection, async)
+      if async
+        async_collection = collection.clone
+        async_collection.extend SupplyDrop::AsyncEnumerable
+        async_collection
+      else
+        collection
+      end
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/writer/batched.rb b/vendor/supply_drop/lib/supply_drop/writer/batched.rb
new file mode 100644
index 0000000..e5fc826
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/writer/batched.rb
@@ -0,0 +1,22 @@
+module SupplyDrop
+  module Writer
+    class Batched
+      def initialize(logger)
+        @outputs = {}
+        @logger = logger
+      end
+
+      def collect_output(host, data)
+        @outputs[host] ||= ""
+        @outputs[host] << data
+      end
+
+      def all_output_collected
+        @outputs.keys.sort.each do |host|
+          @logger.info "Puppet output for #{host}"
+          @logger.debug @outputs[host], host
+        end
+      end
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/writer/file.rb b/vendor/supply_drop/lib/supply_drop/writer/file.rb
new file mode 100644
index 0000000..61454d8
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/writer/file.rb
@@ -0,0 +1,23 @@
+module SupplyDrop
+  module Writer
+    class File
+      def initialize(writer, file)
+        @wrapped_writer = writer
+        @logger = Capistrano::Logger.new(:output => file)
+        @logger.level = Capistrano::Logger::TRACE
+        @file_writer = Batched.new(@logger)
+      end
+
+      def collect_output(host, data)
+        @wrapped_writer.collect_output(host, data)
+        @file_writer.collect_output(host, data)
+      end
+
+      def all_output_collected
+        @wrapped_writer.all_output_collected
+        @file_writer.all_output_collected
+        @logger.close
+      end
+    end
+  end
+end
diff --git a/vendor/supply_drop/lib/supply_drop/writer/streaming.rb b/vendor/supply_drop/lib/supply_drop/writer/streaming.rb
new file mode 100644
index 0000000..e180ec8
--- /dev/null
+++ b/vendor/supply_drop/lib/supply_drop/writer/streaming.rb
@@ -0,0 +1,16 @@
+module SupplyDrop
+  module Writer
+    class Streaming
+      def initialize(logger)
+        @logger = logger
+      end
+
+      def collect_output(host, data)
+        @logger.debug data, host
+      end
+
+      def all_output_collected
+      end
+    end
+  end
+end