mongo 2.17.2 → 2.18.1

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

data/lib/mongo/server/connection.rb

@@ -24,6 +24,7 @@ module Mongo
     class Connection < ConnectionBase
       include Monitoring::Publishable
       include Retryable
+      include Id
       extend Forwardable
 
       # The ping command.
@@ -105,6 +106,7 @@ module Mongo
         end
 
         @id = server.next_connection_id
+        @global_id = self.class.next_id
         @monitoring = server.monitoring
         @options = options.freeze
         @server = server
@@ -112,6 +114,7 @@ module Mongo
         @last_checkin = nil
         @auth_mechanism = nil
         @pid = Process.pid
+        @pinned = false
 
         publish_cmap_event(
           Monitoring::Event::Cmap::ConnectionCreated.new(address, id)
@@ -129,6 +132,10 @@ module Mongo
       # @since 2.9.0
       attr_reader :id
 
+      # @return [ Integer ] The global ID for the connection. This will be unique
+      # across all connections.
+      attr_reader :global_id
+
       # The connection pool from which this connection was created.
       # May be nil.
       #
@@ -154,6 +161,32 @@ module Mongo
         !!@error
       end
 
+      # Whether the connection is used by a transaction or cursor operations.
+      #
+      # Pinned connections should not be disconnected and removed from a
+      # connection pool if they are idle or stale.
+      #
+      # # @return [ true | false ] Whether connection is pinned.
+      #
+      # @api private
+      def pinned?
+        @pinned
+      end
+
+      # Mark the connection as pinned.
+      #
+      # @api private
+      def pin
+        @pinned = true
+      end
+
+      # Mark the connection as not pinned.
+      #
+      # @api private
+      def unpin
+        @pinned = false
+      end
+
       # Establishes a network connection to the target address.
       #
       # If the connection is already established, this method does nothing.

data/lib/mongo/server/connection_base.rb

@@ -193,6 +193,7 @@ module Mongo
             command_failed(nil, address, operation_id, message.payload,
               e.message, total_duration,
               started_event: started_event,
+              server_connection_id: description.server_connection_id,
               service_id: description.service_id,
             )
             raise
@@ -201,6 +202,7 @@ module Mongo
             command_completed(result, address, operation_id, message.payload,
               total_duration,
               started_event: started_event,
+              server_connection_id: description.server_connection_id,
               service_id: description.service_id,
             )
           end

data/lib/mongo/server/connection_common.rb

@@ -92,7 +92,7 @@ module Mongo
       #
       # @api private
       def handshake_command(handshake_document)
-        if handshake_document['apiVersion']
+        if handshake_document['apiVersion'] || handshake_document['loadBalanced']
           Protocol::Msg.new(
             [], {}, handshake_document.merge({'$db' => Database::ADMIN})
           )
@@ -163,6 +163,9 @@ module Mongo
         if respond_to?(:generation)
           # Non-monitoring connections
           e.generation = generation
+          if respond_to?(:global_id)
+            e.connection_global_id = global_id
+          end
           if respond_to?(:description)
             e.service_id = service_id
           end