mongo 2.20.1 → 2.21.1

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

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

@@ -57,7 +57,7 @@ module Mongo
       #
       # @return [ true | false ] Whether the :secondary_ok flag should be added.
       def add_secondary_ok_flag?(connection)
-        # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#topology-type-single
+        # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.md#topology-type-single
         if connection.description.standalone?
           # Read preference is never sent to standalones.
           false
@@ -83,7 +83,7 @@ module Mongo
       #
       # $readPreference is only sent when the server is a mongos,
       # following the rules described in
-      # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#passing-read-preference-to-mongos.
+      # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.md#passing-read-preference-to-mongos.
       # The topology does not matter for figuring out whether to send
       # $readPreference since the decision is always made based on
       # server type.

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

@@ -42,7 +42,7 @@ module Mongo
       # Adds error labels to exceptions raised in the yielded to block,
       # which should perform MongoDB operations and raise Mongo::Errors on
       # failure. This method handles network errors (Error::SocketError)
-      # and server-side errors (Error::OperationFailure); it does not
+      # and server-side errors (Error::OperationFailure::Family); it does not
       # handle server selection errors (Error::NoServerAvailable), for which
       # labels are added in the server selection code.
       #
@@ -65,7 +65,7 @@ module Mongo
       rescue Mongo::Error::SocketTimeoutError => e
         maybe_add_retryable_write_error_label!(e, connection, context)
         raise e
-      rescue Mongo::Error::OperationFailure => e
+      rescue Mongo::Error::OperationFailure::Family => e
         if context.committing_transaction?
           if e.write_retryable? || e.wtimeout? || (e.write_concern_error? &&
               !Session::UNLABELED_WRITE_CONCERN_CODES.include?(e.write_concern_error_code)
@@ -104,7 +104,7 @@ module Mongo
       # raised during execution of operations on servers.
       def add_server_diagnostics(connection)
         yield
-      rescue Error::SocketError, Error::SocketTimeoutError
+      rescue Error::SocketError, Error::SocketTimeoutError, Error::TimeoutError
         # Diagnostics should have already been added by the connection code,
         # do not add them again.
         raise

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

@@ -114,7 +114,7 @@ module Mongo
       end
 
       def apply_txn_opts!(selector)
-        session.add_txn_opts!(selector, read_command?(selector))
+        session.add_txn_opts!(selector, read_command?(selector), context)
       end
 
       def suppress_read_write_concern!(selector)
@@ -154,7 +154,7 @@ module Mongo
       #
       # $readPreference is only sent when the server is a mongos,
       # following the rules described in
-      # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#passing-read-preference-to-mongos.
+      # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.md#passing-read-preference-to-mongos.
       # The topology does not matter for figuring out whether to send
       # $readPreference since the decision is always made based on
       # server type.
@@ -167,7 +167,7 @@ module Mongo
       def add_read_preference(sel, connection)
         Lint.assert_type(connection, Server::Connection)
 
-        # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.rst#topology-type-single
+        # https://github.com/mongodb/specifications/blob/master/source/server-selection/server-selection.md#topology-type-single
         read_doc = if connection.description.standalone?
           # Read preference is never sent to standalones.
           nil

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

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module Mongo
+  module Operation
+    # Defines the behavior of operations that have the default timeout
+    # behavior described by the client-side operation timeouts (CSOT)
+    # spec.
+    #
+    # @api private
+    module Timed
+      # If a timeout is active (as defined by the current context), and it has
+      # not yet expired, add :maxTimeMS to the spec.
+      #
+      # @param [ Hash ] spec The spec to modify
+      # @param [ Connection ] connection The connection that will be used to
+      #   execute the operation
+      #
+      # @return [ Hash ] the spec
+      #
+      # @raises [ Mongo::Error::TimeoutError ] if the current timeout has
+      #   expired.
+      def apply_relevant_timeouts_to(spec, connection)
+        with_max_time(connection) do |max_time_sec|
+          return spec if max_time_sec.nil?
+          return spec if connection.description.mongocryptd?
+
+          spec.tap { spec[:maxTimeMS] = (max_time_sec * 1_000).to_i }
+        end
+      end
+
+      # A helper method that computes the remaining timeout (in seconds) and
+      # yields it to the associated block. If no timeout is present, yields
+      # nil. If the timeout has expired, raises Mongo::Error::TimeoutError.
+      #
+      # @param [ Connection ] connection The connection that will be used to
+      #   execute the operation
+      #
+      # @return [ Hash ] the result of yielding to the block (which must be
+      #   a Hash)
+      def with_max_time(connection)
+        if context&.timeout?
+          max_time_sec = context.remaining_timeout_sec - connection.server.minimum_round_trip_time
+          raise Mongo::Error::TimeoutError if max_time_sec <= 0
+
+          yield max_time_sec
+        else
+          yield nil
+        end
+      end
+    end
+  end
+end

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

@@ -35,7 +35,10 @@ module Mongo
       #
       # @since 2.5.2
       def execute(server, context:)
-        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)
         end
       end

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

