mongo 2.20.0 → 2.21.0

data/lib/mongo/grid/stream/write.rb

@@ -83,6 +83,12 @@ module Mongo
             @options.freeze
             @filename = @options[:filename]
             @open = true
+            @timeout_holder = CsotTimeoutHolder.new(
+              operation_timeouts: {
+                operation_timeout_ms: options[:timeout_ms],
+                inherited_timeout_ms: fs.database.timeout_ms
+              }
+            )
           end
 
           # Write to the GridFS bucket from the source stream or a string.
@@ -107,7 +113,12 @@ module Mongo
             end
             chunks = File::Chunk.split(io, file_info, @n)
             @n += chunks.size
-            chunks_collection.insert_many(chunks) unless chunks.empty?
+            unless chunks.empty?
+              chunks_collection.insert_many(
+                chunks,
+                timeout_ms: @timeout_holder.remaining_timeout_ms!
+              )
+            end
             self
           end
 
@@ -124,7 +135,10 @@ module Mongo
           def close
             ensure_open!
             update_length
-            files_collection.insert_one(file_info, @options)
+            files_collection.insert_one(
+              file_info,
+              @options.merge(timeout_ms: @timeout_holder.remaining_timeout_ms!)
+            )
             @open = false
             file_id
           end
@@ -166,7 +180,10 @@ module Mongo
           #
           # @since 2.1.0
           def abort
-            fs.chunks_collection.find({ :files_id => file_id }, @options).delete_many
+            fs.chunks_collection.find(
+              { :files_id => file_id },
+              @options.merge(timeout_ms: @timeout_holder.remaining_timeout_ms!)
+            ).delete_many
             (@open = false) || true
           end
 
@@ -200,7 +217,7 @@ module Mongo
           end
 
           def ensure_indexes!
-            fs.send(:ensure_indexes!)
+            fs.send(:ensure_indexes!, @timeout_holder)
           end
 
           def ensure_open!

data/lib/mongo/index/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
   module Index
 
@@ -25,6 +27,8 @@ module Mongo
       extend Forwardable
       include Enumerable
       include Retryable
+      include Mongo::CursorHost
+      include Cursor::NonTailable
 
       # @return [ Collection ] collection The indexes collection.
       attr_reader :collection
@@ -33,6 +37,12 @@ module Mongo
       #   when sending the listIndexes command.
       attr_reader :batch_size
 
+      # @return [ Integer | nil | The timeout_ms value that was passed as an
+      #   option to the view.
+      #
+      # @api private
+      attr_reader :operation_timeout_ms
+
       def_delegators :@collection, :cluster, :database, :read_preference, :write_concern, :client
       def_delegators :cluster, :next_primary
 
@@ -90,7 +100,7 @@ module Mongo
       # @since 2.0.0
       def drop_one(name, options = {})
         raise Error::MultiIndexDrop.new if name == Index::ALL
-        drop_by_name(name, comment: options[:comment])
+        drop_by_name(name, options)
       end
 
       # Drop all indexes on the collection.
@@ -107,7 +117,7 @@ module Mongo
       #
       # @since 2.0.0
       def drop_all(options = {})
-        drop_by_name(Index::ALL, comment: options[:comment])
+        drop_by_name(Index::ALL, options)
       end
 
       # Creates an index on the collection.
@@ -161,7 +171,7 @@ module Mongo
         if session = @options[:session]
           create_options[:session] = session
         end
-        %i(commit_quorum session comment).each do |key|
+        %i(commit_quorum session comment timeout_ms max_time_ms).each do |key|
           if value = options.delete(key)
             create_options[key] = value
           end
@@ -210,7 +220,7 @@ module Mongo
           options = models.pop
         end
 
-        client.send(:with_session, @options.merge(options)) do |session|
+        client.with_session(@options.merge(options)) do |session|
           server = next_primary(nil, session)
 
           indexes = normalize_models(models, server)
@@ -229,8 +239,12 @@ module Mongo
             write_concern: write_concern,
             comment: options[:comment],
           }
-
-          Operation::CreateIndex.new(spec).execute(server, context: Operation::Context.new(client: client, session: session))
+          context = Operation::Context.new(
+            client: client,
+            session: session,
+            operation_timeouts: operation_timeouts(options)
+          )
+          Operation::CreateIndex.new(spec).execute(server, context: context)
         end
       end
 
