mongo 2.14.0 → 2.16.0

data/lib/mongo/server/connection.rb

@@ -1,3 +1,6 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -83,11 +86,24 @@ module Mongo
       # @param [ Mongo::Server ] server The server the connection is for.
       # @param [ Hash ] options The connection options.
       #
-      # @option options [ Integer ] :generation Connection pool's generation
-      #   for this connection.
+      # @option options [ Integer ] :generation The generation of this
+      #   connection. The generation should only be specified in this option
+      #   when not in load-balancing mode, and it should be the generation
+      #   of the connection pool when the connection is created. In
+      #   load-balancing mode, the generation is set on the connection
+      #   after the handshake completes.
+      # @option options [ Hash ] :server_api The requested server API version.
+      #   This hash can have the following items:
+      #   - *:version* -- string
+      #   - *:strict* -- boolean
+      #   - *:deprecation_errors* -- boolean
       #
       # @since 2.0.0
       def initialize(server, options = {})
+        if server.load_balancer? && options[:generation]
+          raise ArgumentError, "Generation cannot be set when server is a load balancer"
+        end
+
         @id = server.next_connection_id
         @monitoring = server.monitoring
         @options = options.freeze
@@ -165,6 +181,15 @@ module Mongo
           # authenticated and be usable.
           @socket, @description, @compressor = do_connect
 
+          if server.load_balancer?
+            if Lint.enabled?
+              unless service_id
+                raise Error::InternalDriverError, "The connection is to a load balancer and it must have service_id set here, but does not"
+              end
+            end
+            @generation = connection_pool.generation_manager.generation(service_id: service_id)
+          end
+
           publish_cmap_event(
             Monitoring::Event::Cmap::ConnectionReady.new(address, id)
           )
@@ -177,7 +202,7 @@ module Mongo
       # Separate method to permit easier mocking in the test suite.
       #
       # @return [ Array<Socket, Server::Description> ] Connected socket and
-      #   a server description instance from the ismaster response of the
+      #   a server description instance from the hello response of the
       #   returned socket.
       private def do_connect
         socket = add_server_diagnostics do
@@ -306,7 +331,12 @@ module Mongo
           yield
         rescue Error::SocketError => e
           @error = e
-          @server.unknown!(generation: e.generation, stop_push_monitor: true)
+          @server.unknown!(
+            generation: e.generation,
+            # or description.service_id?
+            service_id: e.service_id,
+            stop_push_monitor: true,
+          )
           raise
         rescue Error::SocketTimeoutError => e
           @error = e

data/lib/mongo/server/connection_base.rb

@@ -1,3 +1,6 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
 # Copyright (C) 2019-2020 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -27,7 +30,7 @@ module Mongo
       include Monitoring::Publishable
 
       # The maximum allowed size in bytes that a user-supplied document may
-      # take up when serialized, if the server's ismaster response does not
+      # take up when serialized, if the server's hello response does not
       # include maxBsonObjectSize field.
       #
       # The commands that are sent to the server may exceed this size by
@@ -63,7 +66,7 @@ module Mongo
                      :update_cluster_time
 
       # Returns the server description for this connection, derived from
-      # the isMaster response for the handshake performed on this connection.
+      # the hello response for the handshake performed on this connection.
       #
       # @note A connection object that hasn't yet connected (handshaken and
       #   authenticated, if authentication is required) does not have a
@@ -85,12 +88,20 @@ module Mongo
         :max_message_size,
         :mongos?
 
+      # @return [ nil | Object ] The service id, if any.
+      def service_id
+        description&.service_id
+      end
+
       # Connection pool generation from which this connection was created.
       # May be nil.
       #
       # @return [ Integer | nil ] Connection pool generation.
       def generation
-        options[:generation]
+        # If the connection is to a load balancer, @generation is set
+        # after handshake completes. If the connection is to another server
+        # type, generation is specified during connection creation.
+        @generation || options[:generation]
       end
 
       def app_metadata
@@ -105,7 +116,7 @@ module Mongo
           if same
             @server.app_metadata
           else
-            AppMetadata.new(options)
+            AppMetadata.new(options.merge(purpose: @server.app_metadata.purpose))
           end
         end
       end
