mongo 2.20.1 → 2.21.1

data/lib/mongo/crypt/kms/azure/credentials_retriever.rb

@@ -34,13 +34,16 @@ module Mongo
           #   request. This is used for testing.
           # @param [String | nil] metadata_host Azure metadata host. This
           #   is used for testing.
+          # @param [ CsotTimeoutHolder | nil ] timeout_holder CSOT timeout.
           #
           # @return [ KMS::Azure::AccessToken ] Azure access token.
           #
           # @raise [KMS::CredentialsNotFound] If credentials could not be found.
-          def self.fetch_access_token(extra_headers: {}, metadata_host: nil)
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout.
+          def self.fetch_access_token(extra_headers: {}, metadata_host: nil, timeout_holder: nil)
             uri, req = prepare_request(extra_headers, metadata_host)
-            parsed_response = fetch_response(uri, req)
+            parsed_response = fetch_response(uri, req, timeout_holder)
             Azure::AccessToken.new(
               parsed_response.fetch('access_token'),
               Integer(parsed_response.fetch('expires_in'))
@@ -78,13 +81,16 @@ module Mongo
           #
           # @param [URI] uri URI to Azure metadata host.
           # @param [Net::HTTP::Get] req Request object.
+          # @param [ CsotTimeoutHolder | nil ] timeout_holder CSOT timeout.
           #
           # @return [Hash] Parsed response.
           #
           # @raise [KMS::CredentialsNotFound] If cannot fetch response or
           #   response is invalid.
-          def self.fetch_response(uri, req)
-            resp = do_request(uri, req)
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout.
+          def self.fetch_response(uri, req, timeout_holder)
+            resp = do_request(uri, req, timeout_holder)
             if resp.code != '200'
               raise KMS::CredentialsNotFound,
                     "Azure metadata host responded with code #{resp.code}"
@@ -100,12 +106,22 @@ module Mongo
           #
           # @param [URI] uri URI to Azure metadata host.
           # @param [Net::HTTP::Get] req Request object.
+          # @param [ CsotTimeoutHolder | nil ] timeout_holder CSOT timeout.
           #
           # @return [Net::HTTPResponse] Response object.
           #
           # @raise [KMS::CredentialsNotFound] If cannot execute request.
-          def self.do_request(uri, req)
-            ::Timeout.timeout(10) do
+          # @raise Error::TimeoutError if credentials cannot be retrieved within
+          #   the timeout.
+          def self.do_request(uri, req, timeout_holder)
+            timeout_holder&.check_timeout!
+            timeout = timeout_holder&.remaining_timeout_sec || 10
+            exception_class = if timeout_holder&.csot?
+                                Error::TimeoutError
+                              else
+                                nil
+                              end
+            ::Timeout.timeout(timeout, exception_class) do
               Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
                 http.request(req)
               end

data/lib/mongo/crypt/kms/gcp/credentials_retriever.rb

@@ -29,14 +29,20 @@ module Mongo
 
           DEFAULT_HOST = 'metadata.google.internal'
 
-          def self.fetch_access_token
+          # Fetch GCP access token.
+          #
+          # @param [ CsotTimeoutHolder | nil ] timeout_holder CSOT timeout.
+          #
+          # @return [ String ] GCP access token.
+          #
+          # @raise [ KMS::CredentialsNotFound ]
+          # @raise [ Error::TimeoutError ]
+          def self.fetch_access_token(timeout_holder = nil)
             host = ENV.fetch(METADATA_HOST_ENV) { DEFAULT_HOST }
             uri = URI("http://#{host}/computeMetadata/v1/instance/service-accounts/default/token")
             req = Net::HTTP::Get.new(uri)
             req['Metadata-Flavor'] = 'Google'
-            resp = Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
-              http.request(req)
-            end
+            resp = fetch_response(uri, req, timeout_holder)
             if resp.code != '200'
               raise KMS::CredentialsNotFound,
                 "GCE metadata host responded with code #{resp.code}"
@@ -50,6 +56,25 @@ module Mongo
               raise KMS::CredentialsNotFound,
                     "Could not receive GCP metadata response; #{e.class}: #{e.message}"
           end
+
+          def self.fetch_response(uri, req, timeout_holder)
+            timeout_holder&.check_timeout!
+            if timeout_holder&.timeout?
+              ::Timeout.timeout(timeout_holder.remaining_timeout_sec, Error:TimeoutError) do
+                do_fetch(uri, req)
+              end
+            else
+              do_fetch(uri, req)
+            end
+          end
+          private_class_method :fetch_response
+
+          def self.do_fetch(uri, req)
+            Net::HTTP.start(uri.hostname, uri.port, use_ssl: false) do |http|
+              http.request(req)
+            end
+          end
+          private_class_method :do_fetch
         end
       end
     end

data/lib/mongo/csot_timeout_holder.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+# Copyright (C) 2024 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
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  # This class stores operation timeout and provides corresponding helper methods.
+  #
+  # @api private
+  class CsotTimeoutHolder
+    def initialize(session: nil, operation_timeouts: {})
+      @deadline = calculate_deadline(operation_timeouts, session)
+      @operation_timeouts = operation_timeouts
+      @timeout_sec = (@deadline - Utils.monotonic_time if @deadline)
+    end
+
+    attr_reader :deadline, :timeout_sec, :operation_timeouts
+
+    # @return [ true | false ] Whether CSOT is enabled for the operation
+    def csot?
+      !deadline.nil?
+    end
+
+    # @return [ true | false ] Returns false if CSOT is not enabled, or if
+    #   CSOT is set to 0 (means unlimited), otherwise true.
+    def timeout?
+      ![ nil, 0 ].include?(@deadline)
+    end
+
+    # @return [ Float | nil ] Returns the remaining seconds of the timeout
+    #   set for the operation; if no timeout is set, or the timeout is 0
+    #   (means unlimited) returns nil.
+    def remaining_timeout_sec
+      return nil unless timeout?
+
+      deadline - Utils.monotonic_time
+    end
+
+    def remaining_timeout_sec!
+      check_timeout!
+      remaining_timeout_sec
+    end
+
+    # @return [ Integer | nil ] Returns the remaining milliseconds of the timeout
+    #   set for the operation; if no timeout is set, or the timeout is 0
+    #   (means unlimited) returns nil.
+    def remaining_timeout_ms
+      seconds = remaining_timeout_sec
+      return nil if seconds.nil?
+
+      (seconds * 1_000).to_i
+    end
+
+    def remaining_timeout_ms!
+      check_timeout!
+      remaining_timeout_ms
+    end
+
+    # @return [ true | false ] Whether the timeout for the operation expired.
+    #   If no timeout set, this method returns false.
+    def timeout_expired?
+      if timeout?
+        Utils.monotonic_time >= deadline
+      else
+        false
+      end
+    end
+
+    # Check whether the operation timeout expired, and raises an appropriate
+    # error if yes.
+    #
+    # @raise [ Error::TimeoutError ]
+    def check_timeout!
+      return unless timeout_expired?
+
+      raise Error::TimeoutError, "Operation took more than #{timeout_sec} seconds"
+    end
+
+    private
+
+    def calculate_deadline(opts = {}, session = nil)
+      check_no_override_inside_transaction!(opts, session)
+      return session&.with_transaction_deadline if session&.with_transaction_deadline
+
+      if (operation_timeout_ms = opts[:operation_timeout_ms])
+        calculate_deadline_from_timeout_ms(operation_timeout_ms)
+      elsif (inherited_timeout_ms = opts[:inherited_timeout_ms])
+        calculate_deadline_from_timeout_ms(inherited_timeout_ms)
+      end
+    end
+
+    def check_no_override_inside_transaction!(opts, session)
+      return unless opts[:operation_timeout_ms] && session&.with_transaction_deadline
+
+      raise ArgumentError, 'Cannot override timeout_ms inside with_transaction block'
+    end
+
+    def calculate_deadline_from_timeout_ms(operation_timeout_ms)
+      if operation_timeout_ms.positive?
+        Utils.monotonic_time + (operation_timeout_ms / 1_000.0)
+      elsif operation_timeout_ms.zero?
+        0
+      elsif operation_timeout_ms.negative?
+        raise ArgumentError, "timeout_ms must be a non-negative integer but #{operation_timeout_ms} given"
+      end
+    end
+  end
+end

data/lib/mongo/cursor/kill_spec.rb

@@ -31,7 +31,8 @@ module Mongo
         db_name:,
         connection_global_id:,
         server_address:,
-        session:
+        session:,
+        connection: nil
       )
         @cursor_id = cursor_id
         @coll_name = coll_name
