mongo 2.20.1 → 2.21.0

data/lib/mongo/collection/view/aggregation/behavior.rb

@@ -0,0 +1,131 @@
+# frozen_string_literal: true
+
+module Mongo
+  class Collection
+    class View
+      class Aggregation
+        # Distills the behavior common to aggregator classes, like
+        # View::Aggregator and View::ChangeStream.
+        module Behavior
+          extend Forwardable
+          include Enumerable
+          include Immutable
+          include Iterable
+          include Explainable
+          include Loggable
+          include Retryable
+
+          # @return [ View ] view The collection view.
+          attr_reader :view
+
+          # Delegate necessary operations to the view.
+          def_delegators :view, :collection, :read, :cluster, :cursor_type, :limit, :batch_size
+
+          # Delegate necessary operations to the collection.
+          def_delegators :collection, :database, :client
+
+          # Set to true if disk usage is allowed during the aggregation.
+          #
+          # @example Set disk usage flag.
+          #   aggregation.allow_disk_use(true)
+          #
+          # @param [ true, false ] value The flag value.
+          #
+          # @return [ true, false, Aggregation ] The aggregation if a value was
+          #   set or the value if used as a getter.
+          #
+          # @since 2.0.0
+          def allow_disk_use(value = nil)
+            configure(:allow_disk_use, value)
+          end
+
+          # Get the explain plan for the aggregation.
+          #
+          # @example Get the explain plan for the aggregation.
+          #   aggregation.explain
+          #
+          # @return [ Hash ] The explain plan.
+          #
+          # @since 2.0.0
+          def explain
+            self.class.new(view, pipeline, options.merge(explain: true)).first
+          end
+
+          # Whether this aggregation will write its result to a database collection.
+          #
+          # @return [ Boolean ] Whether the aggregation will write its result
+          #   to a collection.
+          #
+          # @api private
+          def write?
+            pipeline.any? { |op| op.key?('$out') || op.key?(:$out) || op.key?('$merge') || op.key?(:$merge) }
+          end
+
+          # @return [ Integer | nil ] the timeout_ms value that was passed as
+          #   an option to this object, or which was inherited from the view.
+          #
+          # @api private
+          def timeout_ms
+            @timeout_ms || view.timeout_ms
+          end
+
+          private
+
+          # Common setup for all classes that include this behavior; the
+          # constructor should invoke this method.
+          def perform_setup(view, options, forbid: [])
+            @view = view
+
+            @timeout_ms = options.delete(:timeout_ms)
+            @options = BSON::Document.new(options).freeze
+
+            yield
+
+            validate_timeout_mode!(options, forbid: forbid)
+          end
+
+          def server_selector
+            @view.send(:server_selector)
+          end
+
+          def aggregate_spec(session, read_preference)
+            Builder::Aggregation.new(
+              pipeline,
+              view,
+              options.merge(session: session, read_preference: read_preference)
+            ).specification
+          end
+
+          # Skip, sort, limit, projection are specified as pipeline stages
+          # rather than as options.
+          def cache_options
+            {
+              namespace: collection.namespace,
+              selector: pipeline,
+              read_concern: view.read_concern,
+              read_preference: view.read_preference,
+              collation: options[:collation],
+              # Aggregations can read documents from more than one collection,
+              # so they will be cleared on every write operation.
+              multi_collection: true,
+            }
+          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] || @timeout_ms
+                result[:operation_timeout_ms] = opts.delete(:timeout_ms) || @timeout_ms
+              else
+                result[:inherited_timeout_ms] = view.timeout_ms
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mongo/collection/view/aggregation.rb

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/collection/view/aggregation/behavior'
+
 module Mongo
   class Collection
     class View
@@ -23,46 +25,11 @@ module Mongo
       #
       # @since 2.0.0
       class Aggregation
-        extend Forwardable
-        include Enumerable
-        include Immutable
-        include Iterable
-        include Explainable
-        include Loggable
-        include Retryable
+        include Behavior
 