@@ -124,7 +135,7 @@ module Mongo
       #
       # @param [ Array<Message> ] messages A one-element array containing
       #   the message to dispatch.
-      # @param [ Integer ] operation_id The operation id to link messages.
+      # @param [ Operation::Context ] context The operation context.
       # @param [ Hash ] options
       #
       # @option options [ Boolean ] :deserialize_as_bson Whether to deserialize
@@ -133,33 +144,40 @@ module Mongo
       #
       # @return [ Protocol::Message | nil ] The reply if needed.
       #
+      # @raise [ Error::SocketError | Error::SocketTimeoutError ] When there is a network error.
+      #
       # @since 2.0.0
-      def dispatch(messages, operation_id = nil, client = nil, options = {})
+      def dispatch(messages, context, options = {})
         # The monitoring code does not correctly handle multiple messages,
         # and the driver internally does not send more than one message at
         # a time ever. Thus prohibit multiple message use for now.
         if messages.length != 1
           raise ArgumentError, 'Can only dispatch one message at a time'
         end
+        if description.unknown?
+          raise Error::InternalDriverError, "Cannot dispatch a message on a connection with unknown description: #{description.inspect}"
+        end
         message = messages.first
-        deliver(message, client, options)
+        deliver(message, context, options)
       end
 
       private
 
-      def deliver(message, client, options = {})
+      # @raise [ Error::SocketError | Error::SocketTimeoutError ] When there is a network error.
+      def deliver(message, context, options = {})
         if Lint.enabled? && !@socket
           raise Error::LintError, "Trying to deliver a message over a disconnected connection (to #{address})"
         end
-        buffer = serialize(message, client)
+        buffer = serialize(message, context)
         ensure_connected do |socket|
           operation_id = Monitoring.next_operation_id
