mongo 2.20.0 → 2.21.0

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

data/lib/mongo/database/view.rb

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/cursor/nontailable'
+
 module Mongo
   class Database
 
@@ -25,6 +27,8 @@ module Mongo
       extend Forwardable
       include Enumerable
       include Retryable
+      include Mongo::CursorHost
+      include Cursor::NonTailable
 
       def_delegators :@database, :cluster, :read_preference, :client
       # @api private
@@ -56,6 +60,10 @@ module Mongo
       #   to run the command when access control is enforced.
       # @option options [ Object ] :comment A user-provided
       #   comment to attach to this command.
+      # @option options [ Integer ] :timeout_ms The operation timeout in milliseconds.
+      #    Must be a non-negative integer. An explicit value of 0 means infinite.
+      #    The default value is unset which means the value is inherited from
+      #    the database or the client.
       #
       #   See https://mongodb.com/docs/manual/reference/command/listCollections/
       #   for more information and usage.
@@ -66,9 +74,14 @@ module Mongo
       # @since 2.0.0
       def collection_names(options = {})
         @batch_size = options[:batch_size]
-        session = client.send(:get_session, options)
-        cursor = read_with_retry_cursor(session, ServerSelector.primary, self) do |server|
-          send_initial_query(server, session, options.merge(name_only: true))
+        session = client.get_session(options)
+        context = Operation::Context.new(
+          client: client,
+          session: session,
+          operation_timeouts: operation_timeouts(options)
+        )
+        cursor = read_with_retry_cursor(session, ServerSelector.primary, self, context: context) do |server|
+          send_initial_query(server, session, context, options.merge(name_only: true))
         end
         cursor.map do |info|
           if cursor.initial_result.connection_description.features.list_collections_enabled?
@@ -112,20 +125,33 @@ module Mongo
       #
       # @since 2.0.5
       def list_collections(options = {})
-        session = client.send(:get_session, options)
+        session = client.get_session(options)
         collections_info(session, ServerSelector.primary, options)
       end
 
       # Create the new database view.
       #
       # @example Create the new database view.
-      #   View::Index.new(database)
+      #   Database::View.new(database)
       #
       # @param [ Database ] database The database.
+      # @param [ Hash ] options The options to configure the view with.
+      #
+      # @option options [ :cursor_lifetime | :iteration ] :timeout_mode How to interpret
+      #   :timeout_ms (whether it applies to the lifetime of the cursor, or per
+      #   iteration).
+      # @option options [ Integer ] :timeout_ms The operation timeout in milliseconds.
+      #    Must be a non-negative integer. An explicit value of 0 means infinite.
+      #    The default value is unset which means the value is inherited from
+      #    the database or the client.
       #
       # @since 2.0.0
-      def initialize(database)
+      def initialize(database, options = {})
         @database = database
+        @operation_timeout_ms = options.delete(:timeout_ms)
+
+        validate_timeout_mode!(options)
+
         @batch_size =  nil
         @limit = nil
         @collection = @database[Database::COMMAND]
@@ -134,6 +160,12 @@ module Mongo
       # @api private
       attr_reader :database
 
+      # @return [ Integer | nil | The timeout_ms value that was passed as an
+      #   option to the view.
+      #
+      # @api private
+      attr_reader :operation_timeout_ms
+
       # Execute an aggregation on the database view.
       #
       # @example Aggregate documents.
@@ -152,15 +184,41 @@ module Mongo
         Collection::View::Aggregation.new(self, pipeline, options)
       end
 
+      # The timeout_ms value to use for this operation; either specified as an
+      # option to the view, or inherited from the database.
+      #
+      # @return [ Integer | nil ] the timeout_ms for this operation
+      def timeout_ms
+        operation_timeout_ms || database.timeout_ms
+      end
+
+      # @return [ Hash ] timeout_ms value set on the operation level (if any).
+      #
+      # @api private
+      def operation_timeouts(opts = {})
+        {}.tap do |result|
+          if opts[:timeout_ms] || operation_timeout_ms
+            result[:operation_timeout_ms] = opts.delete(:timeout_ms) || operation_timeout_ms
+          else
+            result[:inherited_timeout_ms] = database.timeout_ms
+          end
+        end
+      end
+
       private
 
       def collections_info(session, server_selector, options = {}, &block)
         description = nil
-        cursor = read_with_retry_cursor(session, server_selector, self) do |server|
+        context = Operation::Context.new(
+          client: client,
+          session: session,
+          operation_timeouts: operation_timeouts(options)
+        )
+        cursor = read_with_retry_cursor(session, server_selector, self, context: context) do |server|
           # TODO take description from the connection used to send the query
           # once https://jira.mongodb.org/browse/RUBY-1601 is fixed.
           description = server.description
-          send_initial_query(server, session, options)
+          send_initial_query(server, session, context, options)
         end
         # On 3.0+ servers, we get just the collection names.
         # On 2.6 server, we get collection names prefixed with the database
@@ -224,17 +282,26 @@ module Mongo
       #   types (where possible).
       #
       # @return [ Operation::Result ] Result of the query.
-      def send_initial_query(server, session, options = {})
+      def send_initial_query(server, session, context, options = {})
         opts = options.dup
         execution_opts = {}
         if opts.key?(:deserialize_as_bson)
           execution_opts[:deserialize_as_bson] = opts.delete(:deserialize_as_bson)
         end
-        initial_query_op(session, opts).execute(
-          server,
-          context: Operation::Context.new(client: client, session: session),
-          options: execution_opts
-        )
+        if server.load_balancer?
+          connection = server.pool.check_out(context: context)
+          initial_query_op(session, opts).execute_with_connection(
+            connection,
+            context: context,
+            options: execution_opts
+          )
+        else
+          initial_query_op(session, opts).execute(
+            server,
+            context: context,
+            options: execution_opts
+          )
+        end
       end
     end
   end