-        # @return [ View ] view The collection view.
-        attr_reader :view
         # @return [ Array<Hash> ] pipeline The aggregation pipeline.
         attr_reader :pipeline
 
-        # Delegate necessary operations to the view.
-        def_delegators :view, :collection, :read, :cluster
-
-        # Delegate necessary operations to the collection.
-        def_delegators :collection, :database, :client
-
-        # The reroute message.
-        #
-        # @since 2.1.0
-        # @deprecated
-        REROUTE = 'Rerouting the Aggregation operation to the primary server.'.freeze
-
-        # Set to true if disk usage is allowed during the aggregation.
-        #
-        # @example Set disk usage flag.
-        #   aggregation.allow_disk_use(true)
-        #
-        # @param [ true, false ] value The flag value.
-        #
-        # @return [ true, false, Aggregation ] The aggregation if a value was
-        #   set or the value if used as a getter.
-        #
-        # @since 2.0.0
-        def allow_disk_use(value = nil)
-          configure(:allow_disk_use, value)
-        end
-
         # Initialize the aggregation for the provided collection view, pipeline
         # and options.
         #
@@ -86,59 +53,29 @@ module Mongo
         # @option options [ Hash ] :let Mapping of variables to use in the pipeline.
         #   See the server documentation for details.
         # @option options [ Integer ] :max_time_ms The maximum amount of time in
-        #   milliseconds to allow the aggregation to run.
-        # @option options [ true, false ] :use_cursor Indicates whether the command
-        #   will request that the server provide results using a cursor. Note that
-        #   as of server version 3.6, aggregations always provide results using a
-        #   cursor and this option is therefore not valid.
+        #   milliseconds to allow the aggregation to run. This option is deprecated, use
+        #   :timeout_ms instead.
         # @option options [ Session ] :session The session to use.
+        # @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(view, pipeline, options = {})
-          @view = view
-          @pipeline = pipeline.dup
-          unless Mongo.broken_view_aggregate || view.filter.empty?
-            @pipeline.unshift(:$match => view.filter)
+          perform_setup(view, options) do
+            @pipeline = pipeline.dup
+            unless Mongo.broken_view_aggregate || view.filter.empty?
+              @pipeline.unshift(:$match => view.filter)
+            end
           end
-          @options = BSON::Document.new(options).freeze
-        end
-
-        # Get the explain plan for the aggregation.
-        #
-        # @example Get the explain plan for the aggregation.
-        #   aggregation.explain
-        #
-        # @return [ Hash ] The explain plan.
-        #
-        # @since 2.0.0
-        def explain
-          self.class.new(view, pipeline, options.merge(explain: true)).first
-        end
-
-        # Whether this aggregation will write its result to a database collection.
-        #
-        # @return [ Boolean ] Whether the aggregation will write its result
-        #   to a collection.
-        #
-        # @api private
-        def write?
-          pipeline.any? { |op| op.key?('$out') || op.key?(:$out) || op.key?('$merge') || op.key?(:$merge) }
         end
 
         private
 
-        def server_selector
-          @view.send(:server_selector)
-        end
-
-        def aggregate_spec(session, read_preference)
-          Builder::Aggregation.new(
-            pipeline,
-            view,
-            options.merge(session: session, read_preference: read_preference)
-          ).specification
-        end
-
         def new(options)
           Aggregation.new(view, pipeline, options)
         end
@@ -180,32 +117,29 @@ module Mongo
 
         end
 
-        def send_initial_query(server, session)
-          server.with_connection do |connection|
+        def send_initial_query(server, context)
+          if server.load_balancer?
+            # Connection will be checked in when cursor is drained.
+            connection = server.pool.check_out(context: context)
             initial_query_op(
-              session,
+              context.session,
               effective_read_preference(connection)
             ).execute_with_connection(
               connection,
-              context: Operation::Context.new(client: client, session: session)
+              context: context
             )
