mongo 2.17.4 → 2.18.0.beta1

data/lib/mongo/operation/kill_cursors.rb

@@ -18,7 +18,6 @@
 require 'mongo/operation/kill_cursors/command_builder'
 require 'mongo/operation/kill_cursors/command'
 require 'mongo/operation/kill_cursors/op_msg'
-require 'mongo/operation/kill_cursors/legacy'
 
 module Mongo
   module Operation

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

@@ -31,7 +31,10 @@ module Mongo
         private
 
         def selector(connection)
-          (spec[SELECTOR] || {}).merge(listCollections: 1)
+          (spec[SELECTOR] || {}).merge({
+            listCollections: 1,
+            comment: spec[:comment]
+          }).compact
         end
       end
     end

data/lib/mongo/operation/map_reduce/result.rb

@@ -131,6 +131,22 @@ module Mongo
           0
         end
 
+        # Get the number of documents returned by the server in this batch.
+        #
+        # Map/Reduce operation returns documents inline without using
+        # cursors; as such, the standard Mongo::Reply#returned_count does
+        # not work correctly for Map/Reduce.
+        #
+        # Note that the Map/Reduce operation is limited to max BSON document
+        # size (16 MB) in its inline result set.
+        #
+        # @return [ Integer ] The number of documents returned.
+        #
+        # @api public
+        def returned_count
+          reply.documents.length
+        end
+
         private
 
         def first_document

data/lib/mongo/operation/result.rb

@@ -97,9 +97,12 @@ module Mongo
       #   this result is for. This parameter is allowed to be nil for
       #   compatibility with existing mongo_kerberos library, but should
       #   always be not nil in the driver proper.
+      # @param [ Integer ] connection_global_id
+      #   Global id of the connection on which the operation that
+      #   this result is for was performed.
       #
       # @api private
-      def initialize(replies, connection_description = nil)
+      def initialize(replies, connection_description = nil, connection_global_id = nil)
         if replies
           if replies.is_a?(Array)
             if replies.length != 1
@@ -114,6 +117,7 @@ module Mongo
           end
           @replies = [ reply ]
           @connection_description = connection_description
+          @connection_global_id = connection_global_id
         end
       end
 
@@ -128,6 +132,12 @@ module Mongo
       # @api private
       attr_reader :connection_description
 
+      # @return [ Object ] Global is of the connection that
+      #   the operation was performed on that this result is for.
+      #
+      # @api private
+      attr_reader :connection_global_id
+
       # @api private
       def_delegators :parser,
         :not_master?, :node_recovering?, :node_shutting_down?
@@ -148,6 +158,15 @@ module Mongo
         !!@replies
       end
 
+      # Whether the result contains cursor_id
+      #
+      # @return [ true, false ] If the result contains cursor_id.
+      #
+      # @api private
+      def has_cursor_id?
+        acknowledged? && replies.last.respond_to?(:cursor_id)
+      end
+
       # Get the cursor id if the response is acknowledged.
       #
       # @note Cursor ids of 0 indicate there is no cursor on the server.
@@ -238,10 +257,7 @@ module Mongo
         end
       end
 
-      # Get the count of documents returned by the server.
-      #
-      # @example Get the number returned.
-      #   result.returned_count
+      # Get the number of documents returned by the server in this batch.
       #
       # @return [ Integer ] The number of documents returned.
       #

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

@@ -27,7 +27,8 @@ module Mongo
       include ResponseHandling
 
       def do_execute(connection, context, options = {})
-        unpin_maybe(session) do
+        session&.materialize_if_needed
+        unpin_maybe(session, connection) do
           add_error_labels(connection, context) do
             add_server_diagnostics(connection) do
               get_result(connection, context, options).tap do |result|
@@ -35,12 +36,16 @@ module Mongo
                   if session.in_transaction? &&
                     connection.description.load_balancer?
                   then
-                    if session.pinned_service_id
-                      unless session.pinned_service_id == connection.service_id
-                        raise Error::InternalDriverError, "Expected operation to use service #{session.pinned_session_id} but it used #{connection.service_id}"
+                    if session.pinned_connection_global_id
+                      unless session.pinned_connection_global_id == connection.global_id
+                        raise(
+                          Error::InternalDriverError,
+                          "Expected operation to use connection #{session.pinned_connection_global_id} but it used #{connection.global_id}"
+                        )
                       end
                     else
