retryable.rb in mongo-2.5.0

- old
+ new

@@ -28,11 +28,10 @@
     #     ...
     #   end
     #
     # @note This only retries read operations on socket errors.
     #
-    # @param [ Integer ] attempt The retry attempt count - for internal use.
     # @param [ Proc ] block The block to execute.
     #
     # @yieldparam [ Server ] server The server to which the write should be sent.
     #
     # @return [ Result ] The result of the operation.
@@ -97,27 +96,71 @@
     # @param [ Proc ] block The block to execute.
     #
     # @return [ Result ] The result of the operation.
     #
     # @since 2.1.0
-    def write_with_retry(session, server_selector)
+    def write_with_retry(session, write_concern, &block)
+      unless retry_write_allowed?(session, write_concern)
+        return legacy_write_with_retry(&block)
+      end
+
+      server = cluster.next_primary
+      unless server.retry_writes?
+        return legacy_write_with_retry(server, &block)
+      end
+
+      begin
+        txn_num = session.next_txn_num
+        yield(server, txn_num)
+      rescue Error::SocketError, Error::SocketTimeoutError => e
+        retry_write(e, txn_num, &block)
+      rescue Error::OperationFailure => e
+        raise e unless e.write_retryable?
+        retry_write(e, txn_num, &block)
+      end
+    end
+
+    private
+
+    def retry_write_allowed?(session, write_concern)
+      session && session.retry_writes? &&
+          (write_concern.nil? || write_concern.acknowledged?)
+    end
+
+    def retry_write(original_error, txn_num, &block)
+      cluster.scan!
+      server = cluster.next_primary
+      raise original_error unless (server.retry_writes? && txn_num)
+      log_retry(original_error)
+      yield(server, txn_num)
+    rescue Error::SocketError, Error::SocketTimeoutError => e
+      cluster.scan!
+      raise e
+    rescue Error::OperationFailure => e
+      raise original_error unless e.write_retryable?
+      cluster.scan!
+      raise e
+    rescue
+      raise original_error
+    end
+
+    def legacy_write_with_retry(server = nil)
       attempt = 0
       begin
         attempt += 1
-        yield(server_selector.call)
+        yield(server || cluster.next_primary)
       rescue Error::OperationFailure => e
+        server = nil
         raise(e) if attempt > Cluster::MAX_WRITE_RETRIES
         if e.write_retryable?
           log_retry(e)
           cluster.scan!
           retry
         else
           raise(e)
         end
       end
     end
-
-    private
 
     # Log a warning so that any application slow down is immediately obvious.
     def log_retry(e)
       Logger.logger.warn "Retry due to: #{e.class.name} #{e.message}"
     end