RUBY-2746 RUBY-2851 RUBY-2737 Provide options to limit number of mongos servers used in connecting to sharded clusters (#2380)

* RUBY-2746 add spec tests

* RUBY-2746 add another spec test

* RUBY-2746 implemented srvServiceName, 5 tests still failing

* RUBY-2746 added srvMaxHosts implementation and fix test for it

* RUBY-2746 add type that raises error on no parse, ALL TESTS PASSING

* RUBY-2746 move uri validation for replica set

* RUBY-2746 uri options passing, still waiting on bad input failure direction

* RUBY-2746 remove offending tests and integer_raise

* RUBY-2746 clean up old test files and combine runners

* RUBY-2746 add support for ruby options

* RUBY-2746 add tests for ruby options

* RUBY-2746 cleanup comments

* RUBY-2746 fix tests, naming, method signatures, set srv_max_hosts to nil on 0

* RUBY-2746 don't set srv_max_hosts on 0

* RUBY-2746 add docs for new options

* RUBY-2746 fix bug, add test for srv_max_hosts

* RUBY-2746 add tests for different valus of srvMaxHosts

* RUBY-2746 add srvServiceName test

* RUBY-2746 fix spec tests, make srv_max_hosts accept 0

* RUBY-2746 fix srv tests, disregard sorting

* RUBY-2746 fix tests sorting

* RUBY-2746 update comments and docs

Author: neilshweky

docs/reference/create-client.txt

@@ -235,7 +235,7 @@ routers explicitly, the driver will not discover remaining routers that
 may be configured and will not attempt to connect to them.
 
 The driver will automatically balance the operation load among the routers
-it is aware of. 
+it is aware of.
 
 To connect to a MongoDB Atlas cluster which is deployed as a sharded cluster,
 connect to the URI:
@@ -537,7 +537,7 @@ Ruby Options
    * - ``:max_idle_time``
      - The maximum time, in seconds, that a connection can be idle before it
        is closed by the connection pool.
-       
+
        *Warning:* when connected to a load balancer, the driver uses existing
        connections for iterating cursors (which includes change streams)
        and executing transactions. Setting an idle time via this option may
@@ -669,11 +669,11 @@ Ruby Options
        - ``:version`` (String)
        - ``:strict`` (true or false)
        - ``:deprecation_errors`` (true or false)
-       
+
        Note that the server API version can only be specified as a Ruby option,
        not as a URI option, and it cannot be overridden for database and
        collection objects.
-       
+
        If server API version is changed on a client (such as via the ``with``
        call), the entire API version hash is replaced with the new specification
        (the old and the new individual fields are NOT merged).
@@ -694,6 +694,24 @@ Ruby Options
      - ``Float``
      - none
 
+   * - ``:srv_max_hosts``
+     - The maximum number of mongoses that the driver will communicate with
+       for sharded topologies. If this option is set to 0, there will
+       be no maximum number of mongoses. If the given URI resolves
+       to more hosts than ``:srv_max_hosts``, the client will ramdomly
+       choose an ``:srv_max_hosts`` sized subset of hosts. Note that the
+       hosts that the driver ignores during client construction will never
+       be used. If the hosts chosen by the driver become unavailable, the
+       client will quit working completely, even though the deployment has
+       other functional mongoses.
+     - ``Integer``
+     - 0
+
+   * - ``:srv_service_name``
+     - The service name to use in the SRV DNS query.
+     - ``String``
+     - mongodb
+
    * - ``:ssl``
      - Tell the client to connect to the servers via TLS.
      - ``Boolean``
@@ -1012,6 +1030,12 @@ URI options are explained in detail in the :manual:`Connection URI reference
        invalid values (e.g. negative and non-numeric values), invalid values
        provided via this URI option are ignored with a warning.
 
+  * - srvMaxHosts=Integer
+    - ``:srv_max_hosts => Integer``
+
+  * - srvServiceName=String
+    - ``:srv_service_name => String``
+
    * - ssl=Boolean
      - ``:ssl => true|false``
 
@@ -1874,7 +1898,7 @@ the highest compression at the same CPU consumption compared to the other
 compressors. For maximum server compatibility all three compressors can be
 specified, e.g. as ``compressors: ["zstd", "snappy", "zlib"]``.
 
-``zstd`` compressor requires the 
+``zstd`` compressor requires the
 `zstd-ruby <https://rubygems.org/gems/zstd-ruby>`_ library to be installed.
 ``snappy`` compressor requires the
 `snappy <https://rubygems.org/gems/snappy>`_ library to be installed.

lib/mongo/client.rb

@@ -92,6 +92,8 @@ class Client
       :server_api,
       :server_selection_timeout,
       :socket_timeout,
+      :srv_max_hosts,
+      :srv_service_name,
       :ssl,
       :ssl_ca_cert,
       :ssl_ca_cert_object,
@@ -340,6 +342,13 @@ def hash
     #   for selecting a server for an operation.
     # @option options [ Float ] :socket_timeout The timeout, in seconds, to
     #   execute operations on a socket.
+    # @option options [ Integer ] :srv_max_hosts The maximum number of mongoses
+    #   that the driver will communicate with for sharded topologies. If this
+    #   option is 0, then there will be no maximum number of mongoses. If the
+    #   given URI resolves to more hosts than ``:srv_max_hosts``, the client
+    #   will ramdomly choose an ``:srv_max_hosts`` sized subset of hosts.
+    # @option options [ String ] :srv_service_name The service name to use in
+    #   the SRV DNS query.
     # @option options [ true, false ] :ssl Whether to use TLS.
     # @option options [ String ] :ssl_ca_cert The file containing concatenated
     #   certificate authority certificates used to validate certs passed from the
@@ -532,7 +541,7 @@ def initialize(addresses_or_uri, options = nil)
       end
 =end
       @options.freeze
-      validate_options!(addresses)
+      validate_options!(addresses, is_srv: uri.is_a?(URI::SRVProtocol))
       validate_authentication_options!
 
       database_options = @options.dup
@@ -1226,6 +1235,12 @@ def validate_new_options!(opts)
             end
 
             _options[key] = compressors unless compressors.empty?
+          elsif key == :srv_max_hosts
+            if v && (!v.is_a?(Integer) || v < 0)
+              log_warn("#{v} is not a valid integer for srv_max_hosts")
+            else
+              _options[key] = v
+            end
           else
             _options[key] = v
           end
@@ -1239,7 +1254,7 @@ def validate_new_options!(opts)
     # Validates all options after they are set on the client.
     # This method is intended to catch combinations of options which are
     # not allowed.
-    def validate_options!(addresses = nil)
+    def validate_options!(addresses = nil, is_srv: nil)
       if options[:write] && options[:write_concern] && options[:write] != options[:write_concern]
         raise ArgumentError, "If :write and :write_concern are both given, they must be identical: #{options.inspect}"
       end
@@ -1349,6 +1364,26 @@ def validate_options!(addresses = nil)
           end
         end
       end
+
+      if options[:srv_max_hosts] && options[:srv_max_hosts] > 0
+        if options[:replica_set]
+          raise ArgumentError, ":srv_max_hosts > 0 cannot be used with :replica_set option"
+        end
+
+        if options[:load_balanced]
+          raise ArgumentError, ":srv_max_hosts > 0 cannot be used with :load_balanced=true"
+        end
+      end
+
+      unless is_srv.nil? || is_srv
+        if options[:srv_max_hosts]
+          raise ArgumentError, ":srv_max_hosts cannot be used on non-SRV URI"
+        end
+
+        if options[:srv_service_name]
+          raise ArgumentError, ":srv_service_name cannot be used on non-SRV URI"
+        end
+      end
     end
 
     # Validates all authentication-related options after they are set on the client

lib/mongo/srv/resolver.rb

@@ -29,6 +29,16 @@ class Resolver
       #   before querying SRV records.
       RECORD_PREFIX = '_mongodb._tcp.'.freeze
 
+      # Generates the record prefix with a custom SRV service name if it is
+      # provided.
+      #
+      # @option srv_service_name [ String | nil ] The SRV service name to use
+      #   in the record prefix.
+      # @return [ String ] The generated record prefix.
+      def record_prefix(srv_service_name=nil)
+        return srv_service_name ? "_#{srv_service_name}._tcp." : RECORD_PREFIX
+      end
+
       # Creates a new Resolver.
       #
       # @option opts [ Float ] :timeout The timeout, in seconds, to use for
@@ -51,13 +61,18 @@ def timeout
         options[:timeout] || Monitor::DEFAULT_TIMEOUT
       end
 
-      # Obtains all of the SRV records for a given hostname.
+      # Obtains all of the SRV records for a given hostname. If a srv_max_hosts
+      # is specified and it is greater than 0, return maximum srv_max_hosts records.
       #
       # In the event that a record with a mismatched domain is found or no
       # records are found, if the :raise_on_invalid option is true,
       # an exception will be raised, otherwise a warning will be logged.
       #
       # @param [ String ] hostname The hostname whose records should be obtained.
+      # @param [ String | nil ] srv_service_name The SRV service name for the DNS query.
+      #   If nil, 'mongodb' is used.
+      # @param [ Integer | nil ] srv_max_hosts The maximum number of records to return.
+      #   If this value is nil, return all of the records.
       #
       # @raise [ Mongo::Error::MismatchedDomain ] If the :raise_in_invalid
       #   Resolver option is true and a record with a domain name that does
@@ -66,8 +81,8 @@ def timeout
       #   option is true and no records are found.
       #
       # @return [ Mongo::Srv::Result ] SRV lookup result.
-      def get_records(hostname)
-        query_name = RECORD_PREFIX + hostname
+      def get_records(hostname, srv_service_name=nil, srv_max_hosts=nil)
+        query_name = record_prefix(srv_service_name) + hostname
         resources = @resolver.getresources(query_name, Resolv::DNS::Resource::IN::SRV)
 
         # Collect all of the records into a Result object, raising an error
@@ -97,6 +112,12 @@ def get_records(hostname)
           end
         end
 
+        # if srv_max_hosts is in [1, #addresses)
+        if (1...result.address_strs.length).include? srv_max_hosts
+          sampled_records = resources.shuffle.first(srv_max_hosts)
+          result = Srv::Result.new(hostname)
+          sampled_records.each { |record| result.add_record(record) }
+        end
         result
       end
 

lib/mongo/uri.rb

@@ -538,6 +538,26 @@ def validate_uri_options!
           raise_invalid_error_no_fmt!("loadBalanced=true cannot be used with replicaSet option")
         end
       end
+
+      unless self.is_a?(URI::SRVProtocol)
+        if uri_options[:srv_max_hosts]
+          raise_invalid_error_no_fmt!("srvMaxHosts cannot be used on non-SRV URI")
+        end
+
+        if uri_options[:srv_service_name]
+          raise_invalid_error_no_fmt!("srvServiceName cannot be used on non-SRV URI")
+        end
+      end
+
+      if uri_options[:srv_max_hosts] && uri_options[:srv_max_hosts] > 0
+        if uri_options[:replica_set]
+          raise_invalid_error_no_fmt!("srvMaxHosts > 0 cannot be used with replicaSet option")
+        end
+
+        if options[:load_balanced]
+          raise_invalid_error_no_fmt!("srvMaxHosts > 0 cannot be used with loadBalanced=true")
+        end
+      end
     end
   end
 end

lib/mongo/uri/options_mapper.rb

@@ -269,6 +269,8 @@ def self.uri_option(uri_key, name, **extra)
       uri_option 'directConnection', :direct_connection, type: :bool
       uri_option 'connect', :connect, type: :symbol
       uri_option 'loadBalanced', :load_balanced, type: :bool
+      uri_option 'srvMaxHosts', :srv_max_hosts, type: :integer
+      uri_option 'srvServiceName', :srv_service_name
 
       # Auth Options
       uri_option 'authSource', :auth_source

lib/mongo/uri/srv_protocol.rb

@@ -147,7 +147,7 @@ def parse!(remaining)
         validate_srv_hostname(hostname)
         @query_hostname = hostname
 
-        @srv_result = resolver.get_records(hostname)
+        @srv_result = resolver.get_records(hostname, uri_options[:srv_service_name], uri_options[:srv_max_hosts])
         if srv_result.empty?
           raise Error::NoSRVRecords.new(NO_SRV_RECORDS % hostname)
         end