-          command_started(address, operation_id, message.payload,
+          started_event = command_started(address, operation_id, message.payload,
             socket_object_id: socket.object_id, connection_id: id,
             connection_generation: generation,
             server_connection_id: description.server_connection_id,
+            service_id: description.service_id,
           )
-          start = Time.now
+          start = Utils.monotonic_time
           result = nil
           begin
             result = add_server_diagnostics do
@@ -171,30 +189,38 @@ module Mongo
               end
             end
           rescue Exception => e
-            total_duration = Time.now - start
-            command_failed(nil, address, operation_id, message.payload, e.message, total_duration)
+            total_duration = Utils.monotonic_time - start
+            command_failed(nil, address, operation_id, message.payload,
+              e.message, total_duration,
+              started_event: started_event,
+              service_id: description.service_id,
+            )
             raise
           else
-            total_duration = Time.now - start
-            command_completed(result, address, operation_id, message.payload, total_duration)
+            total_duration = Utils.monotonic_time - start
+            command_completed(result, address, operation_id, message.payload,
+              total_duration,
+              started_event: started_event,
+              service_id: description.service_id,
+            )
           end
-          if client && result
-            result = result.maybe_decrypt(client)
+          if result && context.decrypt?
+            result = result.maybe_decrypt(context)
           end
           result
         end
       end
 
-      def serialize(message, client, buffer = BSON::ByteBuffer.new)
+      def serialize(message, context, buffer = BSON::ByteBuffer.new)
         # Driver specifications only mandate the fixed 16MiB limit for
         # serialized BSON documents. However, the server returns its
-        # active serialized BSON document size limit in the ismaster response,
+        # active serialized BSON document size limit in the hello response,
         # which is +max_bson_object_size+ below. The +DEFAULT_MAX_BSON_OBJECT_SIZE+
         # is the 16MiB value mandated by the specifications which we use
-        # only as the default if the server's ismaster did not contain
+        # only as the default if the server's hello did not contain
         # maxBsonObjectSize.
         max_bson_size = max_bson_object_size || DEFAULT_MAX_BSON_OBJECT_SIZE
-        if client && client.encrypter && client.encrypter.encrypt?
+        if context.encrypt?
           # The client-side encryption specification requires bulk writes to
           # be split at a reduced maxBsonObjectSize. If this message is a bulk
           # write and its size exceeds the reduced size limit, the serializer

data/lib/mongo/server/connection_common.rb

@@ -1,3 +1,6 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
 # Copyright (C) 2020 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -49,8 +52,67 @@ module Mongo
       # @api private
       attr_reader :pid
 
+      # Build a document that should be used for connection handshake.
+      #
+      # @param [ Server::AppMetadata ] app_metadata Application metadata
+      # @param [ BSON::Document ] speculative_auth_doc The speculative
+      #   authentication document, if any.
+      # @param [ true | false ] load_balancer Whether the connection is to
+      #   a load balancer.
+      # @param server_api [ Hash | nil ] server_api Server API version.
+      #
+      # @return [BSON::Document] Document that should be sent to a server
+      #     for handshake purposes.
+      #
+      # @api private
+      def handshake_document(app_metadata, speculative_auth_doc: nil, load_balancer: false, server_api: nil)
+        serv_api = app_metadata.server_api || server_api
+        document = if serv_api
+                     HELLO_DOC.merge(Utils.transform_server_api(serv_api))
+                   else
+                     LEGACY_HELLO_DOC
+                   end
+        document.merge(app_metadata.validated_document).tap do |doc|
+          if speculative_auth_doc
+            doc.update(speculativeAuthenticate: speculative_auth_doc)
+          end
+          if load_balancer
+            doc.update(loadBalanced: true)
+          end
+        end
+      end
+
+      # Build a command that should be used for connection handshake.
+      #
+      # @param [ BSON::Document ] handshake_document Document that should be
+      #   sent to a server for handshake purpose.
+      #
+      # @return [ Protocol::Message ] Command that should be sent to a server
+      #   for handshake purposes.
+      #
+      # @api private
+      def handshake_command(handshake_document)
+        if handshake_document['apiVersion']
+          Protocol::Msg.new(
+            [], {}, handshake_document.merge({'$db' => Database::ADMIN})
+          )
+        else
+          Protocol::Query.new(
+            Database::ADMIN,
+            Database::COMMAND,
+            handshake_document,
+            :limit => -1
+          )
+        end
+      end
+
+
       private
 
+      HELLO_DOC = BSON::Document.new({ hello: 1 }).freeze
+
+      LEGACY_HELLO_DOC = BSON::Document.new({ isMaster: 1, helloOk: true }).freeze
+
       attr_reader :socket
 
       def set_compressor!(reply)
@@ -88,14 +150,22 @@ module Mongo
         # Server::Monitor::Connection does not reference its server, but
         # knows its address. Server::Connection delegates the address to its
         # server.
-        note = "on #{address.seed}"
+        note = +"on #{address.seed}"
         if respond_to?(:id)
           note << ", connection #{generation}:#{id}"
         end
+        # Non-monitoring connections have service id.
+        # Monitoring connections do not.
+        if respond_to?(:service_id) && service_id
+          note << ", service id #{service_id}"
+        end
         e.add_note(note)
         if respond_to?(:generation)
           # Non-monitoring connections
           e.generation = generation
+          if respond_to?(:description)
+            e.service_id = service_id
+          end
         end
         raise e
       end

data/lib/mongo/server/connection_pool/generation_manager.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
+# Copyright (C) 2021 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
+  class Server
+    class ConnectionPool
+
+      # @api private
+      class GenerationManager
+
+        def initialize(server:)
+          @map = Hash.new { |hash, key| hash[key] = 1 }
+          @server = server
+          @lock = Mutex.new
+        end
+
+        attr_reader :server
+
+        def generation(service_id: nil)
+          if service_id
+            unless server.load_balancer?
+              raise ArgumentError, "Generation scoping to services is only available in load-balanced mode, but the server at #{server.address} is not a load balancer"
+            end
+          else
+            if server.load_balancer?
+              raise ArgumentError, "The server at #{server.address} is a load balancer and therefore does not have a single global generation"
+            end
+          end
+          @lock.synchronize do
+            @map[service_id]
+          end
+        end
+
+        def bump(service_id: nil)
+          @lock.synchronize do
+            if service_id
+              @map[service_id] += 1
+            else
+              # When service id is not supplied, one of two things may be
+              # happening;
+              #
+              # 1. The pool is not to a load balancer, in which case we only
+              #    need to increment the generation for the nil service_id.
+              # 2. The pool is to a load balancer, in which case we need to
+              #    increment the generation for each service.
+              #
+              # Incrementing everything in the map accomplishes both tasks.
+              @map.each do |k, v|
+                @map[k] += 1
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/mongo/server/connection_pool/populator.rb

@@ -1,3 +1,6 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
 # Copyright (C) 2019-2020 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

data/lib/mongo/server/connection_pool.rb

@@ -1,3 +1,6 @@
+# frozen_string_literal: true
+# encoding: utf-8
+
 # Copyright (C) 2014-2020 MongoDB Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -11,7 +14,6 @@
 # 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.
-require 'mongo/server/connection_pool/populator'
 
 module Mongo
   class Server
@@ -108,7 +110,7 @@ module Mongo
         @server = server
         @options = options.freeze
 
-        @generation = 1
+        @generation_manager = GenerationManager.new(server: server)
         @closed = false
 
         # A connection owned by this pool should be either in the
@@ -186,12 +188,14 @@ module Mongo
         @max_idle_time ||= options[:max_idle_time]
       end
 
+      # @api private
+      attr_reader :generation_manager
+
       # @return [ Integer ] generation Generation of connections currently
       #   being used by the queue.
       #
-      # @since 2.9.0
       # @api private
-      attr_reader :generation
+      def_delegator :generation_manager, :generation
 
       # Size of the connection pool.
       #
@@ -274,7 +278,7 @@ module Mongo
       #   and remains so for longer than the wait timeout.
       #
       # @since 2.9.0
-      def check_out
+      def check_out(service_id: nil)
         check_invariants
 
         publish_cmap_event(
@@ -291,7 +295,7 @@ module Mongo
           raise Error::PoolClosedError.new(@server.address, self)
         end
 
-        deadline = Time.now + wait_timeout
+        deadline = Utils.monotonic_time + wait_timeout
         pid = Process.pid
         connection = nil
         # It seems that synchronize sets up its own loop, thus a simple break
@@ -303,7 +307,12 @@ module Mongo
             # a connection while this thread is waiting for one.
             @lock.synchronize do
               until @available_connections.empty?
-                connection = @available_connections.pop
+                connection = next_available_connection(service_id: service_id)
+
+                # If service_id is not nil, connection may be nil here
+                # even if there are available connections in the pool
+                # (they could be to other services).
+                break unless connection
 
                 if connection.pid != pid
                   log_warn("Detected PID change - Mongo client should have been reconnected (old pid #{connection.pid}, new pid #{pid}")
@@ -312,7 +321,7 @@ module Mongo
                   next
                 end
 
-                if connection.generation != generation
+                if connection.generation != generation(service_id: connection.service_id)
                   # Stale connections should be disconnected in the clear
                   # method, but if any don't, check again here
                   connection.disconnect!(reason: :stale)
@@ -332,16 +341,26 @@ module Mongo
                 throw(:done)
               end
 
-              # Ruby does not allow a thread to lock a mutex which it already
-              # holds.
-              if unsynchronized_size < max_size
-                connection = create_connection
-                @pending_connections << connection
-                throw(:done)
+              if service_id
+                # If we need a connection to a particular service, we can't
+                # create one if we don't already have one, but we can wait
+                # for an in-progress operation to return such a connection
+                # to the pool, or for the populator to create a suitable
+                # connection.
+              else
+                # If we are below pool capacity, create a new connection.
+                #
+                # Ruby does not allow a thread to lock a mutex which it already
+                # holds.
+                if unsynchronized_size < max_size
+                  connection = create_connection
+                  @pending_connections << connection
+                  throw(:done)
+                end
               end
             end
 
-            wait = deadline - Time.now
+            wait = deadline - Utils.monotonic_time
             if wait <= 0
               publish_cmap_event(
                 Monitoring::Event::Cmap::ConnectionCheckOutFailed.new(
@@ -351,8 +370,14 @@ module Mongo
               )
 
               msg = @lock.synchronize do
+                service_id_msg = if service_id
+                  " for service #{service_id}"
+                else
+                  ''
+                end
+
                 "Timed out attempting to check out a connection " +
-                  "from pool for #{@server.address} after #{wait_timeout} sec. " +
+                  "from pool for #{@server.address}#{service_id_msg} after #{wait_timeout} sec. " +
                   "Connections in pool: #{@available_connections.length} available, " +
                   "#{@checked_out_connections.length} checked out, " +
                   "#{@pending_connections.length} pending " +
@@ -446,7 +471,7 @@ module Mongo
             # Connection was closed - for example, because it experienced
             # a network error. Nothing else needs to be done here.
             @populate_semaphore.signal
-          elsif connection.generation != @generation
+          elsif connection.generation != generation(service_id: connection.service_id)
             connection.disconnect!(reason: :stale)
             @populate_semaphore.signal
           else
@@ -470,8 +495,12 @@ module Mongo
       # @option options [ true | false ] :lazy If true, do not close any of
       #   the idle connections and instead let them be closed during a
       #   subsequent check out operation.
+      # @option options [ Object ] :service_id Discard state for the specified
+      #   service id only.
       # @option options [ true | false ] :stop_populator Whether to stop
       #   the populator background thread. For internal driver use only.
+      # @option options [ Object ] :service_id Clear connections with
+      #   the specified service id only.
       #
       # @return [ true ] true.
       #
@@ -485,18 +514,38 @@ module Mongo
           stop_populator
         end
 
+        service_id = options && options[:service_id]
+
         @lock.synchronize do
-          @generation += 1
+          @generation_manager.bump(service_id: service_id)
 
           publish_cmap_event(
-            Monitoring::Event::Cmap::PoolCleared.new(@server.address)
+            Monitoring::Event::Cmap::PoolCleared.new(
+              @server.address,
+              service_id: service_id,
+            )
           )
 
           unless options && options[:lazy]
-            until @available_connections.empty?
-              connection = @available_connections.pop
-              connection.disconnect!(reason: :stale)
-              @populate_semaphore.signal
+            if service_id
+              loop do
+                conn = @available_connections.detect do |conn|
+                  conn.service_id == service_id
+                end
+                if conn
+                  @available_connections.delete(conn)
+                  conn.disconnect!(reason: :stale)
+                  @populate_semaphore.signal
+                else
+                  break
+                end
+              end
+            else
+              until @available_connections.empty?
+                connection = @available_connections.pop
+                connection.disconnect!(reason: :stale)
+                @populate_semaphore.signal
+              end
             end
           end
         end
@@ -583,10 +632,10 @@ module Mongo
       # @return [ Object ] The result of the block.
       #
       # @since 2.0.0
-      def with_connection
+      def with_connection(service_id: nil)
         raise_if_closed!
 
-        connection = check_out
+        connection = check_out(service_id: service_id)
         yield(connection)
       ensure
         if connection
@@ -684,8 +733,6 @@ module Mongo
       # @return [ Proc ] The Finalizer.
       def self.finalize(available_connections, pending_connections, populator)
         proc do
-          populator.stop!
-
           available_connections.each do |connection|
             connection.disconnect!(reason: :pool_closed)
           end
@@ -704,9 +751,33 @@ module Mongo
 
       private
 
+      # Returns the next available connection, optionally scoped to the
+      # specified service. If no suitable connections are available,
+      # returns nil.
+      def next_available_connection(service_id: nil)
+        if service_id
+          conn = @available_connections.detect do |conn|
+            conn.service_id == service_id
+          end
+          if conn
+            @available_connections.delete(conn)
+          end
+          conn
+        else
+          @available_connections.pop
+        end
+      end
+
       def create_connection
-        connection = Connection.new(@server, options.merge(generation: generation,
-          connection_pool: self))
+        opts = options.merge(
+          connection_pool: self,
+          # Do not pass app metadata - this will be retrieved by the connection
+          # based on the auth needs.
+        )
+        unless @server.load_balancer?
+          opts[:generation] = generation
+        end
+        connection = Connection.new(@server, opts)
       end
 
       # Create a connection, connect it, and add it to the pool.
@@ -767,7 +838,11 @@ module Mongo
           raise
         end
       rescue Error::SocketError, Error::SocketTimeoutError => exc
-        @server.unknown!(generation: exc.generation, stop_push_monitor: true)
+        @server.unknown!(
+          generation: exc.generation,
+          service_id: exc.service_id,
+          stop_push_monitor: true,
+        )
         raise
       end
 
@@ -795,3 +870,6 @@ module Mongo
     end
   end
 end
+
+require 'mongo/server/connection_pool/generation_manager'
+require 'mongo/server/connection_pool/populator'