-                      session.pin_to_service(connection.service_id)
+                      session.pin_to_connection(connection.global_id)
+                      connection.pin
                     end
                   end
 
@@ -48,6 +53,16 @@ module Mongo
                     session.snapshot_timestamp = result.snapshot_timestamp
                   end
                 end
+
+                if result.has_cursor_id? &&
+                  connection.description.load_balancer?
+                then
+                  if result.cursor_id == 0
+                    connection.unpin
+                  else
+                    connection.pin
+                  end
+                end
                 process_result(result, connection)
               end
             end
@@ -82,7 +97,7 @@ module Mongo
         message = build_message(connection, context)
         message = message.maybe_encrypt(connection, context)
         reply = connection.dispatch([ message ], context, options)
-        [reply, connection.description]
+        [reply, connection.description, connection.global_id]
       end
 
       # @param [ Mongo::Server::Connection ] connection The connection on which

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

@@ -32,11 +32,23 @@ module Mongo
       #
       # @return [ Mongo::Operation::Result ] The operation result.
       def execute(server, context:, options: {})
-        server.with_connection(service_id: context.service_id) do |connection|
-          operation = final_operation(connection)
-          operation.execute(connection, context: context, options: options)
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          execute_with_connection(connection, context: context, options: options)
         end
       end
+
+      # Execute the operation.
+      #
+      # @param [ Mongo::Server::Connection ] connection The connection to send
+      #   the operation through.
+      # @param [ Operation::Context ] context The operation context.
+      # @param [ Hash ] options Operation execution options.
+      #
+      # @return [ Mongo::Operation::Result ] The operation result.
+      def execute_with_connection(connection, context:, options: {})
+        operation = final_operation(connection)
+        operation.execute(connection, context: context, options: options)
+      end
     end
   end
 end

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

@@ -30,7 +30,7 @@ module Mongo
       #   the operation is performed.
       # @param [ Mongo::Operation::Context ] context The operation context.
       def validate_result(result, connection, context)
-        unpin_maybe(context.session) do
+        unpin_maybe(context.session, connection) do
           add_error_labels(connection, context) do
             add_server_diagnostics(connection) do
               result.validate!
@@ -81,19 +81,20 @@ module Mongo
         end
       end
 
-      # Unpins the session if the session is pinned and the yielded to block
-      # raises errors that are required to unpin the session.
+      # Unpins the session and/or the connection if  the yielded to block
+      # raises errors that are required to unpin the session and the connection.
       #
       # @note This method takes the session as an argument because this module
       #   is included in BulkWrite which does not store the session in the
       #   receiver (despite Specifiable doing so).
       #
       # @param [ Session | nil ] session Session to consider.
-      def unpin_maybe(session)
+      # @param [ Connection | nil ] connection Connection to unpin.
+      def unpin_maybe(session, connection)
         yield
       rescue Mongo::Error => e
         if session
-          session.unpin_maybe(e)
+          session.unpin_maybe(e, connection)
         end
         raise
       end

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

@@ -257,14 +257,10 @@ module Mongo
         super.tap do |message|
           if session = context.session
             # Serialize the message to detect client-side problems,
-            # such as invalid BSON keys or too large messages.
-            # The message will be serialized again
+            # such as invalid BSON keys. The message will be serialized again
             # later prior to being sent to the connection.
-            buf = BSON::ByteBuffer.new
-            message.serialize(buf)
-            if buf.length > connection.max_message_size
-              raise Error::MaxMessageSize.new(connection.max_message_size)
-            end
+            message.serialize(BSON::ByteBuffer.new)
+
             session.update_state!
           end
         end

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

@@ -35,21 +35,27 @@ module Mongo
       #
       # @since 2.5.2
       def execute(server, context:)
-        server.with_connection(service_id: context.service_id) do |connection|
-          validate!(connection)
-          op = if connection.features.op_msg_enabled?
-              self.class::OpMsg.new(spec)
-            elsif !acknowledged_write?
-              self.class::Legacy.new(spec)
-            else
-              self.class::Command.new(spec)
-            end
-
-          result = op.execute(connection, context: context)
-          validate_result(result, connection, context)
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          execute_with_connection(connection, context: context)
         end
       end
 