@@ -263,9 +277,15 @@ module Mongo
       #
       # @since 2.0.0
       def each(&block)
-        session = client.send(:get_session, @options)
-        cursor = read_with_retry_cursor(session, ServerSelector.primary, self) do |server|
-          send_initial_query(server, session)
+        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)
         end
         if block_given?
           cursor.each do |doc|
@@ -283,22 +303,53 @@ module Mongo
       #
       # @param [ Collection ] collection The collection.
       # @param [ Hash ] options Options for getting a list of indexes.
-      #   Only relevant for when the listIndexes command is used with server
-      #   versions >=2.8.
       #
       # @option options [ Integer ] :batch_size The batch size for results
       #   returned from the listIndexes command.
+      # @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 collection or the database or the client.
       #
       # @since 2.0.0
       def initialize(collection, options = {})
         @collection = collection
+        @operation_timeout_ms = options.delete(:timeout_ms)
+
+        validate_timeout_mode!(options)
+
         @batch_size = options[:batch_size]
         @options = options
       end
 
+      # The timeout_ms value to use for this operation; either specified as an
+      # option to the view, or inherited from the collection.
+      #
+      # @return [ Integer | nil ] the timeout_ms for this operation
+      def timeout_ms
+        operation_timeout_ms || collection.timeout_ms
+      end
+
+      # @return [ Hash ] timeout_ms value set on the operation level (if any),
+      #   and/or timeout_ms that is set on collection/database/client 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] = collection.timeout_ms
+          end
+        end
+      end
+
       private
 
-      def drop_by_name(name, comment: nil)
+      def drop_by_name(name, opts = {})
         client.send(:with_session, @options) do |session|
           spec = {
             db_name: database.name,
@@ -307,9 +358,14 @@ module Mongo
             session: session,
             write_concern: write_concern,
           }
-          spec[:comment] = comment unless comment.nil?
+          spec[:comment] = opts[:comment] unless opts[:comment].nil?
           server = next_primary(nil, session)
-          Operation::DropIndex.new(spec).execute(server, context: Operation::Context.new(client: client, session: session))
+          context = Operation::Context.new(
+            client: client,
+            session: session,
+            operation_timeouts: operation_timeouts(opts)
+          )
+          Operation::DropIndex.new(spec).execute(server, context: context)
         end
       end
 
@@ -347,8 +403,13 @@ module Mongo
         end
       end
 
-      def send_initial_query(server, session)
-        initial_query_op(session).execute(server, context: Operation::Context.new(client: client, session: session))
+      def send_initial_query(server, session, context)
+        if server.load_balancer?
+          connection = server.pool.check_out(context: context)
+          initial_query_op(session).execute_with_connection(connection, context: context)
+        else
+          initial_query_op(session).execute(server, context: context)
+        end
       end
     end
   end

data/lib/mongo/operation/context.rb

@@ -34,8 +34,15 @@ module Mongo
     # operations.
     #
     # @api private
-    class Context
-      def initialize(client: nil, session: nil, connection_global_id: nil, options: nil)
+    class Context < CsotTimeoutHolder
+      def initialize(
+        client: nil,
+        session: nil,
+        connection_global_id: nil,
+        operation_timeouts: {},
+        view: nil,
+        options: nil
+      )
         if options
           if client
             raise ArgumentError, 'Client and options cannot both be specified'
@@ -52,14 +59,33 @@ module Mongo
 
         @client = client
         @session = session
+        @view = view
         @connection_global_id = connection_global_id
         @options = options
+        super(session: session, operation_timeouts: operation_timeouts)
       end
 
       attr_reader :client
       attr_reader :session
+      attr_reader :view
       attr_reader :options
 
+      # Returns a new Operation::Context with the deadline refreshed
+      # and relative to the current moment.
+      #
+      # @return [ Operation::Context ] the refreshed context
+      def refresh(connection_global_id: @connection_global_id, timeout_ms: nil, view: nil)
+        operation_timeouts = @operation_timeouts
+        operation_timeouts = operation_timeouts.merge(operation_timeout_ms: timeout_ms) if timeout_ms
+
+        self.class.new(client: client,
+                       session: session,
+                       connection_global_id: connection_global_id,
+                       operation_timeouts: operation_timeouts,
+                       view: view || self.view,
+                       options: options)
+      end
+
       def connection_global_id
         @connection_global_id || session&.pinned_connection_global_id
       end
