riak.rb

namespace :riak do
  desc "wait for service to be up before continuing"
  task :wait_for_service, :roles => :riak do
    sudo 'riak-admin wait-for-service riak_kv riak@`hostname -f`'
  end

  desc "restart each node in the cluster serially"
  task :rolling_restart, :roles => :riak do
    serially do
      sudo '/etc/init.d/riak restart'
      wait_for_service
    end
  end
end

def serially(&block)
  original = ENV['HOSTS']
  find_servers_for_task(self.current_task).each do |server|
    ENV['HOSTS'] = server.host
    yield
  end
ensure
  ENV['HOSTS'] = original
end