mongo 2.20.1 → 2.21.1

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

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'mongo/cursor_host'
+
 module Mongo
   class Collection
     class View
@@ -24,13 +26,7 @@ module Mongo
       #
       # @since 2.0.0
       module Iterable
-
-        # Returns the cursor associated with this view, if any.
-        #
-        # @return [ nil | Cursor ] The cursor, if any.
-        #
-        # @api private
-        attr_reader :cursor
+        include Mongo::CursorHost
 
         # Iterate through documents returned by a query with this +View+.
         #
@@ -45,48 +41,21 @@ module Mongo
         #
         # @yieldparam [ Hash ] Each matching document.
         def each
-          # If the caching cursor is closed and was not fully iterated,
-          # the documents we have in it are not the complete result set and
-          # we have no way of completing that iteration.
-          # Therefore, discard that cursor and start iteration again.
-          # The case of the caching cursor not being closed and not having
-          # been fully iterated isn't tested - see RUBY-2773.
-          @cursor = if use_query_cache? && cached_cursor && (
-            cached_cursor.fully_iterated? || !cached_cursor.closed?
-          )
-            cached_cursor
-          else
-            session = client.send(:get_session, @options)
-            select_cursor(session).tap do |cursor|
-              if use_query_cache?
-                # No need to store the cursor in the query cache if there is
-                # already a cached cursor stored at this key.
-                QueryCache.set(cursor, **cache_options)
-              end
-            end
-          end
+          @cursor = prefer_cached_cursor? ? cached_cursor : new_cursor_for_iteration
+          return @cursor.to_enum unless block_given?
 
-          if use_query_cache?
-            # If a query with a limit is performed, the query cache will
-            # re-use results from an earlier query with the same or larger
-            # limit, and then impose the lower limit during iteration.
-            limit_for_cached_query = respond_to?(:limit) ? QueryCache.normalized_limit(limit) : nil
-          end
+          limit_for_cached_query = compute_limit_for_cached_query
 
-          if block_given?
-            # Ruby versions 2.5 and older do not support arr[0..nil] syntax, so
-            # this must be a separate conditional.
-            cursor_to_iterate = if limit_for_cached_query
-              @cursor.to_a[0...limit_for_cached_query]
-            else
-              @cursor
-            end
-
-            cursor_to_iterate.each do |doc|
-              yield doc
-            end
+          # Ruby versions 2.5 and older do not support arr[0..nil] syntax, so
+          # this must be a separate conditional.
+          cursor_to_iterate = if limit_for_cached_query
+            @cursor.to_a[0...limit_for_cached_query]
           else
-            @cursor.to_enum
+            @cursor
+          end
+
+          cursor_to_iterate.each do |doc|
+            yield doc
           end
         end
 
@@ -100,7 +69,7 @@ module Mongo
         #
         # @return [ nil ] Always nil.
         #
-        # @raise [ Error::OperationFailure ] If the server cursor close fails.
+        # @raise [ Error::OperationFailure::Family ] If the server cursor close fails.
         #
         # @since 2.1.0
         def close_query
@@ -113,18 +82,25 @@ module Mongo
         private
 
         def select_cursor(session)
+          context = Operation::Context.new(
+            client: client,
+            session: session,
+            operation_timeouts: operation_timeouts,
+            view: self
+          )
+
           if respond_to?(:write?, true) && write?
             server = server_selector.select_server(cluster, nil, session, write_aggregation: true)
-            result = send_initial_query(server, session)
+            result = send_initial_query(server, context)
 
             if use_query_cache?
-              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
           else
-            read_with_retry_cursor(session, server_selector, view) do |server|
-              send_initial_query(server, session)
+            read_with_retry_cursor(session, server_selector, view, context: context) do |server|
+              send_initial_query(server, context)
             end
           end
         end
@@ -168,18 +144,13 @@ module Mongo
             batch_size: batch_size,
             hint: options[:hint],
             max_scan: options[:max_scan],
