Sync latest CSOT spec tests (#3001)

Author: comandeo-mongo

lib/mongo/client.rb

@@ -36,6 +36,7 @@ class Client
       :write, :write_concern,
       :retry_reads, :max_read_retries, :read_retry_interval,
       :retry_writes, :max_write_retries,
+      :timeout_ms,
 
       # Options which cannot currently be here:
       #

lib/mongo/collection/view/change_stream.rb

@@ -58,6 +58,16 @@ class ChangeStream
         # @api private
         attr_reader :cursor
 
+        # Refreshes the CSOT timeout for the next iteration. Delegates to the
+        # underlying cursor's refresh_timeout! method so that each call to
+        # try_next starts with a fresh timeout deadline, as required by the
+        # CSOT spec for tailable awaitData cursors.
+        #
+        # @api private
+        def refresh_timeout!
+          @cursor&.refresh_timeout!
+        end
+
         # Initialize the change stream for the provided collection view, pipeline
         # and options.
         #

lib/mongo/csot_timeout_holder.rb

@@ -101,9 +101,10 @@ def calculate_deadline(opts = {}, session = nil)
     end
 
     def check_no_override_inside_transaction!(opts, session)
-      return unless opts[:operation_timeout_ms] && session&.with_transaction_deadline
+      return unless opts[:operation_timeout_ms] && session&.inside_with_transaction?
 
-      raise ArgumentError, 'Cannot override timeout_ms inside with_transaction block'
+      raise Mongo::Error::InvalidTransactionOperation,
+            'timeoutMS cannot be overridden inside a withTransaction callback'
     end
 
     def calculate_deadline_from_timeout_ms(operation_timeout_ms)

lib/mongo/cursor.rb

@@ -415,6 +415,18 @@ def fully_iterated?
       !!@fully_iterated
     end
 
+    # Refreshes the cursor's CSOT context so that the next getMore starts
+    # with a fresh timeout deadline. Used by tailable awaitData cursors to
+    # implement per-iteration timeout refresh as required by the CSOT spec.
+    # Only refreshes if this is a tailable awaitData cursor with an active timeout.
+    #
+    # @api private
+    def refresh_timeout!
+      return unless view.cursor_type == :tailable_await && context.timeout?
+
+      @context = @context.refresh(view: view)
+    end
+
     private
 
     def explicitly_closed?
@@ -522,6 +534,10 @@ def execute_operation(op, context: nil)
     # @return [ Operation::Context ] the (possibly-refreshed) context.
     def possibly_refreshed_context
       return context if view.timeout_mode == :cursor_lifetime
+      # For tailable await cursors with CSOT, the timeout budget is shared across
+      # all getMore commands (cumulative). Only update the view reference; keep
+      # the same deadline so remaining time decreases across getMores.
+      return context.with(view: view) if view.cursor_type == :tailable_await
       context.refresh(view: view)
     end
 

lib/mongo/cursor_host.rb

@@ -49,19 +49,19 @@ def validate_timeout_mode!(options, forbid: [])
           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
+
+        # "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 max_await_time_ms.positive? && max_await_time_ms >= timeout_ms
+          raise ArgumentError, ':max_await_time_ms must not be >= :timeout_ms'
+        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'

lib/mongo/operation/get_more/op_msg.rb

@@ -54,9 +54,17 @@ def apply_get_more_timeouts_to(spec, timeout_ms)
             # 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.
+            # it as the maxTimeMS field on getMore commands, capped at remaining
+            # CSOT timeout if less than maxAwaitTimeMS.
             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
+            if max_await_time_ms
+              effective_ms = if timeout_ms && timeout_ms < max_await_time_ms
+                               timeout_ms
+                             else
+                               max_await_time_ms
+                             end
+              spec[:maxTimeMS] = effective_ms
+            end
           end
 
           spec

lib/mongo/server/connection_pool.rb

@@ -1308,6 +1308,12 @@ def retrieve_and_connect_connection(connection_global_id, context =  nil)
         end
 
         connection
+      rescue Error::ConnectionCheckOutTimeout
+        # Per the CSOT spec, if a connection checkout fails because the CSOT
+        # deadline expired (rather than a configured waitQueueTimeout), the
+        # error must be a CSOT TimeoutError so that callers can distinguish it.
+        context&.check_timeout!
+        raise
       end
 
       # Waits for a connection to become available, or raises is no connection

lib/mongo/session.rb

@@ -106,6 +106,8 @@ def initialize(server_session, client, options = {})
       @cluster_time = nil
       @state = NO_TRANSACTION_STATE
       @with_transaction_deadline = nil
+      @with_transaction_timeout_ms = nil
+      @inside_with_transaction = false
     end
 
     # @return [ Hash ] The options for this session.
@@ -452,6 +454,8 @@ def end_session
     #
     # @since 2.7.0
     def with_transaction(options = nil)
+      @inside_with_transaction = true
+      @with_transaction_timeout_ms = options&.dig(:timeout_ms) || @options[:default_timeout_ms] || @client.timeout_ms
       @with_transaction_deadline = calculate_with_transaction_deadline(options)
       deadline = if @with_transaction_deadline
                    # CSOT enabled, so we have a customer defined deadline.
@@ -471,7 +475,7 @@ def with_transaction(options = nil)
           if overload_encountered
             delay = @client.retry_policy.backoff_delay(overload_error_count)
             if backoff_would_exceed_deadline?(deadline, delay)
-              raise(last_error)
+              raise Mongo::Error::TimeoutError, 'CSOT timeout expired waiting to retry withTransaction'
             end
             unless @client.retry_policy.should_retry_overload?(overload_error_count, delay)
               raise(last_error)
@@ -480,7 +484,7 @@ def with_transaction(options = nil)
           else
             backoff = backoff_seconds_for_retry(transaction_attempt)
             if backoff_would_exceed_deadline?(deadline, backoff)
-              raise(last_error)
+              raise Mongo::Error::TimeoutError, 'CSOT timeout expired waiting to retry withTransaction'
             end
             sleep(backoff)
           end
@@ -499,7 +503,10 @@ def with_transaction(options = nil)
         rescue Exception => e
           if within_states?(STARTING_TRANSACTION_STATE, TRANSACTION_IN_PROGRESS_STATE)
             log_warn("Aborting transaction due to #{e.class}: #{e}")
-            @with_transaction_deadline = nil
+            # CSOT: if the deadline is already expired, clear it so that
+            # abort_transaction uses a fresh timeout (not the expired deadline).
+            # If the deadline is not yet expired, keep it so abort uses remaining time.
+            @with_transaction_deadline = nil if @with_transaction_deadline && deadline_expired?(deadline)
             abort_transaction
             transaction_in_progress = false
           end
@@ -528,6 +535,15 @@ def with_transaction(options = nil)
             return rv
           end
 
+          # CSOT: if the timeout has expired before we can commit, abort the
+          # transaction instead and raise a client-side timeout error.
+          if @with_transaction_deadline && deadline_expired?(deadline)
+            transaction_in_progress = false
+            @with_transaction_deadline = nil
+            abort_transaction
+            raise Mongo::Error::TimeoutError, 'CSOT timeout expired before transaction could be committed'
+          end
+
           begin
             commit_transaction(commit_options)
             transaction_in_progress = false
@@ -610,6 +626,8 @@ def with_transaction(options = nil)
         end
       end
       @with_transaction_deadline = nil
+      @with_transaction_timeout_ms = nil
+      @inside_with_transaction = false
     end
 
     # Places subsequent operations in this session into a new transaction.
@@ -1282,6 +1300,12 @@ def txn_num
     # @api private
     attr_reader :with_transaction_deadline
 
+    # @return [ Boolean ] Whether we are currently inside a with_transaction block.
+    # @api private
+    def inside_with_transaction?
+      @inside_with_transaction
+    end
+
     private
 
     # Get the read concern the session will use when starting a transaction.
@@ -1355,9 +1379,14 @@ def check_transactions_supported!
 
     def operation_timeouts(opts)
       {
-        inherited_timeout_ms: @client.timeout_ms
+        inherited_timeout_ms: @with_transaction_timeout_ms || @client.timeout_ms
       }.tap do |result|
-        if @with_transaction_deadline.nil?
+        if @inside_with_transaction
+          if opts[:timeout_ms]
+            raise Mongo::Error::InvalidTransactionOperation,
+              'timeoutMS cannot be overridden inside a withTransaction callback'
+          end
+        else
           if timeout_ms = opts[:timeout_ms]
             result[:operation_timeout_ms] = timeout_ms
           elsif default_timeout_ms = options[:default_timeout_ms]

spec/mongo/session_transaction_spec.rb

@@ -141,7 +141,7 @@ class SessionTransactionSpecError < StandardError; end
             exc.add_label('TransientTransactionError')
             raise exc
           end
-        end.to raise_error(Mongo::Error::OperationFailure, 'timeout test')
+        end.to raise_error(Mongo::Error::TimeoutError)
       end
     end
 

spec/runners/unified/change_stream_operations.rb

@@ -23,12 +23,23 @@ def create_change_stream(op)
     def iterate_until_document_or_error(op)
       object_id = op.use!('object')
       object = entities.get_any(object_id)
-      object.try_next
+      # Per CSOT spec, timeoutMS is refreshed for each "next call" on a
+      # tailable awaitData cursor. Refresh at the start of this iteration
+      # so that each separate iterateUntilDocumentOrError operation gets a
+      # fresh deadline, while getMores within the loop share it cumulatively.
+      object.refresh_timeout! if object.respond_to?(:refresh_timeout!)
+      loop do
+        doc = object.try_next
+        return doc if doc
+      end
     end
 
     def iterate_once(op)
       stream_id = op.use!('object')
       stream = entities.get_any(stream_id)
+      # Per CSOT spec, timeoutMS is refreshed for each "next call" on a
+      # tailable awaitData cursor or change stream.
+      stream.refresh_timeout! if stream.respond_to?(:refresh_timeout!)
       stream.try_next
     end