-
-
Save iamjwc/d5ea38c973abdaddeeb0 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
require 'mongo' | |
module Mongo | |
class RetryUtil | |
RETRIABLE_EXCEPTIONS = [ | |
Mongo::ConnectionFailure, | |
Mongo::ConnectionTimeoutError, | |
Mongo::OperationFailure, | |
] | |
class << self | |
def exponential_retry(reconnect_max_sleep_time = 20) | |
def time_to_sleep(n) | |
Math.exp(n * 0.25) - 1 | |
end | |
elapsed_time = 0 | |
# If the current attempt is less than the number of reconnect attempts, | |
# then we should return true and keep retrying. If not, we have already | |
# retried a bunch, so return false and stop retrying. | |
def should_retry_exponentially?(current_attempt, elapsed_time, max) | |
time_to_sleep(current_attempt) + elapsed_time <= max | |
end | |
# Return if we should retry or not. If we should retry, sleep for proper amount of time. | |
lambda do |n| | |
if should_retry_exponentially?(n, elapsed_time, reconnect_max_sleep_time) | |
elapsed_time += time_to_sleep(n) | |
sleep(time_to_sleep(n)) | |
true | |
end | |
end | |
end | |
def retry_with(&blk) | |
@retry_with = blk if blk | |
@retry_with | |
end | |
def should_retry?(number_of_retries) | |
@retry_with && @retry_with.call(number_of_retries) | |
end | |
def with_retry(&blk) | |
number_of_retries = 0 | |
begin | |
yield | |
rescue *RETRIABLE_EXCEPTIONS | |
if self.should_retry?(number_of_retries) | |
number_of_retries += 1 | |
retry | |
else | |
raise $! | |
end | |
end | |
end | |
end | |
# By default, retry_with exponential_retry. | |
retry_with(&exponential_retry) | |
end | |
end | |
class Mongo::Collection | |
alias :insert_without_retry :insert | |
def insert(*args) | |
Mongo::RetryUtil.with_retry { insert_without_retry(*args) } | |
end | |
alias :<< :insert | |
alias :update_without_retry :update | |
def update(*args) | |
Mongo::RetryUtil.with_retry { update_without_retry(*args) } | |
end | |
alias :remove_without_retry :remove | |
def remove(*args) | |
Mongo::RetryUtil.with_retry { remove_without_retry(*args) } | |
end | |
end | |
class Mongo::Cursor | |
alias :next_document_without_retry :next_document | |
def next_document(*args) | |
Mongo::RetryUtil.with_retry { next_document_without_retry(*args) } | |
end | |
end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
require File.expand_path(File.dirname(__FILE__) + '/../../spec_helper') | |
require File.expand_path(File.dirname(__FILE__) + '/../../../lib/mongo/retriable_commands') | |
describe Mongo::RetryUtil do | |
it "should call .retry_with block if Mongo::ConnectionFailure is thrown" do | |
called = false | |
Mongo::RetryUtil.retry_with {|n| called = true } | |
called.should be_false | |
Mongo::RetryUtil.with_retry { raise Mongo::ConnectionFailure if !called } | |
called.should be_true | |
end | |
it "should reraise the error if it shouldn't retry" do | |
Mongo::RetryUtil.retry_with {|n| false } | |
lambda { | |
Mongo::RetryUtil.with_retry { raise Mongo::ConnectionFailure } | |
}.should raise_error(Mongo::ConnectionFailure) | |
end | |
describe "checking commands" do | |
before do | |
@called = false | |
Mongo::RetryUtil.retry_with {|n| @called = true; false } | |
@coll = Mongo::Collection.allocate | |
@cur = Mongo::Cursor.allocate | |
end | |
it "should be called if Collection#insert raises ConnectionFailure" do | |
def @coll.insert_without_retry(*args) | |
raise Mongo::ConnectionFailure | |
end | |
@called.should be_false | |
@coll.insert rescue nil | |
@called.should be_true | |
end | |
it "should be called if Collection#update raises ConnectionFailure" do | |
def @coll.update_without_retry(*args) | |
raise Mongo::ConnectionFailure | |
end | |
@called.should be_false | |
@coll.update rescue nil | |
@called.should be_true | |
end | |
it "should be called if Collection#remove raises ConnectionFailure" do | |
def @coll.remove_without_retry(*args) | |
raise Mongo::ConnectionFailure | |
end | |
@called.should be_false | |
@coll.remove rescue nil | |
@called.should be_true | |
end | |
it "should be called if Cursor#next_document raises ConnectionFailure" do | |
def @cur.next_document_without_retry(*args) | |
raise Mongo::ConnectionFailure | |
end | |
$c = @cur | |
def @coll.find(*args) | |
$c | |
end | |
@called.should be_false | |
@coll.find_one rescue nil | |
@called.should be_true | |
end | |
end | |
end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment