Class: Mongo::Session

# File 'lib/mongo/session.rb', line 1357

def causal_consistency?
  @causal_consistency ||= if @options.key?(:causal_consistency)
                            !!@options[:causal_consistency]
                          else
                            true
                          end
end

# File 'lib/mongo/session.rb', line 128

attr_reader :client

# File 'lib/mongo/session.rb', line 130

attr_reader :cluster

# File 'lib/mongo/session.rb', line 902

def committing_transaction?
  !!@committing_transaction
end

# File 'lib/mongo/session.rb', line 158

def dirty?
  @server_session&.dirty?
end

# File 'lib/mongo/session.rb', line 255

def ended?
  !!@ended
end

# File 'lib/mongo/session.rb', line 190

def explicit?
  !implicit?
end

# File 'lib/mongo/session.rb', line 178

def implicit?
  @implicit ||= !!(@options.key?(:implicit) && @options[:implicit] == true)
end

# File 'lib/mongo/session.rb', line 894

def in_transaction?
  within_states?(STARTING_TRANSACTION_STATE, TRANSACTION_IN_PROGRESS_STATE)
end

# File 'lib/mongo/session.rb', line 1311

def inside_with_transaction?
  @inside_with_transaction
end

# File 'lib/mongo/session.rb', line 1252

def materialized?
  raise Error::SessionEnded if ended?

  !@server_session.nil?
end

# File 'lib/mongo/session.rb', line 141

attr_reader :operation_time

# File 'lib/mongo/session.rb', line 123

attr_reader :options

# File 'lib/mongo/session.rb', line 292

attr_reader :pinned_connection_global_id

# File 'lib/mongo/session.rb', line 286

attr_reader :pinned_server

# File 'lib/mongo/session.rb', line 298

attr_accessor :recovery_token

# File 'lib/mongo/session.rb', line 210

def retry_reads?
  client.options[:retry_reads] != false
end

# File 'lib/mongo/session.rb', line 225

def retry_writes?
  !!client.options[:retry_writes] && (cluster.replica_set? || cluster.sharded? || cluster.load_balanced?)
end

# File 'lib/mongo/session.rb', line 134

def snapshot?
  !!options[:snapshot]
end

# File 'lib/mongo/session.rb', line 1291

attr_reader :snapshot_timestamp

# File 'lib/mongo/session.rb', line 1301

def snapshot_timestamp=(value)
  @snapshot_timestamp ||= value
end

# File 'lib/mongo/session.rb', line 882

def starting_transaction?
  within_states?(STARTING_TRANSACTION_STATE)
end

# File 'lib/mongo/session.rb', line 1307

attr_reader :with_transaction_deadline

# File 'lib/mongo/session.rb', line 1003

def add_autocommit!(command)
  command.tap do |c|
    c[:autocommit] = false if in_transaction?
  end
end

# File 'lib/mongo/session.rb', line 1018

def add_start_transaction!(command)
  command.tap do |c|
    c[:startTransaction] = true if starting_transaction?
  end
end

# File 'lib/mongo/session.rb', line 1033

def add_txn_num!(command)
  command.tap do |c|
    c[:txnNumber] = BSON::Int64.new(@server_session.txn_num) if in_transaction?
  end
end

# File 'lib/mongo/session.rb', line 1048

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?
      # https://jira.mongodb.org/browse/SPEC-1161: transaction's
      # read concern overrides collection/database/client read concerns,
      # even if transaction's read concern is not set.
      # Read concern here is the one sent to the server and may
      # include afterClusterTime.
      if rc = c[:readConcern]
        rc = rc.dup
        rc.delete(:level)
      end
      if txn_read_concern
        if rc
          rc.update(txn_read_concern)
        else
          rc = txn_read_concern.dup
        end
      end
      if rc.nil? || rc.empty?
        c.delete(:readConcern)
      else
        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.
    c[:readConcern] = Options::Mapper.transform_values_to_strings(c[:readConcern]) if c[:readConcern]

    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 @already_committed
        wc = BSON::Document.new(c[:writeConcern] || txn_write_concern || {})
        wc.merge!(w: :majority)
        wc[:wtimeout] ||= 10_000
        c[:writeConcern] = wc
      elsif txn_write_concern
        c[:writeConcern] ||= txn_write_concern
      end
    end

    # A non-numeric write concern w value needs to be sent as a string rather than a symbol.
    if c[:writeConcern] && c[:writeConcern][:w] && c[:writeConcern][:w].is_a?(Symbol)
      c[:writeConcern][:w] = c[:writeConcern][:w].to_s
    end

    # Ignore wtimeout if csot
    c[:writeConcern]&.delete(:wtimeout) if context&.csot?

    # We must not send an empty (server default) write concern.
    c.delete(:writeConcern) if c[:writeConcern] && c[:writeConcern].empty?
  end
end

# File 'lib/mongo/session.rb', line 1227

def advance_operation_time(new_operation_time)
  @operation_time = if @operation_time
                      [ @operation_time, new_operation_time ].max
                    else
                      new_operation_time
                    end
end

# File 'lib/mongo/session.rb', line 1432

def backoff_seconds_for_retry(transaction_attempt)
  exponential = BACKOFF_INITIAL * (1.5**(transaction_attempt - 1))
  Random.rand * [ exponential, BACKOFF_MAX ].min
end

# File 'lib/mongo/session.rb', line 1437

def backoff_would_exceed_deadline?(deadline, backoff_seconds)
  return false if deadline.zero?

  Utils.monotonic_time + backoff_seconds >= deadline
end

# File 'lib/mongo/session.rb', line 1402

def calculate_with_transaction_deadline(opts)
  calc = lambda { |timeout|
    if timeout == 0
      0
    else
      Utils.monotonic_time + (timeout / 1000.0)
    end
  }
  if timeout_ms = opts&.dig(:timeout_ms)
    calc.call(timeout_ms)
  elsif default_timeout_ms = @options[:default_timeout_ms]
    calc.call(default_timeout_ms)
  elsif @client.timeout_ms
    calc.call(@client.timeout_ms)
  end
end

# File 'lib/mongo/session.rb', line 1351

def causal_consistency_doc
  return unless operation_time && causal_consistency?

  { afterClusterTime: operation_time }
end

# File 'lib/mongo/session.rb', line 1371

def check_if_ended!
  raise Mongo::Error::InvalidSession.new(SESSION_ENDED_ERROR_MSG) if ended?
end

# File 'lib/mongo/session.rb', line 1336

def check_if_no_transaction!
  return unless within_states?(NO_TRANSACTION_STATE)

  raise Mongo::Error::InvalidTransactionOperation.new(
    Mongo::Error::InvalidTransactionOperation::NO_TRANSACTION_STARTED
  )
end

# File 'lib/mongo/session.rb', line 1375

def check_matching_cluster!(client)
  return unless cluster != client.cluster

  raise Mongo::Error::InvalidSession.new(MISMATCHED_CLUSTER_ERROR_MSG)
end

# File 'lib/mongo/session.rb', line 1381

def check_transactions_supported!
  raise Mongo::Error::TransactionsNotSupported, 'standalone topology' if cluster.single?
end

# File 'lib/mongo/session.rb', line 728

def commit_transaction(options = nil)
  QueryCache.clear
  check_if_ended!
  check_if_no_transaction!

  if within_states?(TRANSACTION_ABORTED_STATE)
    raise Mongo::Error::InvalidTransactionOperation.new(
      Mongo::Error::InvalidTransactionOperation.cannot_call_after_msg(
        :abortTransaction, :commitTransaction
      )
    )
  end

  options ||= {}

  begin
    # If commitTransaction is called twice, we need to run the same commit
    # operation again, so we revert the session to the previous state.
    if within_states?(TRANSACTION_COMMITTED_STATE)
      @state = @last_commit_skipped ? STARTING_TRANSACTION_STATE : TRANSACTION_IN_PROGRESS_STATE
      @already_committed = true
    end

    if starting_transaction?
      @last_commit_skipped = true
    else
      @last_commit_skipped = false
      @committing_transaction = true

      write_concern = options[:write_concern] || txn_options[:write_concern]
      write_concern = WriteConcern.get(write_concern) if write_concern && !write_concern.is_a?(WriteConcern::Base)

      context = Operation::Context.new(
        client: @client,
        session: self,
        operation_timeouts: operation_timeouts(options)
      )
      write_with_retry(write_concern, ending_transaction: true,
                                      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] ||= 10_000
            write_concern = WriteConcern.get(wco)
          else
            write_concern = WriteConcern.get(w: :majority, wtimeout: 10_000)
          end
        end
        spec = {
          selector: { commitTransaction: 1 },
          db_name: 'admin',
          session: self,
          txn_num: txn_num,
          write_concern: write_concern,
        }
        operation = Operation::Command.new(spec)
        tracer.trace_operation(operation, context, op_name: 'commitTransaction') do
          operation.execute_with_connection(connection, context: context)
        end
      end
    end
    # Finish the transaction span before changing state
    tracer.finish_transaction_span(self)
  ensure
    @state = TRANSACTION_COMMITTED_STATE
    @committing_transaction = false
  end

  # No official return value, but return true so that in interactive
  # use the method hints that it succeeded.
  true
end

# File 'lib/mongo/session.rb', line 1419

def deadline_expired?(deadline)
  if deadline.zero?
    false
  else
    Utils.monotonic_time >= deadline
  end
end

# File 'lib/mongo/session.rb', line 150

def dirty!(mark = true)
  @server_session&.dirty!(mark)
end

# File 'lib/mongo/session.rb', line 385

def end_session
  if !ended? && @client
    if within_states?(TRANSACTION_IN_PROGRESS_STATE)
      begin
        abort_transaction
      rescue Mongo::Error, Error::AuthError
      end
    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

# File 'lib/mongo/session.rb', line 363

def inspect
  "#<Mongo::Session:0x#{object_id} session_id=#{session_id} options=#{@options}>"
end

# File 'lib/mongo/session.rb', line 1446

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

# File 'lib/mongo/session.rb', line 1241

def materialize_if_needed
  raise Error::SessionEnded if ended?

  return unless implicit? && !@server_session

  @server_session = cluster.session_pool.checkout

  self
end

# File 'lib/mongo/session.rb', line 1267

def next_txn_num
  raise Error::SessionEnded if ended?

  @server_session.next_txn_num
end

# File 'lib/mongo/session.rb', line 1385

def operation_timeouts(opts)
  {
    inherited_timeout_ms: @with_transaction_timeout_ms || @client.timeout_ms
  }.tap do |result|
    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

# File 'lib/mongo/session.rb', line 936

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

# File 'lib/mongo/session.rb', line 919

def pin_to_server(server)
  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

# File 'lib/mongo/session.rb', line 1201

def process(result)
  unless implicit?
    set_operation_time(result)
    if cluster_time_doc = result.cluster_time
      advance_cluster_time(cluster_time_doc)
    end
  end
  @server_session.set_last_use!

  if (doc = result.reply && result.reply.documents.first) && doc[:recoveryToken]
    self.recovery_token = doc[:recoveryToken]
  end

  result
end

# File 'lib/mongo/session.rb', line 1152

def revert_to_starting_transaction!
  return unless within_states?(TRANSACTION_IN_PROGRESS_STATE)

  @state = STARTING_TRANSACTION_STATE
end

# File 'lib/mongo/session.rb', line 267

def session_id
  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
  # will not have a session_id until materialized, thus calls to
  # session_id might fail. An application should not have an opportunity
  # 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.
  raise Error::SessionNotMaterialized unless materialized?

  @server_session.session_id
end

# File 'lib/mongo/session.rb', line 1365

def set_operation_time(result)
  return unless result && result.operation_time

  @operation_time = result.operation_time
end

# File 'lib/mongo/session.rb', line 664

def start_transaction(options = nil)
  check_transactions_supported!

  if options
    Lint.validate_read_concern_option(options[:read_concern])

    #         # 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

  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
    )
  end

  unpin

  next_txn_num
  @txn_options = (@options[:default_transaction_options] || {}).merge(options || {})

  if txn_write_concern && !WriteConcern.get(txn_write_concern).acknowledged?
    raise Mongo::Error::InvalidTransactionOperation.new(
      Mongo::Error::InvalidTransactionOperation::UNACKNOWLEDGED_WRITE_CONCERN
    )
  end

  @state = STARTING_TRANSACTION_STATE
  @already_committed = false
  tracer.start_transaction_span(self)

  # This method has no explicit return value.
  # We could return nil here but true indicates to the user that the
  # operation succeeded. This is intended for interactive use.
  # Note that the return value is not documented.
  true
end

# File 'lib/mongo/session.rb', line 1116

def suppress_read_write_concern!(command)
  command.tap do |c|
    next unless in_transaction?

    c.delete(:readConcern) unless starting_transaction?
    c.delete(:writeConcern) unless c[:commitTransaction] || c[:abortTransaction]
  end
end

# File 'lib/mongo/session.rb', line 1281

def txn_num
  raise Error::SessionEnded if ended?

  @server_session.txn_num
end

# File 'lib/mongo/session.rb', line 166

def txn_options
  @txn_options or raise ArgumentError, 'There is no active transaction'
end

# File 'lib/mongo/session.rb', line 1327

def txn_read_concern
  # Read concern is inherited from client but not db or collection.
  txn_options[:read_concern] || @client.read_concern
end

# File 'lib/mongo/session.rb', line 240

def txn_read_preference
  rp = txn_options[:read] ||
       @client.read_preference
  Mongo::Lint.validate_underscore_read_preference(rp)
  rp
end

# File 'lib/mongo/session.rb', line 1344

def txn_write_concern
  txn_options[:write_concern] ||
    (@client.write_concern && @client.write_concern.options)
end

# File 'lib/mongo/session.rb', line 949

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
  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

# File 'lib/mongo/session.rb', line 982

def unpin_maybe(error, connection = nil)
  if !within_states?(Session::NO_TRANSACTION_STATE) &&
     error.label?('TransientTransactionError')
    unpin(connection)
  end

  if committing_transaction? &&
     error.label?('UnknownTransactionCommitResult')
    unpin(connection)
  end
end

# File 'lib/mongo/session.rb', line 1162

def update_state!
  case @state
  when STARTING_TRANSACTION_STATE
    @state = TRANSACTION_IN_PROGRESS_STATE
  when TRANSACTION_COMMITTED_STATE, TRANSACTION_ABORTED_STATE
    @state = NO_TRANSACTION_STATE
  end
end

# File 'lib/mongo/session.rb', line 1184

def validate!(client)
  check_if_ended!
  check_matching_cluster!(client)
  self
end

# File 'lib/mongo/session.rb', line 1135

def validate_read_preference!(command)
  return unless in_transaction?
  return unless command['$readPreference']

  mode = command['$readPreference']['mode'] || command['$readPreference'][:mode]

  return unless mode && mode != 'primary'

  raise Mongo::Error::InvalidTransactionOperation.new(
    "read preference in a transaction must be primary (requested: #{mode})"
  )
end

# File 'lib/mongo/session.rb', line 458

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.
               Utils.monotonic_time + 120
             end
  transaction_in_progress = false
  transaction_attempt = 0
  last_error = nil
  overload_error_count = 0
  overload_encountered = false

  loop do
    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}")
        # 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
        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

      raise
    else
      if within_states?(TRANSACTION_ABORTED_STATE, NO_TRANSACTION_STATE, TRANSACTION_COMMITTED_STATE)
        transaction_in_progress = false
        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?)
            transaction_in_progress = false

            raise unless @with_transaction_timeout_ms && deadline_expired?(deadline)

            make_timeout_error_from(e, 'CSOT timeout expired during withTransaction commit')
          end

          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
            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
        else
          transaction_in_progress = false
          raise
        end
      rescue Error::AuthError
        transaction_in_progress = false
        raise
      end
    end
  end

  # No official return value, but return true so that in interactive
  # use the method hints that it succeeded.
  true
ensure
  if transaction_in_progress
    log_warn('with_transaction callback broke out of with_transaction loop, aborting transaction')
    begin
      abort_transaction
    rescue Error::OperationFailure::Family, Error::InvalidTransactionOperation
    end
  end
  @with_transaction_deadline = nil
  @with_transaction_timeout_ms = nil
  @inside_with_transaction = false
end

# File 'lib/mongo/session.rb', line 1332

def within_states?(*states)
  states.include?(@state)
end