mongo 2.20.1 → 2.21.0

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)

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

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

data/lib/mongo/retryable/write_worker.rb

@@ -74,7 +74,11 @@ module Mongo
         # If we are here, session is not nil. A session being nil would have
         # failed retry_write_allowed? check.
 
-        server = select_server(cluster, ServerSelector.primary, session)
+        server = select_server(
+          cluster, ServerSelector.primary,
+          session,
+          timeout: context.remaining_timeout_sec
+        )
 
         unless ending_transaction || server.retry_writes?
           return legacy_write_with_retry(server, context: context, &block)
@@ -110,7 +114,7 @@ module Mongo
             server.with_connection(connection_global_id: context.connection_global_id) do |connection|
               yield connection, nil, context
             end
-          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
@@ -170,6 +174,7 @@ module Mongo
       # @api private
       def legacy_write_with_retry(server = nil, context:)
         session = context.session
+        context.check_timeout!
 
         # This is the pre-session retry logic, and is not subject to
         # current retryable write specifications.
@@ -177,12 +182,20 @@ module Mongo
         attempt = 0
         begin
           attempt += 1
-          server ||= select_server(cluster, ServerSelector.primary, session)
-          server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          server ||= select_server(
+            cluster,
+            ServerSelector.primary,
+            session,
+            timeout: context.remaining_timeout_sec
+          )
+          server.with_connection(
+            connection_global_id: context.connection_global_id,
+            context: context
+          ) do |connection|
             # Legacy retries do not use txn_num
             yield connection, nil, context.dup
           end
-        rescue Error::OperationFailure => e
+        rescue Error::OperationFailure::Family => e
           e.add_note('legacy retry')
           e.add_note("attempt #{attempt}")
           server = nil
@@ -220,7 +233,10 @@ module Mongo
         txn_num = nil
         connection_succeeded = false
 
-        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|
           connection_succeeded = true
 
           session.materialize_if_needed
@@ -230,10 +246,10 @@ module Mongo
           # it later for the retry as well.
           yield connection, txn_num, context.dup
         end
-      rescue *retryable_exceptions, Error::PoolError, Auth::Unauthorized, Error::OperationFailure => e
+      rescue *retryable_exceptions, Error::PoolError, Auth::Unauthorized, Error::OperationFailure::Family => e
         e.add_notes('modern retry', 'attempt 1')
 
-        if e.is_a?(Error::OperationFailure)
+        if e.is_a?(Error::OperationFailure::Family)
           ensure_retryable!(e)
         else
           ensure_labeled_retryable!(e, connection_succeeded, session)
@@ -256,6 +272,8 @@ module Mongo
       #
       # @return [ Result ] The result of the operation.
       def retry_write(original_error, txn_num, context:, failed_server: nil, &block)
+        context&.check_timeout!
+
         session = context.session
 
         # We do not request a scan of the cluster here, because error handling
@@ -263,7 +281,13 @@ module Mongo
         # server description and/or topology as necessary (specifically,
         # a socket error or a not master error should have marked the respective
         # server unknown). Here we just need to wait for server selection.
-        server = select_server(cluster, ServerSelector.primary, session, failed_server)
+        server = select_server(
+          cluster,
+          ServerSelector.primary,
+          session,
+          failed_server,
+          timeout: context.remaining_timeout_sec
+        )
 
         unless server.retry_writes?
           # Do not need to add "modern retry" here, it should already be on
@@ -279,14 +303,21 @@ module Mongo
           raise Error::RaiseOriginalError
         end
 
+        attempt = attempt ? attempt + 1 : 2
         log_retry(original_error, message: 'Write retry')
         server.with_connection(connection_global_id: context.connection_global_id) do |connection|
           yield(connection, txn_num, context)
         end
       rescue *retryable_exceptions, Error::PoolError => e
-        fail_on_retryable!(e, original_error)
-      rescue Error::OperationFailure => e
-        fail_on_operation_failure!(e, original_error)
+        maybe_fail_on_retryable(e, original_error, context, attempt)
+        failed_server = server
+        retry
+      rescue Error::OperationFailure::Family => e
+        maybe_fail_on_operation_failure(e, original_error, context, attempt)
+        failed_server = server
+        retry
+      rescue Mongo::Error::TimeoutError
+        raise
       rescue Error, Error::AuthError => e
         fail_on_other_error!(e, original_error)
       rescue Error::RaiseOriginalError
@@ -332,10 +363,10 @@ module Mongo
 
       # Raise either e, or original_error, depending on whether e is
       # write_retryable.