@@ -39,6 +40,7 @@ module Mongo
         @connection_global_id = connection_global_id
         @server_address = server_address
         @session = session
+        @connection = connection
       end
 
       attr_reader :cursor_id,
@@ -46,7 +48,8 @@ module Mongo
       :db_name,
       :connection_global_id,
       :server_address,
-      :session
+      :session,
+      :connection
 
       def ==(other)
         cursor_id == other.cursor_id &&

data/lib/mongo/cursor/nontailable.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Mongo
+  class Cursor
+    # This module is used by cursor-implementing classes to indicate that
+    # the only cursors they generate are non-tailable, and iterable.
+    #
+    # @api private
+    module NonTailable
+      # These views are always non-tailable.
+      #
+      # @return [ nil ] indicating a non-tailable cursor.
+      def cursor_type
+        nil
+      end
+
+      # These views apply timeouts to each iteration of a cursor, as
+      # opposed to the entire lifetime of the cursor.
+      #
+      # @return [ :iterable ] indicating a cursor with a timeout mode of
+      #   "iterable".
+      def timeout_mode
+        :iterable
+      end
+    end
+  end
+end

data/lib/mongo/cursor.rb

@@ -49,6 +49,9 @@ module Mongo
     # @api private
     attr_reader :resume_token
 
+    # @return [ Operation::Context ] context the context for this cursor
+    attr_reader :context
+
     # Creates a +Cursor+ object.
     #
     # @example Instantiate the cursor.
@@ -59,6 +62,8 @@ module Mongo
     # @param [ Server ] server The server this cursor is locked to.
     # @param [ Hash ] options The cursor options.
     #
+    # @option options [ Operation::Context ] :context The operation context
+    #   for this cursor.
     # @option options [ true, false ] :disable_retry Whether to disable
     #   retrying on error when sending getMore operations (deprecated, getMore
     #   operations are no longer retried)
@@ -80,15 +85,25 @@ module Mongo
       if @cursor_id.nil?
         raise ArgumentError, 'Cursor id must be present in the result'
       end
-      @connection_global_id = result.connection_global_id
       @options = options
       @session = @options[:session]
+      @connection_global_id = result.connection_global_id
+      @context = @options[:context]&.with(connection_global_id: connection_global_id_for_context) || fresh_context
       @explicitly_closed = false
       @lock = Mutex.new
-      unless closed?
+      if server.load_balancer?
+        # We need the connection in the cursor only in load balanced topology;
+        # we do not need an additional reference to it otherwise.
+        @connection = @initial_result.connection
+      end
+      if closed?
+        check_in_connection
+      else
         register
-        ObjectSpace.define_finalizer(self, self.class.finalize(kill_spec(@connection_global_id),
-          cluster))
+        ObjectSpace.define_finalizer(
+          self,
+          self.class.finalize(kill_spec(@connection_global_id), cluster)
+        )
       end
     end
 
@@ -98,6 +113,9 @@ module Mongo
     # @api private
     attr_reader :initial_result
 
+    # @api private
+    attr_reader :connection
+
     # Finalize the cursor for garbage collection. Schedules this cursor to be included
     # in a killCursors operation executed by the Cluster's CursorReaper.
     #
@@ -284,9 +302,11 @@ module Mongo
     # the server.
     #
     # @return [ nil ] Always nil.
-    def close
+    def close(opts = {})
       return if closed?
 
+      ctx = context ? context.refresh(timeout_ms: opts[:timeout_ms]) : fresh_context(opts)
+
       unregister
       read_with_one_retry do
         spec = {
@@ -295,11 +315,11 @@ module Mongo
           cursor_ids: [id],
         }
         op = Operation::KillCursors.new(spec)
-        execute_operation(op)
+        execute_operation(op, context: ctx)
       end
 
       nil
-    rescue Error::OperationFailure, Error::SocketError, Error::SocketTimeoutError, Error::ServerNotUsable
+    rescue Error::OperationFailure::Family, Error::SocketError, Error::SocketTimeoutError, Error::ServerNotUsable
       # Errors are swallowed since there is noting can be done by handling them.
     ensure
       end_session
@@ -307,6 +327,7 @@ module Mongo
       @lock.synchronize do
         @explicitly_closed = true
       end
+      check_in_connection
     end
 
     # Get the parsed collection name.
@@ -374,7 +395,7 @@ module Mongo
       # Legacy retryable read logic used to retry getMores, but since
       # doing so may result in silent data loss, the driver no longer retries
       # getMore operations in any circumstance.
-      # https://github.com/mongodb/specifications/blob/master/source/retryable-reads/retryable-reads.rst#qa
+      # https://github.com/mongodb/specifications/blob/master/source/retryable-reads/retryable-reads.md#qa
       process(execute_operation(get_more_operation))
     end
 
@@ -387,6 +408,7 @@ module Mongo
         connection_global_id: connection_global_id,
         server_address: server.address,
         session: @session,
+        connection: @connection
       )
     end
 
@@ -434,15 +456,7 @@ module Mongo
         # 3.2+ servers use batch_size, 3.0- servers use to_return.
         # TODO should to_return be calculated in the operation layer?
         batch_size: batch_size_for_get_more,
-        to_return: to_return,
-        max_time_ms: if view.respond_to?(:max_await_time_ms) &&
-          view.max_await_time_ms &&
-          view.options[:await_data]
-        then
-          view.max_await_time_ms
-        else
-          nil
-        end,
+        to_return: to_return
       }
       if view.respond_to?(:options) && view.options.is_a?(Hash)
         spec[:comment] = view.options[:comment] unless view.options[:comment].nil?
@@ -464,7 +478,10 @@ module Mongo
       # the @cursor_id may be zero (all results fit in the first batch).
       # Thus we need to check both @cursor_id and the cursor_id of the result
       # prior to calling unregister here.
-      unregister if !closed? && result.cursor_id == 0
+      if !closed? && result.cursor_id == 0
+        unregister
+        check_in_connection
+      end
       @cursor_id = set_cursor_id(result)
 
       if result.respond_to?(:post_batch_resume_token)
@@ -495,13 +512,22 @@ module Mongo
       cluster.unregister_cursor(@cursor_id)
     end
 
-    def execute_operation(op)
-      context = Operation::Context.new(
-        client: client,
-        session: @session,
-        connection_global_id: @connection_global_id,
-      )
-      op.execute(@server, context: context)
+    def execute_operation(op, context: nil)
+      op_context = context || possibly_refreshed_context
+      if @connection.nil?
+        op.execute(@server, context: op_context)
+      else
+        op.execute_with_connection(@connection, context: op_context)
+      end
+    end
+
+    # Considers the timeout mode and will either return the cursor's
+    # context directly, or will return a new (refreshed) context.
+    #
+    # @return [ Operation::Context ] the (possibly-refreshed) context.
+    def possibly_refreshed_context
+      return context if view.timeout_mode == :cursor_lifetime
+      context.refresh(view: view)
     end
 
     # Sets @cursor_id from the operation result.
@@ -521,6 +547,42 @@ module Mongo
                    end
     end
 