+      # Execute the operation.
+      #
+      # @param [ Mongo::Server::Connection ] connection The connection to send
+      #   the operation through.
+      # @param [ Operation::Context ] context The operation context.
+      # @param [ Hash ] options Operation execution options.
+      #
+      # @return [ Mongo::Operation::Result ] The operation result.
+      def execute_with_connection(connection, context:)
+        validate!(connection)
+        op = self.class::OpMsg.new(spec)
+
+        result = op.execute(connection, context: context)
+        validate_result(result, connection, context)
+      end
+
       # Execute the bulk write operation.
       #
       # @param [ Mongo::Server::Connection ] connection The connection over

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

@@ -37,7 +37,8 @@ module Mongo
           {
             update: coll_name,
             ordered: ordered?,
-            let: spec[:let]
+            let: spec[:let],
+            comment: spec[:comment]
           }.compact
         end
 

data/lib/mongo/operation/update.rb

@@ -17,7 +17,6 @@
 
 require 'mongo/operation/update/command'
 require 'mongo/operation/update/op_msg'
-require 'mongo/operation/update/legacy'
 require 'mongo/operation/update/result'
 require 'mongo/operation/update/bulk_result'
 

data/lib/mongo/protocol/caching_hash.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
+# Copyright (C) 2014-2022 MongoDB Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Mongo
+  module Protocol
+
+    # A Hash that caches the results of #to_bson.
+    #
+    # @api private
+    class CachingHash
+
+      def initialize(hash)
+        @hash = hash
+      end
+
+      def bson_type
+        Hash::BSON_TYPE
+      end
+
+      # Caches the result of to_bson and writes it to the given buffer on subsequent
+      # calls to this method. If this method is originally called without validation,
+      # and then is subsequently called with validation, we will want to recalculate
+      # the to_bson to trigger the validations.
+      #
+      # @param [ BSON::ByteBuffer ] buffer The encoded BSON buffer to append to.
+      # @param [ true, false ] validating_keys Whether keys should be validated when serializing.
+      #
+      # @return [ BSON::ByteBuffer ] The buffer with the encoded object.
+      def to_bson(buffer = BSON::ByteBuffer.new, validating_keys = BSON::Config.validating_keys?)
+        if !@bytes
+          @bytes = @hash.to_bson(BSON::ByteBuffer.new, validating_keys).to_s
+        elsif needs_validation?(validating_keys)
+          @validated = true
+          return @hash.to_bson(buffer, validating_keys)
+        end
+        @validated ||= validating_keys
+        buffer.put_bytes(@bytes)
+      end
+
+      private
+
+      # Checks the current value for validating keys, and whether or not this
+      # bson has been validated in the past, and decides if we need to recalculate
+      # the to_bson to check the validations.
+      #
+      # @param [ true, false ] validating_keys Whether keys should be validated when serializing.
+      #
+      # @return [ true, false ] Whether or not the bson needs to be recalculated
+      #   with validation.
+      def needs_validation?(validating_keys)
+        !@validated && validating_keys
+      end
+    end
+  end
+end

data/lib/mongo/protocol/msg.rb

@@ -87,7 +87,9 @@ module Mongo
         ] + @sequences.map do |section|
           {type: 1, payload: {
             identifier: section.identifier,
-            sequence: section.documents,
+            sequence: section.documents.map do |doc|
+              CachingHash.new(doc)
+            end,
           }}
         end
         @request_id = nil
@@ -227,6 +229,22 @@ module Mongo
           if cmd.key?('$db') && !enc_cmd.key?('$db')
             enc_cmd['$db'] = cmd['$db']
           end
+          # This code MUST be removed as soon as server starts accepting
+          # contention as int32.
+          if schema = enc_cmd.dig('encryptionInformation', 'schema')
+            enc_cmd['encryptionInformation']['schema'] = schema.map do |coll, params|
+              if params['fields']
+                params['fields'] = params['fields'].map do |field|
+                  if contention = field.dig('queries', 'contention')
+                    field['queries']['contention'] = BSON::Int64.new(contention)
+                  end
+                  field
+                end
+              end
+              [coll, params]
+            end.to_h
+          end
+          # End of code to be removed
 
           Msg.new(@flags, @options, enc_cmd)
         else
@@ -298,6 +316,24 @@ module Mongo
         Msg.new(@flags, @options, main_document, *@sequences)
       end
 
