retriable_commands.rb

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

retriable_commands_spec.rb

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