+          else
+            server.with_connection do |connection|
+              initial_query_op(
+                context.session,
+                effective_read_preference(connection)
+              ).execute_with_connection(
+                connection,
+                context: context
+              )
+            end
           end
         end
-
-        # Skip, sort, limit, projection are specified as pipeline stages
-        # rather than as options.
-        def cache_options
-          {
-            namespace: collection.namespace,
-            selector: pipeline,
-            read_concern: view.read_concern,
-            read_preference: view.read_preference,
-            collation: options[:collation],
-            # Aggregations can read documents from more than one collection,
-            # so they will be cleared on every write operation.
-            multi_collection: true,
-          }
-        end
       end
     end
   end

data/lib/mongo/collection/view/builder/aggregation.rb

@@ -113,17 +113,11 @@ module Mongo
               command[:readConcern] = Options::Mapper.transform_values_to_strings(
                 read_concern)
             end
-            command[:cursor] = cursor if cursor
+            command[:cursor] = batch_size_doc
             command.merge!(Options::Mapper.transform_documents(options, MAPPINGS))
             command
           end
 
-          def cursor
-            if options[:use_cursor] == true || options[:use_cursor].nil?
-              batch_size_doc
-            end
-          end
-
           def batch_size_doc
             value = options[:batch_size] || view.batch_size
             if value == 0 && write?

data/lib/mongo/collection/view/change_stream.rb

@@ -15,6 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/collection/view/aggregation/behavior'
 require 'mongo/collection/view/change_stream/retryable'
 
 module Mongo
@@ -35,7 +36,8 @@ module Mongo
       #
       #
       # @since 2.5.0
-      class ChangeStream < Aggregation
+      class ChangeStream
+        include Aggregation::Behavior
         include Retryable
 
         # @return [ String ] The fullDocument option default value.
@@ -60,6 +62,10 @@ module Mongo
         # @since 2.5.0
         attr_reader :options
 
+        # @return [ Cursor ] the underlying cursor for this operation
+        # @api private
+        attr_reader :cursor
+
         # Initialize the change stream for the provided collection view, pipeline
         # and options.
         #
@@ -125,11 +131,13 @@ module Mongo
         #
         # @since 2.5.0
         def initialize(view, pipeline, changes_for, options = {})
-          @view = view
-          @changes_for = changes_for
-          @change_stream_filters = pipeline && pipeline.dup
-          @options = options && options.dup.freeze
-          @start_after = @options[:start_after]
+          # change stream cursors can only be :iterable, so we don't allow
+          # timeout_mode to be specified.
+          perform_setup(view, options, forbid: %i[ timeout_mode ]) do
+            @changes_for = changes_for
+            @change_stream_filters = pipeline && pipeline.dup
+            @start_after = @options[:start_after]
+          end
 
           # The resume token tracked by the change stream, used only
           # when there is no cursor, or no cursor resume token
@@ -181,24 +189,30 @@ module Mongo
         # @return [ BSON::Document | nil ] A change stream document.
         # @since 2.6.0
         def try_next
+          recreate_cursor! if @timed_out
+
           raise StopIteration.new if closed?
+
           begin
             doc = @cursor.try_next
           rescue Mongo::Error => e
-            if !e.change_stream_resumable?
-              raise
-            end
-
-            # Rerun initial aggregation.
-            # Any errors here will stop iteration and break out of this
-            # method.
+            # "If a next call fails with a timeout error, drivers MUST NOT
+            # invalidate the change stream. The subsequent next call MUST
+            # perform a resume attempt to establish a new change stream on the
+            # server..."
+            #
+            # However, SocketTimeoutErrors are TimeoutErrors, but are also
+            # change-stream-resumable. To preserve existing (specified) behavior,
+            # We only count timeouts when the error is not also
+            # change-stream-resumable.
+            @timed_out = e.is_a?(Mongo::Error::TimeoutError) && !e.change_stream_resumable?
+
+            raise unless @timed_out || e.change_stream_resumable?
 