-            max_time_ms: options[:max_time_ms],
             max_value: options[:max_value],
             min_value: options[:min_value],
             no_cursor_timeout: options[:no_cursor_timeout],
             return_key: options[:return_key],
             show_disk_loc: options[:show_disk_loc],
             comment: options[:comment],
-            oplog_replay: if (v = options[:oplog_replay]).nil?
-              collection.options[:oplog_replay]
-            else
-              v
-            end,
+            oplog_replay: oplog_replay
           }
 
           if spec[:oplog_replay]
@@ -196,14 +167,52 @@ module Mongo
           end
         end
 
-        def send_initial_query(server, session = nil)
-          initial_query_op(session).execute(server, context: Operation::Context.new(client: client, session: session))
+        def send_initial_query(server, context)
+          operation = initial_query_op(context.session)
+          if server.load_balancer?
+            # Connection will be checked in when cursor is drained.
+            connection = server.pool.check_out(context: context)
+            operation.execute_with_connection(connection, context: context)
+          else
+            operation.execute(server, context: context)
+          end
         end
 
         def use_query_cache?
           QueryCache.enabled? && !collection.system_collection?
         end
 
+        # If the caching cursor is closed and was not fully iterated,
+        # the documents we have in it are not the complete result set and
+        # we have no way of completing that iteration.
+        # Therefore, discard that cursor and start iteration again.
+        def prefer_cached_cursor?
+          use_query_cache? &&
+            cached_cursor &&
+            (cached_cursor.fully_iterated? || !cached_cursor.closed?)
+        end
+
+        # Start a new cursor for use when iterating (via #each).
+        def new_cursor_for_iteration
+          session = client.get_session(@options)
+          select_cursor(session).tap do |cursor|
+            if use_query_cache?
+              # No need to store the cursor in the query cache if there is
+              # already a cached cursor stored at this key.
+              QueryCache.set(cursor, **cache_options)
+            end
+          end
+        end
+
+        def compute_limit_for_cached_query
+          return nil unless use_query_cache? && respond_to?(:limit)
+
+          # If a query with a limit is performed, the query cache will
+          # re-use results from an earlier query with the same or larger
+          # limit, and then impose the lower limit during iteration.
+          return QueryCache.normalized_limit(limit)
+        end
+
         # Add tailable cusror options to the command specifiction if needed.
         #
         # @param [ Hash ] spec The command specification.
@@ -216,6 +225,13 @@ module Mongo
             spec[:await_data] = true
           end
         end
+
+        # @return [ true | false | nil ] options[:oplog_replay], if
+        #    set, otherwise the same option from the collection.
+        def oplog_replay
+          v = options[:oplog_replay]
+          v.nil? ? collection.options[:oplog_replay] : v
+        end
       end
     end
   end

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

@@ -51,7 +51,7 @@ module Mongo
         attr_reader :reduce_function
 
         # Delegate necessary operations to the view.
-        def_delegators :view, :collection, :read, :cluster
+        def_delegators :view, :collection, :read, :cluster, :timeout_ms
 
         # Delegate necessary operations to the collection.
         def_delegators :collection, :database, :client
@@ -70,10 +70,18 @@ module Mongo
         # @yieldparam [ Hash ] Each matching document.
         def each
           @cursor = nil
-          session = client.send(:get_session, @options)
+          session = client.get_session(@options)
           server = cluster.next_primary(nil, session)
-          result = send_initial_query(server, session, context: Operation::Context.new(client: client, session: session))
-          result = send_fetch_query(server, session) unless inline?
+          context = Operation::Context.new(client: client, session: session, operation_timeouts: view.operation_timeouts)
+          if server.load_balancer?
+            # Connection will be checked in when cursor is drained.
+            connection = server.pool.check_out(context: context)
+            result = send_initial_query_with_connection(connection, context.session, context: context)
+            result = send_fetch_query_with_connection(connection, session) unless inline?
+          else
+            result = send_initial_query(server, context)
+            result = send_fetch_query(server, session) unless inline?
+          end
           @cursor = Cursor.new(view, result, server, session: session)
           if block_given?
             @cursor.each do |doc|