@@ -45,7 +45,8 @@ module Mongo
         def message(connection)
           updates = validate_updates(connection, send(IDENTIFIER))
           section = Protocol::Msg::Section1.new(IDENTIFIER, updates)
-          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/update_search_index/op_msg.rb

@@ -14,11 +14,11 @@ module Mongo
         # Returns the command to send to the database, describing the
         # desired updateSearchIndex 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)
           {
             updateSearchIndex: coll_name,
             :$db => db_name,

data/lib/mongo/operation.rb

@@ -22,6 +22,7 @@ require 'mongo/operation/shared/specifiable'
 require 'mongo/operation/shared/validatable'
 require 'mongo/operation/shared/object_id_generator'
 require 'mongo/operation/shared/op_msg_executable'
+require 'mongo/operation/shared/timed'
 
 require 'mongo/operation/op_msg_base'
 require 'mongo/operation/command'

data/lib/mongo/protocol/compressed.rb

@@ -21,7 +21,7 @@ module Mongo
     # MongoDB Wire protocol Compressed message.
     #
     # This is a bi-directional message that compresses another opcode.
-    # See https://github.com/mongodb/specifications/blob/master/source/compression/OP_COMPRESSED.rst
+    # See https://github.com/mongodb/specifications/blob/master/source/compression/OP_COMPRESSED.md
     #
     # @api semipublic
     #

data/lib/mongo/protocol/message.rb

@@ -244,10 +244,7 @@ module Mongo
         # timeout option. For compatibility with whoever might call this
         # method with some other IO-like object, pass options only when they
         # are not empty.
-        read_options = {}
-        if timeout = options[:socket_timeout]
-          read_options[:timeout] = timeout
-        end
+        read_options = options.slice(:timeout, :socket_timeout)
 
         if read_options.empty?
           chunk = io.read(16)

data/lib/mongo/protocol/msg.rb

@@ -226,7 +226,7 @@ module Mongo
 
           db_name = @main_document[DATABASE_IDENTIFIER]
           cmd = merge_sections
-          enc_cmd = context.encrypter.encrypt(db_name, cmd)
+          enc_cmd = context.encrypt(db_name, cmd)
           if cmd.key?('$db') && !enc_cmd.key?('$db')
             enc_cmd['$db'] = cmd['$db']
           end
@@ -251,7 +251,7 @@ module Mongo
       def maybe_decrypt(context)
         if context.decrypt?
           cmd = merge_sections
-          enc_cmd = context.encrypter.decrypt(cmd)
+          enc_cmd = context.decrypt(cmd)
           Msg.new(@flags, @options, enc_cmd)
         else
           self
@@ -359,7 +359,7 @@ module Mongo
 
       # Encapsulates a type 1 OP_MSG section.
       #
-      # @see https://github.com/mongodb/specifications/blob/master/source/message/OP_MSG.rst#sections
+      # @see https://github.com/mongodb/specifications/blob/master/source/message/OP_MSG.md#sections
       #
       # @api private
       class Section1

data/lib/mongo/retryable/read_worker.rb

@@ -60,19 +60,21 @@ module Mongo
       # @param [ Mongo::ServerSelector::Selectable ] server_selector Server
       #   selector for the operation.
       # @param [ CollectionView ] view The +CollectionView+ defining the query.
+      # @param [ Operation::Context | nil ] context the operation context to use
+      #   with the cursor.
       # @param [ Proc ] block The block to execute.
       #
       # @return [ Cursor ] The cursor for the result set.
-      def read_with_retry_cursor(session, server_selector, view, &block)
-        read_with_retry(session, server_selector) do |server|
+      def read_with_retry_cursor(session, server_selector, view, context: nil, &block)
+        read_with_retry(session, server_selector, context) do |server|
           result = yield server
 
           # RUBY-2367: This will be updated to allow the query cache to
           # cache cursors with multi-batch results.
           if QueryCache.enabled? && !view.collection.system_collection?
-            CachingCursor.new(view, result, server, session: session)
+            CachingCursor.new(view, result, server, session: session, context: context)
           else
-            Cursor.new(view, result, server, session: session)
+            Cursor.new(view, result, server, session: session, context: context)
           end
         end
       end
@@ -107,16 +109,18 @@ module Mongo
       #   is being run on.
       # @param [ Mongo::ServerSelector::Selectable | nil ] server_selector
       #   Server selector for the operation.
+      # @param [ Mongo::Operation::Context | nil ] context Context for the
+      #   read operation.
       # @param [ Proc ] block The block to execute.
       #
       # @return [ Result ] The result of the operation.
-      def read_with_retry(session = nil, server_selector = nil, &block)
+      def read_with_retry(session = nil, server_selector = nil, context = nil, &block)
         if session.nil? && server_selector.nil?
           deprecated_legacy_read_with_retry(&block)
         elsif session&.retry_reads?
-          modern_read_with_retry(session, server_selector, &block)
+          modern_read_with_retry(session, server_selector, context, &block)
         elsif client.max_read_retries > 0
-          legacy_read_with_retry(session, server_selector, &block)
+          legacy_read_with_retry(session, server_selector, context, &block)
         else
           read_without_retry(session, server_selector, &block)
         end
@@ -186,17 +190,24 @@ module Mongo
       #   being run on.
       # @param [ Mongo::ServerSelector::Selectable ] server_selector Server
       #   selector for the operation.
+      # @param [ Mongo::Operation::Context ] context Context for the
+      #   read operation.
       # @param [ Proc ] block The block to execute.
       #
       # @return [ Result ] The result of the operation.
-      def modern_read_with_retry(session, server_selector, &block)
-        server = select_server(cluster, server_selector, session)
+      def modern_read_with_retry(session, server_selector, context, &block)
+        server = select_server(
+          cluster,
+          server_selector,
+          session,
+          timeout: context&.remaining_timeout_sec
+        )
         yield server
-      rescue *retryable_exceptions, Error::OperationFailure, Auth::Unauthorized, Error::PoolError => e
+      rescue *retryable_exceptions, Error::OperationFailure::Family, Auth::Unauthorized, Error::PoolError => e
         e.add_notes('modern retry', 'attempt 1')
         raise e if session.in_transaction?
         raise e if !is_retryable_exception?(e) && !e.write_retryable?
-        retry_read(e, session, server_selector, failed_server: server, &block)
+        retry_read(e, session, server_selector, context: context, failed_server: server, &block)
       end
 
       # Attempts to do a "legacy" read with retry. The operation will be
@@ -207,17 +218,19 @@ module Mongo
       #   being run on.
       # @param [ Mongo::ServerSelector::Selectable ] server_selector Server
       #   selector for the operation.
+      # @param [ Mongo::Operation::Context | nil ] context Context for the
+      #   read operation.
       # @param [ Proc ] block The block to execute.
       #
       # @return [ Result ] The result of the operation.
-      def legacy_read_with_retry(session, server_selector, &block)
+      def legacy_read_with_retry(session, server_selector, context = nil, &block)
+        context&.check_timeout!
         attempt = attempt ? attempt + 1 : 1
         yield select_server(cluster, server_selector, session)
-      rescue *legacy_retryable_exceptions, Error::OperationFailure => e
+      rescue *legacy_retryable_exceptions, Error::OperationFailure::Family => e
         e.add_notes('legacy retry', "attempt #{attempt}")
 
         if is_legacy_retryable_exception?(e)
-
           raise e if attempt > client.max_read_retries || session&.in_transaction?
         elsif e.retryable? && !session&.in_transaction?
           raise e if attempt > client.max_read_retries
@@ -245,7 +258,7 @@ module Mongo
 
         begin
           yield server
-        rescue *retryable_exceptions, Error::PoolError, Error::OperationFailure => e
+        rescue *retryable_exceptions, Error::PoolError, Error::OperationFailure::Family => e
           e.add_note('retries disabled')
           raise e
         end
@@ -259,40 +272,67 @@ module Mongo
       #   being run on.
       # @param [ Mongo::ServerSelector::Selectable ] server_selector Server
       #   selector for the operation.
-      # @param [ Mongo::Server ] failed_server The server on which the original
+      # @param [ Mongo::Operation::Context | nil ] :context Context for the
+      #   read operation.
+      # @param [ Mongo::Server | nil ] :failed_server The server on which the original
       #   operation failed.
       # @param [ Proc ] block The block to execute.
       #
       # @return [ Result ] The result of the operation.
-      def retry_read(original_error, session, server_selector, failed_server: nil, &block)
-        begin
-          server = select_server(cluster, server_selector, session, failed_server)
-        rescue Error, Error::AuthError => e
-          original_error.add_note("later retry failed: #{e.class}: #{e}")
-          raise original_error
-        end
+      def retry_read(original_error, session, server_selector, context: nil, failed_server: nil, &block)
+        server = select_server_for_retry(
+          original_error, session, server_selector, context, failed_server
+        )
 
         log_retry(original_error, message: 'Read retry')
 
         begin
+          context&.check_timeout!
+          attempt = attempt ? attempt + 1 : 2
           yield server, true
+        rescue Error::TimeoutError
+          raise
         rescue *retryable_exceptions => e
-          e.add_notes('modern retry', 'attempt 2')
-          raise e
-        rescue Error::OperationFailure, Error::PoolError => e
+          e.add_notes('modern retry', "attempt #{attempt}")
+          if context&.csot?
+            failed_server = server
+            retry
+          else
+            raise e
+          end
+        rescue Error::OperationFailure::Family, Error::PoolError => e
           e.add_note('modern retry')
-          unless e.write_retryable?
+          if e.write_retryable?
+            e.add_note("attempt #{attempt}")
+            if context&.csot?
+              failed_server = server
+              retry
+            else
+              raise e
+            end
+          else
             original_error.add_note("later retry failed: #{e.class}: #{e}")
             raise original_error
           end
-          e.add_note("attempt 2")
-          raise e
         rescue Error, Error::AuthError => e
           e.add_note('modern retry')
           original_error.add_note("later retry failed: #{e.class}: #{e}")
           raise original_error
         end
       end
+
+      def select_server_for_retry(original_error, session, server_selector, context, failed_server)
+        select_server(
+          cluster,
+          server_selector,
+          session,
+          failed_server,
+          timeout: context&.remaining_timeout_sec
+        )
+      rescue Error, Error::AuthError => e
+        original_error.add_note("later retry failed: #{e.class}: #{e}")
+        raise original_error
+      end
     end
   end
 end