+      # Returns the number of documents returned from the server.
+      #
+      # The Msg instance must be for a server reply and the reply must return
+      # an active cursor (either a newly created one or one whose iteration is
+      # continuing via getMore).
+      #
+      # @return [ Integer ] Number of returned documents.
+      def number_returned
+        if doc = documents.first
+          if cursor = doc['cursor']
+            if batch = cursor['firstBatch'] || cursor['nextBatch']
+              return batch.length
+            end
+          end
+        end
+        raise NotImplementedError, "number_returned is only defined for cursor replies"
+      end
+
       private
 
       # Validate that the documents in this message are all smaller than the

data/lib/mongo/protocol.rb

@@ -6,6 +6,7 @@ require 'mongo/protocol/serializers'
 require 'mongo/protocol/registry'
 require 'mongo/protocol/bit_vector'
 require 'mongo/protocol/message'
+require 'mongo/protocol/caching_hash'
 
 # Client Requests
 require 'mongo/protocol/compressed'

data/lib/mongo/query_cache.rb

@@ -280,6 +280,21 @@ module Mongo
       ensure
         QueryCache.clear
       end
+
+      # ActiveJob middleware that activates the query cache for each job.
+      module ActiveJob
+        def self.included(base)
+          base.class_eval do
+            around_perform do |_job, block|
+              QueryCache.cache do
+                block.call
+              end
+            ensure
+              QueryCache.clear
+            end
+          end
+        end
+      end
     end
   end
 end

data/lib/mongo/retryable.rb

@@ -191,25 +191,29 @@ module Mongo
     # @note This only retries operations on not master failures, since it is
     #   the only case we can be sure a partial write did not already occur.
     #
-    # @param [ nil | Session ] session Optional session to use with the operation.
     # @param [ nil | Hash | WriteConcern::Base ] write_concern The write concern.
-    # @param [ true | false ] ending_transaction True if the write operation is abortTransaction or
-    #   commitTransaction, false otherwise.
+    # @param [ true | false ] ending_transaction True if the write operation is
+    #   abortTransaction or commitTransaction, false otherwise.
+    # @param [ Context ] context The context for the operation.
     # @param [ Proc ] block The block to execute.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #   write should be sent.
     # @yieldparam [ Integer ] txn_num Transaction number (NOT the ACID kind).
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @return [ Result ] The result of the operation.
     #
     # @since 2.1.0
-    def write_with_retry(session, write_concern, ending_transaction = false, &block)
+    def write_with_retry(write_concern, ending_transaction: false, context:, &block)
+      session = context.session
+
       if ending_transaction && !session
         raise ArgumentError, 'Cannot end a transaction without a session'
       end
 
       unless ending_transaction || retry_write_allowed?(session, write_concern)
-        return legacy_write_with_retry(nil, session, &block)
+        return legacy_write_with_retry(nil, context: context, &block)
       end
 
       # If we are here, session is not nil. A session being nil would have
@@ -218,23 +222,46 @@ module Mongo
       server = select_server(cluster, ServerSelector.primary, session)
 
       unless ending_transaction || server.retry_writes?
-        return legacy_write_with_retry(server, session, &block)
+        return legacy_write_with_retry(server, context: context, &block)
       end
 
-      txn_num = if session.in_transaction?
-        session.txn_num
-      else
-        session.next_txn_num
-      end
+      txn_num = nil
+
       begin
-        yield(server, txn_num, false)
-      rescue Error::SocketError, Error::SocketTimeoutError => e
+        connection_succeeded = false
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          connection_succeeded = true
+
+          session.materialize_if_needed
+          txn_num = if session.in_transaction?
+            session.txn_num
+          else
+            session.next_txn_num
+          end
+
+          # The context needs to be duplicated here because we will be using
+          # it later for the retry as well.
+          yield(connection, txn_num, context.dup)
+        end
+      rescue Error::SocketError, Error::SocketTimeoutError, Auth::Unauthorized => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
         if !e.label?('RetryableWriteError')
-          raise e
+          # If we get an auth error, it was raised when connecting the connection
+          # and therefore we didn't have the connection yet to add labels.
+          # Therefore, check if it is retryable, and if it is, add the label
+          # and retry it. We also want to retry this if there was a Socket error
+          # when trying to create the connection.
+          if !connection_succeeded && !session.in_transaction? && e.write_retryable?
+            e.add_label('RetryableWriteError')
+          else
+            raise e
+          end
         end