-            # Save cursor's resume token so we can use it
-            # to create a new cursor
             @resume_token = @cursor.resume_token
+            raise e if @timed_out
 
-            close
-            create_cursor!
+            recreate_cursor!(@cursor.context)
             retry
           end
 
@@ -231,14 +245,17 @@ module Mongo
         #   This method ignores any errors that occur when closing the
         #   server-side cursor.
         #
+        # @params [ Hash ] opts Options to be passed to the cursor close
+        #   command.
+        #
         # @return [ nil ] Always nil.
         #
         # @since 2.5.0
-        def close
+        def close(opts = {})
           unless closed?
             begin
-              @cursor.close
-            rescue Error::OperationFailure, Error::SocketError, Error::SocketTimeoutError, Error::MissingConnection
+              @cursor.close(opts)
+            rescue Error::OperationFailure::Family, Error::SocketError, Error::SocketTimeoutError, Error::MissingConnection
               # ignore
             end
             @cursor = nil
@@ -284,6 +301,28 @@ module Mongo
           cursor_resume_token || @resume_token
         end
 
+        # "change streams are an abstraction around tailable-awaitData cursors..."
+        #
+        # @return :tailable_await
+        def cursor_type
+          :tailable_await
+        end
+
+        # "change streams...implicitly use ITERATION mode"
+        #
+        # @return :iteration
+        def timeout_mode
+          :iteration
+        end
+
+        # Returns the value of the max_await_time_ms option that was
+        # passed to this change stream.
+        #
+        # @return [ Integer | nil ] the max_await_time_ms value
+        def max_await_time_ms
+          options[:max_await_time_ms]
+        end
+
         private
 
         def for_cluster?
@@ -298,19 +337,23 @@ module Mongo
           !for_cluster? && !for_database?
         end
 
-        def create_cursor!
+        def create_cursor!(timeout_ms = nil)
           # clear the cache because we may get a newer or an older server
           # (rolling upgrades)
           @start_at_operation_time_supported = nil
 
-          session = client.send(:get_session, @options)
+          session = client.get_session(@options)
+          context = Operation::Context.new(client: client, session: session, view: self, operation_timeouts: timeout_ms ? { operation_timeout_ms: timeout_ms } : operation_timeouts)
+
           start_at_operation_time = nil
           start_at_operation_time_supported = nil
-          @cursor = read_with_retry_cursor(session, server_selector, view) do |server|
+
+          @cursor = read_with_retry_cursor(session, server_selector, self, context: context) do |server|
             server.with_connection do |connection|
               start_at_operation_time_supported = connection.description.server_version_gte?('4.0')
 
-              result = send_initial_query(connection, session)
+              result = send_initial_query(connection, context)
+
               if doc = result.replies.first && result.replies.first.documents.first
                 start_at_operation_time = doc['operationTime']
               else
@@ -324,6 +367,7 @@ module Mongo
               result
             end
           end
+
           @start_at_operation_time = start_at_operation_time
           @start_at_operation_time_supported = start_at_operation_time_supported
         end
@@ -390,11 +434,11 @@ module Mongo
           end
         end
 
-        def send_initial_query(connection, session)
-          initial_query_op(session, view.read_preference)
+        def send_initial_query(connection, context)
+          initial_query_op(context.session, view.read_preference)
             .execute_with_connection(
               connection,
-              context: Operation::Context.new(client: client, session: session),
+              context: context,
             )
         end
 
@@ -412,6 +456,15 @@ module Mongo
         def resuming?
           !!@resuming
         end
+
+        # Recreates the current cursor (typically as a consequence of attempting
+        # to resume the change stream)
+        def recreate_cursor!(context = nil)
+          @timed_out = false
+
+          close
+          create_cursor!(context&.remaining_timeout_ms)
+        end
       end
     end
   end