diff --git a/lib/puppet/util/retryaction.rb b/lib/puppet/util/retryaction.rb new file mode 100644 index 000000000..a50ceb784 --- /dev/null +++ b/lib/puppet/util/retryaction.rb @@ -0,0 +1,47 @@ +module Puppet::Util::RetryAction + class RetryException < Exception; end + class RetryException::NoBlockGiven < RetryException; end + class RetryException::NoRetriesGiven < RetryException;end + class RetryException::RetriesExceeded < RetryException; end + + def self.retry_action( parameters = { :retry_exceptions => nil, :retries => nil } ) + # Retry actions for a specified amount of time. This method will allow the final + # retry to complete even if that extends beyond the timeout period. + unless block_given? + raise RetryException::NoBlockGiven + end + + raise RetryException::NoRetriesGiven if parameters[:retries].nil? + parameters[:retry_exceptions] ||= Hash.new + + failures = 0 + + begin + yield + rescue Exception => e + # If we were giving exceptions to catch, + # catch the excptions we care about and retry. + # All others fail hard + + raise RetryException::RetriesExceeded, "#{parameters[:retries]} exceeded", e.backtrace if parameters[:retries] == 0 + + if (not parameters[:retry_exceptions].keys.empty?) and parameters[:retry_exceptions].keys.include?(e.class) + Puppet.info("Caught exception #{e.class}:#{e}") + Puppet.info(parameters[:retry_exceptions][e.class]) + elsif (not parameters[:retry_exceptions].keys.empty?) + # If the exceptions is not in the list of retry_exceptions re-raise. + raise e + end + + failures += 1 + parameters[:retries] -= 1 + + # Increase the amount of time that we sleep after every + # failed retry attempt. + sleep (((2 ** failures) -1) * 0.1) + + retry + + end + end +end diff --git a/spec/unit/util/retryaction_spec.rb b/spec/unit/util/retryaction_spec.rb new file mode 100644 index 000000000..28c5fed84 --- /dev/null +++ b/spec/unit/util/retryaction_spec.rb @@ -0,0 +1,62 @@ +#! /usr/bin/env ruby +require 'spec_helper' + +require 'puppet/util/retryaction' + +describe Puppet::Util::RetryAction do + let (:exceptions) {{ Puppet::Error => 'Puppet Error Exception' }} + + it 'should retry on any exception if no acceptable exceptions given' do + Puppet::Util::RetryAction.expects(:sleep).with( (((2 ** 1) -1) * 0.1) ) + Puppet::Util::RetryAction.expects(:sleep).with( (((2 ** 2) -1) * 0.1) ) + + expect do + Puppet::Util::RetryAction.retry_action( :retries => 2 ) do + raise ArgumentError, 'Fake Failure' + end + end.to raise_exception(Puppet::Util::RetryAction::RetryException::RetriesExceeded) + end + + it 'should retry on acceptable exceptions' do + Puppet::Util::RetryAction.expects(:sleep).with( (((2 ** 1) -1) * 0.1) ) + Puppet::Util::RetryAction.expects(:sleep).with( (((2 ** 2) -1) * 0.1) ) + + expect do + Puppet::Util::RetryAction.retry_action( :retries => 2, :retry_exceptions => exceptions) do + raise Puppet::Error, 'Fake Failure' + end + end.to raise_exception(Puppet::Util::RetryAction::RetryException::RetriesExceeded) + end + + it 'should not retry on unacceptable exceptions' do + Puppet::Util::RetryAction.expects(:sleep).never + + expect do + Puppet::Util::RetryAction.retry_action( :retries => 2, :retry_exceptions => exceptions) do + raise ArgumentError + end + end.to raise_exception(ArgumentError) + end + + it 'should succeed if nothing is raised' do + Puppet::Util::RetryAction.expects(:sleep).never + + Puppet::Util::RetryAction.retry_action( :retries => 2) do + true + end + end + + it 'should succeed if an expected exception is raised retried and succeeds' do + should_retry = nil + Puppet::Util::RetryAction.expects(:sleep).once + + Puppet::Util::RetryAction.retry_action( :retries => 2, :retry_exceptions => exceptions) do + if should_retry + true + else + should_retry = true + raise Puppet::Error, 'Fake error' + end + end + end +end