mongo 2.23.0 → 2.24.1

data/lib/mongo/session.rb

@@ -1,5 +1,4 @@
 # frozen_string_literal: true
-# rubocop:todo all
 
 # Copyright (C) 2017-2020 MongoDB Inc.
 #
@@ -19,7 +18,6 @@ require 'mongo/session/session_pool'
 require 'mongo/session/server_session'
 
 module Mongo
-
   # A logical session representing a set of sequential operations executed
   # by an application that are related in some way.
   #
@@ -36,10 +34,9 @@ module Mongo
     # Initialize a Session.
     #
     # A session can be explicit or implicit. Lifetime of explicit sessions is
-    # managed by the application - applications explicitry create such sessions
+    # managed by the application - applications explicitly create such sessions
     # and explicitly end them. Implicit sessions are created automatically by
-    # the driver when sending operations to servers that support sessions
-    # (3.6+), and their lifetime is managed by the driver.
+    # the driver, and their lifetime is managed by the driver.
     #
     # When an implicit session is created, it cannot have a server session
     # associated with it. The server session will be checked out of the
@@ -77,6 +74,8 @@ module Mongo
     #     and *:nearest*.
     # @option options [ true | false ] :snapshot Set up the session for
     #   snapshot reads.
+    # @option options [ BSON::Timestamp ] :snapshot_time The desired snapshot
+    #   time for snapshot reads. Only valid when :snapshot is true.
     #
     # @since 2.5.0
     # @api private
@@ -85,24 +84,37 @@ module Mongo
         raise ArgumentError, ':causal_consistency and :snapshot options cannot be both set on a session'
       end
 
+      if options[:snapshot_time] && !options[:snapshot]
+        raise ArgumentError, ':snapshot_time can only be set when :snapshot is true'
+      end
+
+      if options[:snapshot_time] && !options[:snapshot_time].is_a?(BSON::Timestamp)
+        raise ArgumentError, ':snapshot_time must be a BSON::Timestamp'
+      end
+
       if options[:implicit]
         unless server_session.nil?
           raise ArgumentError, 'Implicit session cannot reference server session during construction'
         end
-      else
-        if server_session.nil?
-          raise ArgumentError, 'Explicit session must reference server session during construction'
-        end
+      elsif server_session.nil?
+        raise ArgumentError, 'Explicit session must reference server session during construction'
       end
 
       @server_session = server_session
       options = options.dup
 
-      @client = client.use(:admin)
+      # Implicit sessions only need the cluster and client options (never run
+      # transactions), so avoid creating a Mongo::Client clone to prevent
+      # memory leaks: use the original client directly instead.
+      @client = options[:implicit] ? client : client.use(:admin)
+      @cluster = @client.cluster
       @options = options.dup.freeze
       @cluster_time = nil
       @state = NO_TRANSACTION_STATE
       @with_transaction_deadline = nil
+      @with_transaction_timeout_ms = nil
+      @inside_with_transaction = false
+      @snapshot_timestamp = options[:snapshot_time]
     end
 
     # @return [ Hash ] The options for this session.
@@ -115,9 +127,7 @@ module Mongo
     # @since 2.5.1
     attr_reader :client
 
-    def cluster
-      @client.cluster
-    end
+    attr_reader :cluster
 
     # @return [ true | false ] Whether the session is configured for snapshot
     #   reads.
@@ -154,7 +164,7 @@ module Mongo
     #
     # @since 2.6.0
     def txn_options
-      @txn_options or raise ArgumentError, "There is no active transaction"
+      @txn_options or raise ArgumentError, 'There is no active transaction'
     end
 
     # Is this session an implicit one (not user-created).
@@ -208,8 +218,8 @@ module Mongo
     #
     # @return [ true, false ] If writes will be retried.
     #
-    # @note Retryable writes are only available on server versions at least 3.6
-    #   and with sharded clusters, replica sets, or load-balanced topologies.
+    # @note Retryable writes are only available with sharded clusters, replica
+    #   sets, or load-balanced topologies.
     #
     # @since 2.5.0
     def retry_writes?
@@ -229,7 +239,7 @@ module Mongo
     # @since 2.6.0
     def txn_read_preference
       rp = txn_options[:read] ||
-        @client.read_preference
+           @client.read_preference
       Mongo::Lint.validate_underscore_read_preference(rp)
       rp
     end
@@ -255,9 +265,7 @@ module Mongo
     #
     # @since 2.5.0
     def session_id
-      if ended?
-        raise Error::SessionEnded
-      end
+      raise Error::SessionEnded if ended?
 
       # An explicit session will always have a session_id, because during
       # construction a server session must be provided. An implicit session
@@ -266,9 +274,7 @@ module Mongo
       # to experience this failure because an implicit session shouldn't be
       # accessible to applications due to its lifetime being constrained to
       # operation execution, which is done entirely by the driver.
-      unless materialized?
-        raise Error::SessionNotMaterialized
-      end
+      raise Error::SessionNotMaterialized unless materialized?
 
       @server_session.session_id
     end
@@ -296,20 +302,20 @@ module Mongo
     #
     # @since 2.5.0
     MISMATCHED_CLUSTER_ERROR_MSG = 'The configuration of the client used to create this session does not match that ' +
-        'of the client owning this operation. Please only use this session for operations through its parent ' +
-        'client.'.freeze
+                                   'of the client owning this operation. Please only use this session for operations through its parent ' +
+                                   'client.'
 
     # Error message describing that the session cannot be used because it has already been ended.
     #
     # @since 2.5.0
-    SESSION_ENDED_ERROR_MSG = 'This session has ended and cannot be used. Please create a new one.'.freeze
+    SESSION_ENDED_ERROR_MSG = 'This session has ended and cannot be used. Please create a new one.'
 
     # Error message describing that sessions are not supported by the server version.
     #
     # @since 2.5.0
     # @deprecated
-    SESSIONS_NOT_SUPPORTED = 'Sessions are not supported by the connected servers.'.freeze
-    # Note: SESSIONS_NOT_SUPPORTED is used by Mongoid - do not remove from driver.
+    SESSIONS_NOT_SUPPORTED = 'Sessions are not supported by the connected servers.'
+    # NOTE: SESSIONS_NOT_SUPPORTED is used by Mongoid 7.6 and earlier
 
     # The state of a session in which the last operation was not related to
     # any transaction or no operations have yet occurred.
@@ -384,13 +390,15 @@ module Mongo
           rescue Mongo::Error, Error::AuthError
           end
         end
-        if @server_session
-          @client.cluster.session_pool.checkin(@server_session)
-        end
+        # Release any pinned connection (e.g. after a committed transaction
+        # in load-balanced mode).
+        unpin if pinned_connection_global_id
+        cluster.session_pool.checkin(@server_session) if @server_session
       end
     ensure
       @server_session = nil
       @ended = true
+      @client = nil
     end
 
     # Executes the provided block in a transaction, retrying as necessary.
@@ -448,38 +456,74 @@ module Mongo
     #
     # @since 2.7.0
     def with_transaction(options = nil)
+      @inside_with_transaction = true
+      @with_transaction_timeout_ms = options&.dig(:timeout_ms) || @options[:default_timeout_ms] || @client.timeout_ms
       @with_transaction_deadline = calculate_with_transaction_deadline(options)
       deadline = if @with_transaction_deadline
                    # CSOT enabled, so we have a customer defined deadline.
                    @with_transaction_deadline
                  else
-                    # CSOT not enabled, so we use the default deadline, 120 seconds.
+                   # CSOT not enabled, so we use the default deadline, 120 seconds.
                    Utils.monotonic_time + 120
                  end
       transaction_in_progress = false
+      transaction_attempt = 0
+      last_error = nil
+      overload_error_count = 0
+      overload_encountered = false
+
       loop do
-        commit_options = {}
-        if options
-          commit_options[:write_concern] = options[:write_concern]
+        if transaction_attempt > 0
+          if overload_encountered
+            delay = @client.retry_policy.backoff_delay(overload_error_count)
+            if backoff_would_exceed_deadline?(deadline, delay)
+              make_timeout_error_from(last_error, 'CSOT timeout expired waiting to retry withTransaction')
+            end
+            raise(last_error) unless @client.retry_policy.should_retry_overload?(overload_error_count, delay)
+
+            sleep(delay)
+          else
+            backoff = backoff_seconds_for_retry(transaction_attempt)
+            if backoff_would_exceed_deadline?(deadline, backoff)
+              make_timeout_error_from(last_error, 'CSOT timeout expired waiting to retry withTransaction')
+            end
+
+            sleep(backoff)
+          end
         end
+
+        commit_options = {}
+        commit_options[:write_concern] = options[:write_concern] if options
         start_transaction(options)
         transaction_in_progress = true
+        transaction_attempt += 1
+
         begin
           rv = yield self
         rescue Exception => e
           if within_states?(STARTING_TRANSACTION_STATE, TRANSACTION_IN_PROGRESS_STATE)
             log_warn("Aborting transaction due to #{e.class}: #{e}")
-            @with_transaction_deadline = nil
+            # CSOT: if the deadline is already expired, clear it so that
+            # abort_transaction uses a fresh timeout (not the expired deadline).
+            # If the deadline is not yet expired, keep it so abort uses remaining time.
+            @with_transaction_deadline = nil if @with_transaction_deadline && deadline_expired?(deadline)
             abort_transaction
             transaction_in_progress = false
           end
 
           if deadline_expired?(deadline)
             transaction_in_progress = false
-            raise
+            make_timeout_error_from(e, 'CSOT timeout expired during withTransaction callback')
           end
 
           if e.is_a?(Mongo::Error) && e.label?('TransientTransactionError')
+            last_error = e
+            if e.label?('SystemOverloadedError')
+              overload_encountered = true
+              overload_error_count += 1
+            elsif overload_encountered
+              overload_error_count += 1
+            end
             next
           end
 
@@ -490,32 +534,71 @@ module Mongo
             return rv
           end
 
+          # CSOT: if the timeout has expired before we can commit, abort the
+          # transaction instead and raise a client-side timeout error.
+          if @with_transaction_deadline && deadline_expired?(deadline)
+            transaction_in_progress = false
+            @with_transaction_deadline = nil
+            abort_transaction
+            raise Mongo::Error::TimeoutError, 'CSOT timeout expired before transaction could be committed'
+          end
+
           begin
             commit_transaction(commit_options)
             transaction_in_progress = false
             return rv
           rescue Mongo::Error => e
             if e.label?('UnknownTransactionCommitResult')
-              if  deadline_expired?(deadline) ||
-                e.is_a?(Error::OperationFailure::Family) && e.max_time_ms_expired?
-              then
+              if deadline_expired?(deadline) ||
+                 (e.is_a?(Error::OperationFailure::Family) && e.max_time_ms_expired?)
                 transaction_in_progress = false
-                raise
+
+                raise unless @with_transaction_timeout_ms && deadline_expired?(deadline)
+
+                make_timeout_error_from(e, 'CSOT timeout expired during withTransaction commit')
               end
-              wc_options = case v = commit_options[:write_concern]
-                when WriteConcern::Base
-                  v.options
-                when nil
-                  {}
-                else
-                  v
+
+              if e.label?('SystemOverloadedError')
+                overload_encountered = true
+                overload_error_count += 1
+              elsif overload_encountered
+                overload_error_count += 1
+              end
+
+              if overload_encountered
+                delay = @client.retry_policy.backoff_delay(overload_error_count)
+                if backoff_would_exceed_deadline?(deadline, delay)
+                  transaction_in_progress = false
+                  make_timeout_error_from(e, 'CSOT timeout expired during withTransaction commit')
                 end
+                unless @client.retry_policy.should_retry_overload?(overload_error_count, delay)
+                  transaction_in_progress = false
+                  raise
+                end
+                sleep(delay)
+              end
+
+              wc_options = case v = commit_options[:write_concern]
+                           when WriteConcern::Base
+                             v.options
+                           when nil
+                             {}
+                           else
+                             v
+                           end
               commit_options[:write_concern] = wc_options.merge(w: :majority)
               retry
             elsif e.label?('TransientTransactionError')
               if Utils.monotonic_time >= deadline
                 transaction_in_progress = false
-                raise
+                make_timeout_error_from(e, 'CSOT timeout expired during withTransaction commit')
+              end
+              last_error = e
+              if e.label?('SystemOverloadedError')
+                overload_encountered = true
+                overload_error_count += 1
+              elsif overload_encountered
+                overload_error_count += 1
               end
               @state = NO_TRANSACTION_STATE
               next
@@ -542,6 +625,8 @@ module Mongo
         end
       end
       @with_transaction_deadline = nil
+      @with_transaction_timeout_ms = nil
+      @inside_with_transaction = false
     end
 
     # Places subsequent operations in this session into a new transaction.
@@ -556,7 +641,7 @@ module Mongo
     #
     # @option options [ Integer ] :max_commit_time_ms The maximum amount of
     #   time to allow a single commitTransaction command to run, in milliseconds.
-    #   This options is deprecated, use :timeout_ms instead.
+    #   This option is deprecated, use :timeout_ms instead.
     # @option options [ Hash ] :read_concern The read concern options hash,
     #   with the following optional keys:
     #   - *:level* -- the read preference level as a symbol; valid values
@@ -582,28 +667,25 @@ module Mongo
       if options
         Lint.validate_read_concern_option(options[:read_concern])
 
-=begin
-        # It would be handy to detect invalid read preferences here, but
-        # some of the spec tests require later detection of invalid read prefs.
-        # Maybe we can do this when lint mode is on.
-        mode = options[:read] && options[:read][:mode].to_s
-        if mode && mode != 'primary'
-          raise Mongo::Error::InvalidTransactionOperation.new(
-            "read preference in a transaction must be primary (requested: #{mode})"
-          )
-        end
-=end
+        #         # It would be handy to detect invalid read preferences here, but
+        #         # some of the spec tests require later detection of invalid read prefs.
+        #         # Maybe we can do this when lint mode is on.
+        #         mode = options[:read] && options[:read][:mode].to_s
+        #         if mode && mode != 'primary'
+        #           raise Mongo::Error::InvalidTransactionOperation.new(
+        #             "read preference in a transaction must be primary (requested: #{mode})"
+        #           )
+        #         end
       end
 
-      if snapshot?
-        raise Mongo::Error::SnapshotSessionTransactionProhibited
-      end
+      raise Mongo::Error::SnapshotSessionTransactionProhibited if snapshot?
 
       check_if_ended!
 
       if within_states?(STARTING_TRANSACTION_STATE, TRANSACTION_IN_PROGRESS_STATE)
         raise Mongo::Error::InvalidTransactionOperation.new(
-          Mongo::Error::InvalidTransactionOperation::TRANSACTION_ALREADY_IN_PROGRESS)
+          Mongo::Error::InvalidTransactionOperation::TRANSACTION_ALREADY_IN_PROGRESS
+        )
       end
 
       unpin
@@ -613,7 +695,8 @@ module Mongo
 
       if txn_write_concern && !WriteConcern.get(txn_write_concern).acknowledged?
         raise Mongo::Error::InvalidTransactionOperation.new(
-          Mongo::Error::InvalidTransactionOperation::UNACKNOWLEDGED_WRITE_CONCERN)
+          Mongo::Error::InvalidTransactionOperation::UNACKNOWLEDGED_WRITE_CONCERN
+        )
       end
 
       @state = STARTING_TRANSACTION_STATE
@@ -642,7 +725,7 @@ module Mongo
     # @raise [ Error::InvalidTransactionOperation ] If there is no active transaction.
     #
     # @since 2.6.0
-    def commit_transaction(options=nil)
+    def commit_transaction(options = nil)
       QueryCache.clear
       check_if_ended!
       check_if_no_transaction!
@@ -650,7 +733,9 @@ module Mongo
       if within_states?(TRANSACTION_ABORTED_STATE)
         raise Mongo::Error::InvalidTransactionOperation.new(
           Mongo::Error::InvalidTransactionOperation.cannot_call_after_msg(
-            :abortTransaction, :commitTransaction))
+            :abortTransaction, :commitTransaction
+          )
+        )
       end
 
       options ||= {}
@@ -670,9 +755,7 @@ module Mongo
           @committing_transaction = true
 
           write_concern = options[:write_concern] || txn_options[:write_concern]
-          if write_concern && !write_concern.is_a?(WriteConcern::Base)
-            write_concern = WriteConcern.get(write_concern)
-          end
+          write_concern = WriteConcern.get(write_concern) if write_concern && !write_concern.is_a?(WriteConcern::Base)
 
           context = Operation::Context.new(
             client: @client,
@@ -680,15 +763,14 @@ module Mongo
             operation_timeouts: operation_timeouts(options)
           )
           write_with_retry(write_concern, ending_transaction: true,
-            context: context,
-          ) do |connection, txn_num, context|
-            if context.retry?
+                                          context: context) do |connection, txn_num, context|
+            if context.retry? && !context.overload_only_retry?
               if write_concern
                 wco = write_concern.options.merge(w: :majority)
-                wco[:wtimeout] ||= 10000
+                wco[:wtimeout] ||= 10_000
                 write_concern = WriteConcern.get(wco)
               else
-                write_concern = WriteConcern.get(w: :majority, wtimeout: 10000)
+                write_concern = WriteConcern.get(w: :majority, wtimeout: 10_000)
               end
             end
             spec = {
@@ -738,12 +820,15 @@ module Mongo
       if within_states?(TRANSACTION_COMMITTED_STATE)
         raise Mongo::Error::InvalidTransactionOperation.new(
           Mongo::Error::InvalidTransactionOperation.cannot_call_after_msg(
-            :commitTransaction, :abortTransaction))
+            :commitTransaction, :abortTransaction
+          )
+        )
       end
 
       if within_states?(TRANSACTION_ABORTED_STATE)
         raise Mongo::Error::InvalidTransactionOperation.new(
-          Mongo::Error::InvalidTransactionOperation.cannot_call_twice_msg(:abortTransaction))
+          Mongo::Error::InvalidTransactionOperation.cannot_call_twice_msg(:abortTransaction)
+        )
       end
 
       options ||= {}
@@ -757,21 +842,18 @@ module Mongo
             operation_timeouts: operation_timeouts(options)
           )
           write_with_retry(txn_options[:write_concern],
-            ending_transaction: true, context: context,
-          ) do |connection, txn_num, context|
-            begin
-              operation = Operation::Command.new(
-                selector: { abortTransaction: 1 },
-                db_name: 'admin',
-                session: self,
-                txn_num: txn_num
-              )
-              tracer.trace_operation(operation, context, op_name: 'abortTransaction') do
-                operation.execute_with_connection(connection, context: context)
-              end
-            ensure
-              unpin
+                           ending_transaction: true, context: context) do |connection, txn_num, context|
+            operation = Operation::Command.new(
+              selector: { abortTransaction: 1 },
+              db_name: 'admin',
+              session: self,
+              txn_num: txn_num
+            )
+            tracer.trace_operation(operation, context, op_name: 'abortTransaction') do
+              operation.execute_with_connection(connection, context: context)
             end
+          ensure
+            unpin
           end
         end
 
@@ -835,14 +917,12 @@ module Mongo
     #
     # @api private
     def pin_to_server(server)
-      if server.nil?
-        raise ArgumentError, 'Cannot pin to a nil server'
-      end
-      if Lint.enabled?
-        unless server.mongos?
-          raise Error::LintError, "Attempted to pin the session to server #{server.summary} which is not a mongos"
-        end
+      raise ArgumentError, 'Cannot pin to a nil server' if server.nil?
+
+      if Lint.enabled? && !server.mongos?
+        raise Error::LintError, "Attempted to pin the session to server #{server.summary} which is not a mongos"
       end
+
       @pinned_server = server
     end
 
@@ -850,13 +930,14 @@ module Mongo
     #
     # @param [ Integer ] connection_global_id The global id of connection to pin
     # this session to.
+    # @param [ Connection | nil ] connection The connection object to pin to.
     #
     # @api private
-    def pin_to_connection(connection_global_id)
-      if connection_global_id.nil?
-        raise ArgumentError, 'Cannot pin to a nil connection id'
-      end
+    def pin_to_connection(connection_global_id, connection: nil)
+      raise ArgumentError, 'Cannot pin to a nil connection id' if connection_global_id.nil?
+
       @pinned_connection_global_id = connection_global_id
+      @pinned_connection = connection
     end
 
     # Unpins this session from the pinned server or connection,
@@ -866,9 +947,25 @@ module Mongo
     #
     # @api private
     def unpin(connection = nil)
+      # Idempotent: if there is no pinned state to clear, do nothing. Nested
+      # unpin_maybe handlers (e.g. in BulkWrite#execute_operation wrapping an
+      # OpMsg execution that already calls unpin_maybe in its own do_execute)
+      # can call this method twice for the same error; checking the connection
+      # back into the pool a second time would raise from the pool.
+      return if @pinned_server.nil? && @pinned_connection.nil? && @pinned_connection_global_id.nil?
+
       @pinned_server = nil
       @pinned_connection_global_id = nil
-      connection.unpin unless connection.nil?
+      conn = connection || @pinned_connection
+      if conn
+        conn.unpin(:transaction)
+        # Only check the connection back into the pool if nothing else
+        # still holds a pin on it (e.g. an open cursor).
+        unless conn.pinned?
+          conn.connection_pool.check_in(conn)
+        end
+      end
+      @pinned_connection = nil
     end
 
     # Unpins this session from the pinned server or connection, if the session was pinned
@@ -884,14 +981,12 @@ module Mongo
     # @api private
     def unpin_maybe(error, connection = nil)
       if !within_states?(Session::NO_TRANSACTION_STATE) &&
-        error.label?('TransientTransactionError')
-      then
+         error.label?('TransientTransactionError')
         unpin(connection)
       end
 
       if committing_transaction? &&
-        error.label?('UnknownTransactionCommitResult')
-      then
+         error.label?('UnknownTransactionCommitResult')
         unpin(connection)
       end
     end
@@ -922,9 +1017,7 @@ module Mongo
     # @api private
     def add_start_transaction!(command)
       command.tap do |c|
-        if starting_transaction?
-          c[:startTransaction] = true
-        end
+        c[:startTransaction] = true if starting_transaction?
       end
     end
 
@@ -952,7 +1045,7 @@ module Mongo
     #
     # @since 2.6.0
     # @api private
-    def add_txn_opts!(command, read, context)
+    def add_txn_opts!(command, _read, context)
       command.tap do |c|
         # The read concern should be added to any command that starts a transaction.
         if starting_transaction?
@@ -975,27 +1068,23 @@ module Mongo
           if rc.nil? || rc.empty?
             c.delete(:readConcern)
           else
-            c[:readConcern ] = Options::Mapper.transform_values_to_strings(rc)
+            c[:readConcern] = Options::Mapper.transform_values_to_strings(rc)
           end
         end
 
         # We need to send the read concern level as a string rather than a symbol.
-        if c[:readConcern]
-          c[:readConcern] = Options::Mapper.transform_values_to_strings(c[:readConcern])
-        end
+        c[:readConcern] = Options::Mapper.transform_values_to_strings(c[:readConcern]) if c[:readConcern]
 
-        if c[:commitTransaction]
-          if max_time_ms = txn_options[:max_commit_time_ms]
-            c[:maxTimeMS] = max_time_ms
-          end
+        if c[:commitTransaction] && (max_time_ms = txn_options[:max_commit_time_ms])
+          c[:maxTimeMS] = max_time_ms
         end
 
         # The write concern should be added to any abortTransaction or commitTransaction command.
-        if (c[:abortTransaction] || c[:commitTransaction])
+        if c[:abortTransaction] || c[:commitTransaction]
           if @already_committed
             wc = BSON::Document.new(c[:writeConcern] || txn_write_concern || {})
             wc.merge!(w: :majority)
-            wc[:wtimeout] ||= 10000
+            wc[:wtimeout] ||= 10_000
             c[:writeConcern] = wc
           elsif txn_write_concern
             c[:writeConcern] ||= txn_write_concern
@@ -1008,12 +1097,10 @@ module Mongo
         end
 
         # Ignore wtimeout if csot
-        if context&.csot?
-          c[:writeConcern]&.delete(:wtimeout)
-        end
+        c[:writeConcern]&.delete(:wtimeout) if context&.csot?
 
         # We must not send an empty (server default) write concern.
-        c.delete(:writeConcern) if c[:writeConcern]&.empty?
+        c.delete(:writeConcern) if c[:writeConcern] && c[:writeConcern].empty?
       end
     end
 
@@ -1035,7 +1122,7 @@ module Mongo
       end
     end
 
-    # Ensure that the read preference of a command primary.
+    # Ensure that the read preference of a command is primary.
     #
     # @example
     #   session.validate_read_preference!(command)
@@ -1051,11 +1138,21 @@ module Mongo
 
       mode = command['$readPreference']['mode'] || command['$readPreference'][:mode]
 
-      if mode && mode != 'primary'
-        raise Mongo::Error::InvalidTransactionOperation.new(
-          "read preference in a transaction must be primary (requested: #{mode})"
-        )
-      end
+      return unless mode && mode != 'primary'
+
+      raise Mongo::Error::InvalidTransactionOperation.new(
+        "read preference in a transaction must be primary (requested: #{mode})"
+      )
+    end
+
+    # Reverts the session state to STARTING_TRANSACTION_STATE.
+    # Called before retrying the first command in a transaction so that
+    # startTransaction: true is preserved on the retry.
+    # @api private
+    def revert_to_starting_transaction!
+      return unless within_states?(TRANSACTION_IN_PROGRESS_STATE)
+
+      @state = STARTING_TRANSACTION_STATE
     end
 
     # Update the state of the session due to a (non-commit and non-abort) operation being run.