@@ -279,9 +287,9 @@ module Mongo
           out.respond_to?(:keys) && out.keys.first.to_s.downcase == INLINE
         end
 
-        def send_initial_query(server, session, context:)
+        def send_initial_query(server, context)
           server.with_connection do |connection|
-            send_initial_query_with_connection(connection, session, context: context)
+            send_initial_query_with_connection(connection, context.session, context: context)
           end
         end
 
@@ -305,7 +313,7 @@ module Mongo
           Builder::MapReduce.new(map_function, reduce_function, view, options.merge(session: session)).command_specification
         end
 
-        def fetch_query_op(server, session)
+        def fetch_query_op(session)
           spec = {
             coll_name: out_collection_name,
             db_name: out_database_name,
@@ -319,7 +327,16 @@ module Mongo
         end
 
         def send_fetch_query(server, session)
-          fetch_query_op(server, session).execute(server, context: Operation::Context.new(client: client, session: session))
+          fetch_query_op(session).execute(server, context: Operation::Context.new(client: client, session: session))
+        end
+
+        def send_fetch_query_with_connection(connection, session)
+          fetch_query_op(
+            session
+          ).execute_with_connection(
+            connection,
+            context: Operation::Context.new(client: client, session: session)
+          )
         end
       end
     end

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

@@ -47,12 +47,13 @@ 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 [ 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.
         #
         # @return [ Aggregation ] The aggregation object.
         #
@@ -157,6 +158,10 @@ module Mongo
         # @option opts [ Mongo::Session ] :session The session to use for the operation.
         # @option opts [ Object ] :comment A user-provided
         #   comment to attach to this command.
+        # @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.
         #
         # @return [ Integer ] The document count.
         #
@@ -182,7 +187,12 @@ module Mongo
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
-            read_with_retry(session, selector) do |server|
+            context = Operation::Context.new(
+              client: client,
+              session: session,
+              operation_timeouts: operation_timeouts(opts)
+            )
+            read_with_retry(session, selector, context) do |server|
               Operation::Count.new(
                 selector: cmd,
                 db_name: database.name,
@@ -193,7 +203,10 @@ module Mongo
                 # string key. Note that this isn't documented as valid usage.
                 collation: opts[:collation] || opts['collation'] || collation,
                 comment: opts[:comment],
-              ).execute(server, context: Operation::Context.new(client: client, session: session))
+              ).execute(
+                server,
+                context: context
+              )
             end.n.to_i
           end
         end
@@ -210,12 +223,17 @@ module Mongo
         #   MongoDB to use a specific index for the query. Requires server version 3.6+.
         # @option opts :limit [ Integer ] Max number of docs to count.
         # @option opts :max_time_ms [ Integer ] The maximum amount of time to allow the
-        #   command to run.
+        #   command to run. This option is deprecated, use
+        #   :timeout_ms instead.
         # @option opts [ Hash ] :read The read preference options.
         # @option opts [ Hash ] :collation The collation to use.
         # @option opts [ Mongo::Session ] :session The session to use for the operation.
         # @option ops [ Object ] :comment A user-provided
         #   comment to attach to this command.
+        # @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.
         #
         # @return [ Integer ] The document count.
         #
@@ -227,7 +245,7 @@ module Mongo
           pipeline << { :'$limit' => opts[:limit] } if opts[:limit]
           pipeline << { :'$group' => { _id: 1, n: { :'$sum' => 1 } } }
 
-          opts = opts.slice(:hint, :max_time_ms, :read, :collation, :session, :comment)
+          opts = opts.slice(:hint, :max_time_ms, :read, :collation, :session, :comment, :timeout_ms)
           opts[:collation] ||= collation
 
           first = aggregate(pipeline, opts).first
@@ -247,6 +265,10 @@ module Mongo
         # @option opts [ Hash ] :read The read preference options.
         # @option opts [ Object ] :comment A user-provided
         #   comment to attach to this command.
+        # @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.
         #
         # @return [ Integer ] The document count.
         #
@@ -267,8 +289,12 @@ module Mongo
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
-            read_with_retry(session, selector) do |server|
-              context = Operation::Context.new(client: client, session: session)
+            context = Operation::Context.new(
+              client: client,
+              session: session,
+              operation_timeouts: operation_timeouts(opts)
+            )
+            read_with_retry(session, selector, context) do |server|
               cmd = { count: collection.name }
               cmd[:maxTimeMS] = opts[:max_time_ms] if opts[:max_time_ms]
               if read_concern
@@ -284,7 +310,7 @@ module Mongo
               result.n.to_i
             end
           end
-        rescue Error::OperationFailure => exc
+        rescue Error::OperationFailure::Family => exc
           if exc.code == 26
             # NamespaceNotFound
             # This should only happen with the aggregation pipeline path
@@ -331,7 +357,12 @@ module Mongo
           read_pref = opts[:read] || read_preference
           selector = ServerSelector.get(read_pref || server_selector)
           with_session(opts) do |session|
-            read_with_retry(session, selector) do |server|
+            context = Operation::Context.new(
+              client: client,
+              session: session,
+              operation_timeouts: operation_timeouts(opts)
+            )
+            read_with_retry(session, selector, context) do |server|
               Operation::Distinct.new(
                 selector: cmd,
                 db_name: database.name,
@@ -342,7 +373,10 @@ module Mongo
                 # For some reason collation was historically accepted as a
                 # string key. Note that this isn't documented as valid usage.
                 collation: opts[:collation] || opts['collation'] || collation,
-              ).execute(server, context: Operation::Context.new(client: client, session: session))
+              ).execute(
+                server,
+                context: context
+              )
             end.first['values']
           end
         end