+    # Returns a newly instantiated operation context based on the
+    # default values from the view.
+    def fresh_context(opts = {})
+      Operation::Context.new(client: view.client,
+                             session: @session,
+                             connection_global_id: connection_global_id_for_context,
+                             operation_timeouts: view.operation_timeouts(opts),
+                             view: view)
+    end
+
+    # Because a context must not have a connection_global_id if the session
+    # is already pinned to one, this method checks to see whether or not there's
+    # pinned connection_global_id on the session and returns nil if so.
+    def connection_global_id_for_context
+      if @session&.pinned_connection_global_id
+        nil
+      else
+        @connection_global_id
+      end
+    end
+
+    # Returns the connection that was used to create the cursor back to the
+    # corresponding connection pool.
+    #
+    # In a load balanced topology cursors must use the same connection for the
+    # initial and all subsequent operations. Therefore, the connection is not
+    # checked into the pool after the initial operation is completed, but
+    # only when the cursor is drained.
+    def check_in_connection
+      # Connection nil means the connection has been already checked in.
+      return if @connection.nil?
+      return unless @connection.server.load_balancer?
+
+      @connection.connection_pool.check_in(@connection)
+      @connection = nil
+    end
   end
 end
 

data/lib/mongo/cursor_host.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module Mongo
+  # A shared concern implementing settings and configuration for entities that
+  # "host" (or spawn) cursors.
+  #
+  # The class or module that includes this concern must implement:
+  #   * timeout_ms -- this must return either the operation level timeout_ms
+  #       (if set) or an inherited timeout_ms from a hierarchically higher
+  #       level (if any).
+  module CursorHost
+    # Returns the cursor associated with this view, if any.
+    #
+    # @return [ nil | Cursor ] The cursor, if any.
+    #
+    # @api private
+    attr_reader :cursor
+
+    # @return [ :cursor_lifetime | :iteration ] The timeout mode to be
+    #   used by this object.
+    attr_reader :timeout_mode
+
+    # Ensure the timeout mode is appropriate for other options that
+    # have been given.
+    #
+    # @param [ Hash ] options The options to inspect.
+    # @param [ Array<Symbol> ] forbid The list of options to forbid for this
+    #   class.
+    #
+    # @raise [ ArgumentError ] if inconsistent or incompatible options are
+    #   detected.
+    #
+    # @api private
+    # rubocop:disable Metrics
+    def validate_timeout_mode!(options, forbid: [])
+      forbid.each do |key|
+        raise ArgumentError, "#{key} is not allowed here" if options.key?(key)
+      end
+
+      cursor_type = options[:cursor_type]
+      timeout_mode = options[:timeout_mode]
+
+      if timeout_ms
+        # "Tailable cursors only support the ITERATION value for the
+        # timeoutMode option. This is the default value and drivers MUST
+        # error if the option is set to CURSOR_LIFETIME."
+        if cursor_type
+          timeout_mode ||= :iteration
+          if timeout_mode == :cursor_lifetime
+            raise ArgumentError, 'tailable cursors only support `timeout_mode: :iteration`'
+          end
+
+          # "Drivers MUST error if [the maxAwaitTimeMS] option is set,
+          # timeoutMS is set to a non-zero value, and maxAwaitTimeMS is
+          # greater than or equal to timeoutMS."
+          max_await_time_ms = options[:max_await_time_ms] || 0
+          if cursor_type == :tailable_await && max_await_time_ms >= timeout_ms
+            raise ArgumentError, ':max_await_time_ms must not be >= :timeout_ms'
+          end
+        else
+          # "For non-tailable cursors, the default value of timeoutMode
+          # is CURSOR_LIFETIME."
+          timeout_mode ||= :cursor_lifetime
+        end
+      elsif timeout_mode
+        # "Drivers MUST error if timeoutMode is set and timeoutMS is not."
+        raise ArgumentError, ':timeout_ms must be set if :timeout_mode is set'
+      end
+
+      if timeout_mode == :iteration && respond_to?(:write?) && write?
+        raise ArgumentError, 'timeout_mode=:iteration is not supported for aggregation pipelines with $out or $merge'
+      end
+
+      # set it as an instance variable, rather than updating the options,
+      # because if the cursor type changes (e.g. via #configure()), the new
+      # View instance must be able to select a different default timeout_mode
+      # if no timeout_mode was set initially.
+      @timeout_mode = timeout_mode
+    end
+    # rubocop:enable Metrics
+  end
+end