@@ -1085,8 +1182,8 @@ module Mongo
     # @since 2.5.0
     # @api private
     def validate!(client)
-      check_matching_cluster!(client)
       check_if_ended!
+      check_matching_cluster!(client)
       self
     end
 
@@ -1110,10 +1207,8 @@ module Mongo
       end
       @server_session.set_last_use!
 
-      if doc = result.reply && result.reply.documents.first
-        if doc[:recoveryToken]
-          self.recovery_token = doc[:recoveryToken]
-        end
+      if (doc = result.reply && result.reply.documents.first) && doc[:recoveryToken]
+        self.recovery_token = doc[:recoveryToken]
       end
 
       result
@@ -1130,11 +1225,11 @@ module Mongo
     #
     # @since 2.5.0
     def advance_operation_time(new_operation_time)
-      if @operation_time
-        @operation_time = [ @operation_time, new_operation_time ].max
-      else
-        @operation_time = new_operation_time
-      end
+      @operation_time = if @operation_time
+                          [ @operation_time, new_operation_time ].max
+                        else
+                          new_operation_time
+                        end
     end
 
     # If not already set, populate a session objects's server_session by
@@ -1144,9 +1239,7 @@ module Mongo
     #
     # @api private
     def materialize_if_needed
-      if ended?
-        raise Error::SessionEnded
-      end
+      raise Error::SessionEnded if ended?
 
       return unless implicit? && !@server_session
 
@@ -1157,9 +1250,7 @@ module Mongo
 
     # @api private
     def materialized?
-      if ended?
-        raise Error::SessionEnded
-      end
+      raise Error::SessionEnded if ended?
 
       !@server_session.nil?
     end
@@ -1174,9 +1265,7 @@ module Mongo
     # @since 2.5.0
     # @api private
     def next_txn_num
-      if ended?
-        raise Error::SessionEnded
-      end
+      raise Error::SessionEnded if ended?
 
       @server_session.next_txn_num
     end
@@ -1190,20 +1279,39 @@ module Mongo
     #
     # @since 2.6.0
     def txn_num
-      if ended?
-        raise Error::SessionEnded
-      end
+      raise Error::SessionEnded if ended?
 
       @server_session.txn_num
     end
 
+    # @return [ BSON::Timestamp | nil ] The snapshot time for this session.
+    #   nil if the session is not a snapshot session, or if it is a snapshot
+    #   session for which no :snapshot_time option was provided and no read
+    #   has yet captured atClusterTime from the server.
+    attr_reader :snapshot_timestamp
+
+    # Sets the snapshot time for the session. Once set, subsequent
+    # assignments are ignored: snapshotTime is established at most once per
+    # session, either from the :snapshot_time option at construction or from
+    # the atClusterTime returned by the first find/aggregate/distinct
+    # response. This keeps the property effectively read-only for callers,
+    # per the snapshot-sessions spec rationale.
+    #
     # @api private
-    attr_accessor :snapshot_timestamp
+    def snapshot_timestamp=(value)
+      @snapshot_timestamp ||= value
+    end
 
     # @return [ Integer | nil ] The deadline for the current transaction, if any.
     # @api private
     attr_reader :with_transaction_deadline
 
+    # @return [ Boolean ] Whether we are currently inside a with_transaction block.
+    # @api private
+    def inside_with_transaction?
+      @inside_with_transaction
+    end
+
     private
 
     # Get the read concern the session will use when starting a transaction.
@@ -1229,7 +1337,8 @@ module Mongo
       return unless within_states?(NO_TRANSACTION_STATE)
 
       raise Mongo::Error::InvalidTransactionOperation.new(
-        Mongo::Error::InvalidTransactionOperation::NO_TRANSACTION_STARTED)
+        Mongo::Error::InvalidTransactionOperation::NO_TRANSACTION_STARTED
+      )
     end
 
     def txn_write_concern
@@ -1240,25 +1349,23 @@ module Mongo
     # Returns causal consistency document if the last operation time is
     # known and causal consistency is enabled, otherwise returns nil.
     def causal_consistency_doc