-      def fail_on_retryable!(e, original_error)
+      def maybe_fail_on_retryable(e, original_error, context, attempt)
         if e.write_retryable?
-          e.add_notes('modern retry', 'attempt 2')
-          raise e
+          e.add_notes('modern retry', "attempt #{attempt}")
+          raise e unless context&.deadline
         else
           original_error.add_note("later retry failed: #{e.class}: #{e}")
           raise original_error
@@ -344,11 +375,11 @@ module Mongo
 
       # Raise either e, or original_error, depending on whether e is
       # appropriately labeled.
-      def fail_on_operation_failure!(e, original_error)
+      def maybe_fail_on_operation_failure(e, original_error, context, attempt)
         e.add_note('modern retry')
         if e.label?('RetryableWriteError') && !e.label?('NoWritesPerformed')
-          e.add_note('attempt 2')
-          raise e
+          e.add_note("attempt #{attempt}")
+          raise e unless context&.deadline
         else
           original_error.add_note("later retry failed: #{e.class}: #{e}")
           raise original_error

data/lib/mongo/retryable.rb

@@ -46,8 +46,14 @@ module Mongo
     # @api private
     #
     # @return [ Mongo::Server ] A server matching the server preference.
-    def select_server(cluster, server_selector, session, failed_server = nil)
-      server_selector.select_server(cluster, nil, session, deprioritized: [failed_server].compact)
+    def select_server(cluster, server_selector, session, failed_server = nil, timeout: nil)
+      server_selector.select_server(
+        cluster,
+        nil,
+        session,
+        deprioritized: [failed_server].compact,
+        timeout: timeout
+      )
     end
 
     # Returns the read worker for handling retryable reads.

data/lib/mongo/server/connection.rb

@@ -226,11 +226,11 @@ module Mongo
       # @return [ true ] If the connection succeeded.
       #
       # @since 2.0.0
-      def connect!
+      def connect!(context = nil)
         raise_if_closed!
 
         unless @socket
-          @socket = create_socket
+          @socket = create_socket(context)
           @description, @compressor = do_connect
 
           if server.load_balancer?
@@ -256,10 +256,16 @@ module Mongo
       #
       #
       # @return [ Socket ] The created socket.
-      private def create_socket
+      private def create_socket(context = nil)
         add_server_diagnostics do
-          address.socket(socket_timeout, ssl_options.merge(
-            connection_address: address, connection_generation: generation, pipe: options[:pipe]))
+          opts = ssl_options.merge(
+            connection_address: address,
+            connection_generation: generation,
+            pipe: options[:pipe],
+            connect_timeout: context&.remaining_timeout_sec,
+            csot: !!context&.csot?
+          )
+          address.socket(socket_timeout, opts)
         end
       end
 

data/lib/mongo/server/connection_base.rb

@@ -169,6 +169,7 @@ module Mongo
           raise Error::LintError, "Trying to deliver a message over a disconnected connection (to #{address})"
         end
         buffer = serialize(message, context)
+        check_timeout!(context)
         ensure_connected do |socket|
           operation_id = Monitoring.next_operation_id
           started_event = command_started(address, operation_id, message.payload,
@@ -181,9 +182,10 @@ module Mongo
           result = nil
           begin
             result = add_server_diagnostics do
-              socket.write(buffer.to_s)
+              socket.write(buffer.to_s, timeout: context.remaining_timeout_sec)
               if message.replyable?
-                Protocol::Message.deserialize(socket, max_message_size, message.request_id, options)
+                check_timeout!(context)
+                Protocol::Message.deserialize(socket, max_message_size, message.request_id, options.merge(timeout: context.remaining_timeout_sec))
               else
                 nil
               end
@@ -273,6 +275,24 @@ module Mongo
 
         buffer
       end
+
+      # If timeoutMS is set for the operation context, checks whether there is
+      # enough time left to send the corresponding message to the server
+      # (remaining timeout is bigger than minimum round trip time for
+      # the server)
+      #
+      # @param [ Mongo::Operation::Context ] context Context of the operation.
+      #
+      # @raise [ Mongo::Error::TimeoutError ] if timeout expired or there is
+      #   not enough time to send the message to the server.
+      def check_timeout!(context)
+        return if [nil, 0].include?(context.deadline)
+
+        time_to_execute = context.remaining_timeout_sec - server.minimum_round_trip_time
+        if time_to_execute <= 0
+          raise Mongo::Error::TimeoutError
+        end
+      end
     end
   end
 end