@@ -627,6 +661,15 @@ module Mongo
           configure(:cursor_type, type)
         end
 
+        # The per-operation timeout in milliseconds. Must a positive integer.
+        #
+        # @param [ Integer ] timeout_ms Timeout value.
+        #
+        # @return [ Integer, View ] Either the timeout_ms value or a new +View+.
+        def timeout_ms(timeout_ms = nil)
+          configure(:timeout_ms, timeout_ms)
+        end
+
         # @api private
         def read_concern
           if options[:session] && options[:session].in_transaction?
@@ -656,24 +699,10 @@ module Mongo
           end
         end
 
-        private
-
-        def collation(doc = nil)
-          configure(:collation, doc)
-        end
-
-        def server_selector
-          @server_selector ||= if options[:session] && options[:session].in_transaction?
-            ServerSelector.get(read_preference || client.server_selector)
-          else
-            ServerSelector.get(read_preference || collection.server_selector)
-          end
-        end
-
         def parallel_scan(cursor_count, options = {})
           if options[:session]
             # The session would be overwritten by the one in +options+ later.
-            session = client.send(:get_session, @options)
+            session = client.get_session(@options)
           else
             session = nil
           end
@@ -707,11 +736,31 @@ module Mongo
               session: session,
               connection_global_id: result.connection_global_id,
             )
-            result = op.execute(server, context: context)
+            result = if server.load_balancer?
+                       # Connection will be checked in when cursor is drained.
+                       connection = server.pool.check_out(context: context)
+                       op.execute_with_connection(connection, context: context)
+                     else
+                       op.execute(server, context: context)
+                     end
             Cursor.new(self, result, server, session: session)
           end
         end
 
+        private
+
+        def collation(doc = nil)
+          configure(:collation, doc)
+        end
+
+        def server_selector
+          @server_selector ||= if options[:session] && options[:session].in_transaction?
+            ServerSelector.get(read_preference || client.server_selector)
+          else
+            ServerSelector.get(read_preference || collection.server_selector)
+          end
+        end
+
         def validate_doc!(doc)
           raise Error::InvalidDocument.new unless doc.respond_to?(:keys)
         end