-      if operation_time && causal_consistency?
-        {:afterClusterTime => operation_time}
-      else
-        nil
-      end
+      return unless operation_time && causal_consistency?
+
+      { afterClusterTime: operation_time }
     end
 
     def causal_consistency?
-      @causal_consistency ||= (if @options.key?(:causal_consistency)
-                                 !!@options[:causal_consistency]
-                               else
-                                 true
-                               end)
+      @causal_consistency ||= if @options.key?(:causal_consistency)
+                                !!@options[:causal_consistency]
+                              else
+                                true
+                              end
     end
 
     def set_operation_time(result)
-      if result && result.operation_time
-        @operation_time = result.operation_time
-      end
+      return unless result && result.operation_time
+
+      @operation_time = result.operation_time
     end
 
     def check_if_ended!
@@ -1266,40 +1373,34 @@ module Mongo
     end
 
     def check_matching_cluster!(client)
-      if @client.cluster != client.cluster
-        raise Mongo::Error::InvalidSession.new(MISMATCHED_CLUSTER_ERROR_MSG)
-      end
+      return unless cluster != client.cluster
+
+      raise Mongo::Error::InvalidSession.new(MISMATCHED_CLUSTER_ERROR_MSG)
     end
 
     def check_transactions_supported!
-      raise Mongo::Error::TransactionsNotSupported, "standalone topology" if cluster.single?
-
-      cluster.next_primary.with_connection do |conn|
-        if cluster.replica_set? && !conn.features.transactions_enabled?
-          raise Mongo::Error::TransactionsNotSupported, "server version is < 4.0"
-        end
-        if cluster.sharded? && !conn.features.sharded_transactions_enabled?
-          raise Mongo::Error::TransactionsNotSupported, "sharded transactions require server version >= 4.2"
-        end
-      end
+      raise Mongo::Error::TransactionsNotSupported, 'standalone topology' if cluster.single?
     end
 
     def operation_timeouts(opts)
       {
-        inherited_timeout_ms: @client.timeout_ms
+        inherited_timeout_ms: @with_transaction_timeout_ms || @client.timeout_ms
       }.tap do |result|
-        if @with_transaction_deadline.nil?
-          if timeout_ms = opts[:timeout_ms]
-            result[:operation_timeout_ms] = timeout_ms
-          elsif default_timeout_ms = options[:default_timeout_ms]
-            result[:operation_timeout_ms] = default_timeout_ms
+        if @inside_with_transaction
+          if opts[:timeout_ms]
+            raise Mongo::Error::InvalidTransactionOperation,
+                  'timeoutMS cannot be overridden inside a withTransaction callback'
           end
+        elsif timeout_ms = opts[:timeout_ms]
+          result[:operation_timeout_ms] = timeout_ms
+        elsif default_timeout_ms = options[:default_timeout_ms]
+          result[:operation_timeout_ms] = default_timeout_ms
         end
       end
     end
 
     def calculate_with_transaction_deadline(opts)
-      calc = -> (timeout) {
+      calc = lambda { |timeout|
         if timeout == 0
           0
         else
@@ -1322,5 +1423,36 @@ module Mongo
         Utils.monotonic_time >= deadline
       end
     end
+
+    # Exponential backoff settings for with_transaction retries.
+    BACKOFF_INITIAL = 0.005
+    BACKOFF_MAX = 0.5
+    private_constant :BACKOFF_INITIAL, :BACKOFF_MAX
+
+    def backoff_seconds_for_retry(transaction_attempt)
+      exponential = BACKOFF_INITIAL * (1.5**(transaction_attempt - 1))
+      Random.rand * [ exponential, BACKOFF_MAX ].min
+    end
+
+    def backoff_would_exceed_deadline?(deadline, backoff_seconds)
+      return false if deadline.zero?
+
+      Utils.monotonic_time + backoff_seconds >= deadline
+    end
+
+    # Implements makeTimeoutError(lastError) from the transactions-convenient-api spec.
+    # In CSOT mode raises TimeoutError with last_error's message and labels copied.
+    # In non-CSOT mode re-raises last_error directly.
+    def make_timeout_error_from(last_error, timeout_message)
+      if @with_transaction_timeout_ms
+        timeout_error = Mongo::Error::TimeoutError.new("#{timeout_message}: #{last_error}")
+        if last_error.respond_to?(:labels)
+          last_error.labels.each { |label| timeout_error.add_label(label) }
+        end
+        raise timeout_error
+      end
+
+      raise last_error
+    end
   end
 end