-        retry_write(e, session, txn_num, &block)
+
+        # Context#with creates a new context, which is not necessary here
+        # but the API is less prone to misuse this way.
+        retry_write(e, txn_num, context: context.with(is_retry: true), &block)
       rescue Error::OperationFailure => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
@@ -244,7 +271,9 @@ module Mongo
           raise e
         end
 
-        retry_write(e, session, txn_num, &block)
+        # Context#with creates a new context, which is not necessary here
+        # but the API is less prone to misuse this way.
+        retry_write(e, txn_num, context: context.with(is_retry: true), &block)
       end
     end
 
@@ -258,23 +287,30 @@ module Mongo
     # delegates to legacy_write_with_retry which performs write retries using
     # legacy logic.
     #
-    # @param [ nil | Session ] session Optional session to use with the operation.
     # @param [ nil | Hash | WriteConcern::Base ] write_concern The write concern.
+    # @param [ Context ] context The context for the operation.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #   write should be sent.
+    # @yieldparam [ nil ] txn_num nil as transaction number.
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @api private
-    def nro_write_with_retry(session, write_concern, &block)
+    def nro_write_with_retry(write_concern, context:, &block)
+      session = context.session
+
+      server = select_server(cluster, ServerSelector.primary, session)
       if session && session.client.options[:retry_writes]
-        server = select_server(cluster, ServerSelector.primary, session)
         begin
-          yield server
+          server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+            yield connection, nil, context
+          end
         rescue Error::SocketError, Error::SocketTimeoutError, Error::OperationFailure => e
           e.add_note('retries disabled')
           raise e
         end
       else
-        legacy_write_with_retry(nil, session, &block)
+        legacy_write_with_retry(server, context: context, &block)
       end
     end
 
@@ -287,12 +323,17 @@ module Mongo
     # @param [ Server ] server The server which should be used for the
     #   operation. If not provided, the current primary will be retrieved from
     #   the cluster.
-    # @param [ nil | Session ] session Optional session to use with the operation.
+    # @param [ Context ] context The context for the operation.
     #
-    # @yieldparam [ Server ] server The server to which the write should be sent.
+    # @yieldparam [ Connection ] connection The connection through which the
+    #    write should be sent.
+    # @yieldparam [ nil ] txn_num nil as transaction number.
+    # @yieldparam [ Operation::Context ] context The operation context.
     #
     # @api private
-    def legacy_write_with_retry(server = nil, session = nil)
+    def legacy_write_with_retry(server = nil, context:)
+      session = context.session
+
       # This is the pre-session retry logic, and is not subject to
       # current retryable write specifications.
       # In particular it does not retry on SocketError and SocketTimeoutError.
@@ -300,7 +341,10 @@ module Mongo
       begin
         attempt += 1
         server ||= select_server(cluster, ServerSelector.primary, session)
-        yield server
+        server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+          # Legacy retries do not use txn_num
+          yield connection, nil, context.dup
+        end
       rescue Error::OperationFailure => e
         e.add_note('legacy retry')
         e.add_note("attempt #{attempt}")
@@ -331,7 +375,7 @@ module Mongo
           raise e
         end
         retry_read(e, server_selector, session, &block)
-      rescue Error::OperationFailure => e
+      rescue Error::OperationFailure, Auth::Unauthorized => e
         e.add_note('modern retry')
         e.add_note("attempt 1")
         if session.in_transaction? || !e.write_retryable?
@@ -423,7 +467,9 @@ module Mongo
       raise original_error
     end
 
-    def retry_write(original_error, session, txn_num, &block)
+    def retry_write(original_error, txn_num, context:, &block)
+      session = context.session
+
       # We do not request a scan of the cluster here, because error handling
       # for the error which triggered the retry should have updated the
       # server description and/or topology as necessary (specifically,
@@ -444,7 +490,9 @@ module Mongo
         raise Error::RaiseOriginalError
       end
       log_retry(original_error, message: 'Write retry')
-      yield(server, txn_num, true)
+      server.with_connection(connection_global_id: context.connection_global_id) do |connection|
+        yield(connection, txn_num, context)
+      end
     rescue Error::SocketError, Error::SocketTimeoutError => e
       e.add_note('modern retry')
       e.add_note('attempt 2')