lib/mongo/retryable.rb in mongo-2.5.3 vs lib/mongo/retryable.rb in mongo-2.6.0

@@ -1,6 +1,6 @@
-# Copyright (C) 2015 MongoDB, Inc.
+# Copyright (C) 2015-2018 MongoDB, Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
@@ -28,29 +28,30 @@
     #     ...
     #   end
     #
     # @note This only retries read operations on socket errors.
     #
+    # @param [ Mongo::Session ] session The session that the operation is being run on.
     # @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.
     #
     # @since 2.1.0
-    def read_with_retry
+    def read_with_retry(session = nil)
       attempt = 0
       begin
         attempt += 1
         yield
       rescue Error::SocketError, Error::SocketTimeoutError => e
-        raise(e) if attempt > cluster.max_read_retries
+        raise(e) if attempt > cluster.max_read_retries || (session && session.in_transaction?)
         log_retry(e)
         cluster.scan!
         retry
       rescue Error::OperationFailure => e
-        if cluster.sharded? && e.retryable?
+        if cluster.sharded? && e.retryable? && !(session && session.in_transaction?)
           raise(e) if attempt > cluster.max_read_retries
           log_retry(e)
           sleep(cluster.read_retry_interval)
           retry
         else
@@ -91,41 +92,45 @@
     #   end
     #
     # @note This only retries operations on not master failures, since it is
     #   the only case we can be sure a partial write did not already occur.
     #
+    # @param [ true | false ] ending_transaction True if the write operation is abortTransaction or
+    #   commitTransaction, false otherwise.
     # @param [ Proc ] block The block to execute.
     #
     # @return [ Result ] The result of the operation.
     #
     # @since 2.1.0
-    def write_with_retry(session, write_concern, &block)
-      unless retry_write_allowed?(session, write_concern)
-        return legacy_write_with_retry(&block)
+    def write_with_retry(session, write_concern, ending_transaction = false, &block)
+      unless retry_write_allowed?(session, write_concern) || ending_transaction
+        return legacy_write_with_retry(nil, session, &block)
       end
 
       server = cluster.next_primary
-      unless server.retry_writes?
-        return legacy_write_with_retry(server, &block)
+
+      unless server.retry_writes? || ending_transaction
+        return legacy_write_with_retry(server, session, &block)
       end
 
       begin
-        txn_num = session.next_txn_num
+        txn_num = session.in_transaction? ? session.txn_num : session.next_txn_num
         yield(server, txn_num)
       rescue Error::SocketError, Error::SocketTimeoutError => e
+        raise e if session.in_transaction? && !ending_transaction
         retry_write(e, txn_num, &block)
       rescue Error::OperationFailure => e
-        raise e unless e.write_retryable?
+        raise e if (session.in_transaction? && !ending_transaction) || !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?)
+          (write_concern.nil? || write_concern.acknowledged?) or false
     end
 
     def retry_write(original_error, txn_num, &block)
       cluster.scan!
       server = cluster.next_primary
@@ -141,18 +146,21 @@
       raise e
     rescue
       raise original_error
     end
 
-    def legacy_write_with_retry(server = nil)
+    def legacy_write_with_retry(server = nil, session = nil)
+      # This is the pre-session retry logic, and is not subject to
+      # current retryable write specifications.
+      # In particular it does not retry on SocketError and SocketTimeoutError.
       attempt = 0
       begin
         attempt += 1
         yield(server || cluster.next_primary)
       rescue Error::OperationFailure => e
         server = nil
         raise(e) if attempt > Cluster::MAX_WRITE_RETRIES
-        if e.write_retryable?
+        if e.write_retryable? && !(session && session.in_transaction?)
           log_retry(e)
           cluster.scan!
           retry
         else
           raise(e)