@@ -122,10 +148,18 @@ module Mongo
         client&.encrypter&.encrypt? || false
       end
 
+      def encrypt(db_name, cmd)
+        encrypter.encrypt(db_name, cmd, self)
+      end
+
       def decrypt?
         !!client&.encrypter
       end
 
+      def decrypt(cmd)
+        encrypter.decrypt(cmd, self)
+      end
+
       def encrypter
         if client&.encrypter
           client.encrypter
@@ -133,6 +167,10 @@ module Mongo
           raise Error::InternalDriverError, 'Encrypter should only be accessed when encryption is to be performed'
         end
       end
+
+      def inspect
+        "#<#{self.class} connection_global_id=#{connection_global_id.inspect} deadline=#{deadline.inspect} options=#{options.inspect} operation_timeouts=#{operation_timeouts.inspect}>"
+      end
     end
   end
 end

data/lib/mongo/operation/create_search_indexes/op_msg.rb

@@ -14,11 +14,11 @@ module Mongo
         # Returns the command to send to the database, describing the
         # desired createSearchIndexes operation.
         #
-        # @param [ Mongo::Server ] _server the server that will receive the
+        # @param [ Connection ] _connection the connection that will receive the
         #   command
         #
         # @return [ Hash ] the selector
-        def selector(_server)
+        def selector(_connection)
           {
             createSearchIndexes: coll_name,
             :$db => db_name,

data/lib/mongo/operation/delete/op_msg.rb

@@ -49,7 +49,8 @@ module Mongo
 
         def message(connection)
           section = Protocol::Msg::Section1.new(IDENTIFIER, send(IDENTIFIER))
-          Protocol::Msg.new(flags, {}, command(connection), section)
+          cmd = apply_relevant_timeouts_to(command(connection), connection)
+          Protocol::Msg.new(flags, {}, cmd, section)
         end
       end
     end

data/lib/mongo/operation/drop_search_index/op_msg.rb

@@ -14,11 +14,11 @@ module Mongo
         # Returns the command to send to the database, describing the
         # desired dropSearchIndex operation.
         #
-        # @param [ Mongo::Server ] _server the server that will receive the
+        # @param [ Connection ] _connection the connection that will receive the
         #   command
         #
         # @return [ Hash ] the selector
-        def selector(_server)
+        def selector(_connection)
           {
             dropSearchIndex: coll_name,
             :$db => db_name,

data/lib/mongo/operation/find/op_msg.rb

@@ -31,6 +31,51 @@ module Mongo
 
         private
 
+        # Applies the relevant CSOT timeouts for a find command.
+        # Considers the cursor type and timeout mode and will add (or omit) a
+        # maxTimeMS field accordingly.
+        def apply_relevant_timeouts_to(spec, connection)
+          with_max_time(connection) do |max_time_sec|
+            timeout_ms = max_time_sec ? (max_time_sec * 1_000).to_i : nil
+            apply_find_timeouts_to(spec, timeout_ms) unless connection.description.mongocryptd?
+          end
+        end
+
+        def apply_find_timeouts_to(spec, timeout_ms)
+          view = context&.view
+          return spec unless view
+
+          case view.cursor_type
+          when nil # non-tailable
+            if view.timeout_mode == :cursor_lifetime
+              spec[:maxTimeMS] = timeout_ms || view.options[:max_time_ms]
+            else # timeout_mode == :iterable
+              # drivers MUST honor the timeoutMS option for the initial command
+              # but MUST NOT append a maxTimeMS field to the command sent to the
+              # server
+              if !timeout_ms && view.options[:max_time_ms]
+                spec[:maxTimeMS] = view.options[:max_time_ms]
+              end
+            end
+
+          when :tailable
+            # If timeoutMS is set, drivers...MUST NOT append a maxTimeMS field to any commands.
+            if !timeout_ms && view.options[:max_time_ms]
+              spec[:maxTimeMS] = view.options[:max_time_ms]
+            end
+
+          when :tailable_await
+            # The server supports the maxTimeMS option for the original command.
+            if timeout_ms || view.options[:max_time_ms]
+              spec[:maxTimeMS] = timeout_ms || view.options[:max_time_ms]
+            end
+          end
+
+          spec.tap do |spc|
+            spc.delete(:maxTimeMS) if spc[:maxTimeMS].nil?
+          end
+        end
+
         def selector(connection)
           # The mappings are BSON::Documents and as such store keys as
           # strings, the spec here has symbol keys.

data/lib/mongo/operation/get_more/op_msg.rb

@@ -28,6 +28,39 @@ module Mongo
         include ExecutableTransactionLabel
         include PolymorphicResult
         include CommandBuilder
+
+        private
+
+        # Applies the relevant CSOT timeouts for a getMore command.
+        # Considers the cursor type and timeout mode and will add (or omit) a
+        # maxTimeMS field accordingly.
+        def apply_relevant_timeouts_to(spec, connection)
+          with_max_time(connection) do |max_time_sec|
+            timeout_ms = max_time_sec ? (max_time_sec * 1_000).to_i : nil
+            apply_get_more_timeouts_to(spec, timeout_ms)
+          end
+        end
+
+        def apply_get_more_timeouts_to(spec, timeout_ms)
+          view = context&.view
+          return spec unless view
+
+          if view.cursor_type == :tailable_await
+            # If timeoutMS is set, drivers MUST apply it to the original operation.
+            # Drivers MUST also apply the original timeoutMS value to each next
+            # call on the resulting cursor but MUST NOT use it to derive a
+            # maxTimeMS value for getMore commands. Helpers for operations that
+            # create tailable awaitData cursors MUST also support the
+            # maxAwaitTimeMS option. Drivers MUST error if this option is set,
+            # timeoutMS is set to a non-zero value, and maxAwaitTimeMS is greater
+            # than or equal to timeoutMS. If this option is set, drivers MUST use
+            # it as the maxTimeMS field on getMore commands.
+            max_await_time_ms = view.respond_to?(:max_await_time_ms) ? view.max_await_time_ms : nil
+            spec[:maxTimeMS] = max_await_time_ms if max_await_time_ms
+          end
+
+          spec
+        end
       end
     end
   end

data/lib/mongo/operation/insert/op_msg.rb

@@ -35,7 +35,7 @@ module Mongo
 
         def get_result(connection, context, options = {})
           # This is a Mongo::Operation::Insert::Result
-          Result.new(*dispatch_message(connection, context), @ids)
+          Result.new(*dispatch_message(connection, context), @ids, context: context)
         end
 
         def selector(connection)
@@ -49,7 +49,8 @@ module Mongo
 
         def message(connection)
           section = Protocol::Msg::Section1.new(IDENTIFIER, send(IDENTIFIER))
-          Protocol::Msg.new(flags, {}, command(connection), section)
+          cmd = apply_relevant_timeouts_to(command(connection), connection)
+          Protocol::Msg.new(flags, {}, cmd, section)
         end
       end
     end

data/lib/mongo/operation/insert/result.rb

@@ -47,11 +47,13 @@ module Mongo
         #   Global id of the connection on which the operation that
         #   this result is for was performed.
         # @param [ Array<Object> ] ids The ids of the inserted documents.
+        # @param [ Operation::Context | nil ] context the operation context that
+        #   was active when this result was produced.
         #
         # @since 2.0.0
         # @api private
-        def initialize(replies, connection_description, connection_global_id, ids)
-          super(replies, connection_description, connection_global_id)
+        def initialize(replies, connection_description, connection_global_id, ids, context: nil)
+          super(replies, connection_description, connection_global_id, context: context)
           @inserted_ids = ids
         end
 

data/lib/mongo/operation/list_collections/result.rb

@@ -85,7 +85,7 @@ module Mongo
           if successful?
             self
           else
-            raise Error::OperationFailure.new(
+            raise operation_failure_class.new(
               parser.message,
               self,
               code: parser.code,

data/lib/mongo/operation/map_reduce/result.rb

@@ -108,7 +108,7 @@ module Mongo
         # @example Validate the result.
         #   result.validate!
         #
-        # @raise [ Error::OperationFailure ] If an error is in the result.
+        # @raise [ Error::OperationFailure::Family ] If an error is in the result.
         #
         # @return [ Result ] The result if verification passed.
         #

data/lib/mongo/operation/op_msg_base.rb

@@ -22,11 +22,13 @@ module Mongo
       include Specifiable
       include Executable
       include SessionsSupported
+      include Timed
 
       private
 
       def message(connection)
-        Protocol::Msg.new(flags, options(connection), command(connection))
+        cmd = apply_relevant_timeouts_to(command(connection), connection)
+        Protocol::Msg.new(flags, options(connection), cmd)
       end
     end
   end

data/lib/mongo/operation/result.rb

@@ -100,9 +100,13 @@ module Mongo
       # @param [ Integer ] connection_global_id
       #   Global id of the connection on which the operation that
       #   this result is for was performed.
+      # @param [ Operation::Context | nil ] context the context that was active
+      #   when this result was produced.
       #
       # @api private
-      def initialize(replies, connection_description = nil, connection_global_id = nil)
+      def initialize(replies, connection_description = nil, connection_global_id = nil, context: nil, connection: nil)
+        @context = context
+
         if replies
           if replies.is_a?(Array)
             if replies.length != 1
@@ -118,6 +122,7 @@ module Mongo
           @replies = [ reply ]
           @connection_description = connection_description
           @connection_global_id = connection_global_id
+          @connection = connection
         end
       end
 
@@ -138,6 +143,14 @@ module Mongo
       # @api private
       attr_reader :connection_global_id
 
+      # @return [ Operation::Context | nil ] the operation context (if any)
+      #   that was active when this result was produced.
+      #
+      # @api private
+      attr_reader :context
+
+      attr_reader :connection
+
       # @api private
       def_delegators :parser,
         :not_master?, :node_recovering?, :node_shutting_down?
@@ -320,7 +333,7 @@ module Mongo
       # @example Validate the result.
       #   result.validate!
       #
-      # @raise [ Error::OperationFailure ] If an error is in the result.
+      # @raise [ Error::OperationFailure::Family ] If an error is in the result.
       #
       # @return [ Result ] The result if verification passed.
       #
@@ -330,16 +343,16 @@ module Mongo
         !successful? ? raise_operation_failure : self
       end
 
-      # The exception instance (of the Error::OperationFailure class)
+      # The exception instance (of Error::OperationFailure::Family)
       # that would be raised during processing of this result.
       #
       # This method should only be called when result is not successful.
       #
-      # @return [ Error::OperationFailure ] The exception.
+      # @return [ Error::OperationFailure::Family ] The exception.
       #
       # @api private
       def error
-        @error ||= Error::OperationFailure.new(
+        @error ||= operation_failure_class.new(
           parser.message,
           self,
           code: parser.code,
@@ -453,6 +466,14 @@ module Mongo
 
       private
 
+      def operation_failure_class
+        if context&.csot? && parser.code == 50
+          Error::ServerTimeoutError
+        else
+          Error::OperationFailure
+        end
+      end
+
       def aggregate_returned_count
         replies.reduce(0) do |n, reply|
           n += reply.number_returned

data/lib/mongo/operation/shared/executable.rb

@@ -28,7 +28,18 @@ module Mongo
 
       include ResponseHandling
 
+      # @return [ Operation::Context | nil ] the operation context used to
+      #   execute this operation.
+      attr_accessor :context
+
       def do_execute(connection, context, options = {})
+        # Save the context on the instance, to avoid having to pass it as a
+        # parameter to every single method. There are many legacy methods that
+        # still accept it as a parameter, which are left as-is for now to
+        # minimize the impact of this change. Moving forward, it may be
+        # reasonable to refactor things so this saved reference is used instead.
+        @context = context
+
         session&.materialize_if_needed
         unpin_maybe(session, connection) do
           add_error_labels(connection, context) do
@@ -93,7 +104,7 @@ module Mongo
       end
 
       def get_result(connection, context, options = {})
-        result_class.new(*dispatch_message(connection, context, options))
+        result_class.new(*dispatch_message(connection, context, options), context: context, connection: connection)
       end
 
       # Returns a Protocol::Message or nil as reply.

data/lib/mongo/operation/shared/op_msg_executable.rb

@@ -32,7 +32,10 @@ module Mongo
       #
       # @return [ Mongo::Operation::Result ] The operation result.
       def execute(server, context:, options: {})
-        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+        server.with_connection(
+          connection_global_id: context.connection_global_id,
+          context: context
+        ) do |connection|
           execute_with_connection(connection, context: context, options: options)
         end
       end