PYTHON-2832 Provide options to limit number of mongos servers used in connecting to sharded clusters (#754)

Author: juliusgeo

doc/changelog.rst

@@ -179,6 +179,11 @@ Notable improvements
 
 - Enhanced connection pooling to create connections more efficiently and
   avoid connection storms.
+- :class:`~pymongo.mongo_client.MongoClient` now accepts a URI and keyword
+  argument `srvMaxHosts` that limits the number of mongos-like hosts a client
+  will connect to. More specifically, when a mongodb+srv:// connection string
+  resolves to more than `srvMaxHosts` number of hosts, the client will randomly
+  choose a `srvMaxHosts` sized subset of hosts.
 
 Issues Resolved
 ...............

pymongo/common.py

@@ -629,7 +629,8 @@ def validate_tzinfo(dummy, value):
     'w': validate_non_negative_int_or_basestring,
     'wtimeoutms': validate_non_negative_integer,
     'zlibcompressionlevel': validate_zlib_compression_level,
-    'srvservicename': validate_string
+    'srvservicename': validate_string,
+    'srvmaxhosts': validate_non_negative_integer
 }
 
 # Dictionary where keys are the names of URI options specific to pymongo,

pymongo/mongo_client.py

@@ -652,8 +652,8 @@ def __init__(
         dbase = None
         opts = common._CaseInsensitiveDictionary()
         fqdn = None
-        srv_service_name = keyword_opts.get("srvservicename", None)
-
+        srv_service_name = keyword_opts.get("srvservicename")
+        srv_max_hosts = keyword_opts.get("srvmaxhosts")
         if len([h for h in host if "/" in h]) > 1:
             raise ConfigurationError("host must not contain multiple MongoDB "
                                      "URIs")
@@ -669,7 +669,9 @@ def __init__(
                         keyword_opts.cased_key("connecttimeoutms"), timeout)
                 res = uri_parser.parse_uri(
                     entity, port, validate=True, warn=True, normalize=False,
-                    connect_timeout=timeout, srv_service_name=srv_service_name)
+                    connect_timeout=timeout,
+                    srv_service_name=srv_service_name,
+                    srv_max_hosts=srv_max_hosts)
                 seeds.update(res["nodelist"])
                 username = res["username"] or username
                 password = res["password"] or password
@@ -703,6 +705,7 @@ def __init__(
         if srv_service_name is None:
             srv_service_name = opts.get("srvServiceName", common.SRV_SERVICE_NAME)
 
+        srv_max_hosts = srv_max_hosts or opts.get("srvmaxhosts")
         # Handle security-option conflicts in combined options.
         opts = _handle_security_options(opts)
         # Normalize combined options.
@@ -745,6 +748,7 @@ def __init__(
             srv_service_name=srv_service_name,
             direct_connection=options.direct_connection,
             load_balanced=options.load_balanced,
+            srv_max_hosts=srv_max_hosts
         )
 
         self._topology = Topology(self._topology_settings)

pymongo/settings.py

@@ -41,7 +41,8 @@ def __init__(self,
                  fqdn=None,
                  srv_service_name=common.SRV_SERVICE_NAME,
                  direct_connection=False,
-                 load_balanced=None):
+                 load_balanced=None,
+                 srv_max_hosts=0):
         """Represent MongoClient's configuration.
 
         Take a list of (host, port) pairs and optional replica set name.
@@ -63,7 +64,7 @@ def __init__(self,
         self._fqdn = fqdn
         self._srv_service_name = srv_service_name
         self._heartbeat_frequency = heartbeat_frequency
-
+        self._srv_max_hosts = srv_max_hosts or 0
         self._direct = direct_connection
         self._load_balanced = load_balanced
 

pymongo/srv_resolver.py

@@ -15,6 +15,7 @@
 """Support for resolving hosts and options from mongodb+srv:// URIs."""
 
 import ipaddress
+import random
 
 try:
     from dns import resolver
@@ -25,7 +26,6 @@
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.errors import ConfigurationError
 
-
 # dnspython can return bytes or str from various parts
 # of its API depending on version. We always want str.
 def maybe_decode(text):
@@ -48,11 +48,11 @@ def _resolve(*args, **kwargs):
 
 class _SrvResolver(object):
     def __init__(self, fqdn,
-                 connect_timeout, srv_service_name):
+                 connect_timeout, srv_service_name, srv_max_hosts=0):
         self.__fqdn = fqdn
         self.__srv = srv_service_name
         self.__connect_timeout = connect_timeout or CONNECT_TIMEOUT
-
+        self.__srv_max_hosts = srv_max_hosts or 0
         # Validate the fully qualified domain name.
         try:
             ipaddress.ip_address(fqdn)
@@ -111,7 +111,8 @@ def _get_srv_response_and_hosts(self, encapsulate_errors):
                 raise ConfigurationError("Invalid SRV host: %s" % (node[0],))
             if self.__plist != nlist:
                 raise ConfigurationError("Invalid SRV host: %s" % (node[0],))
-
+        if self.__srv_max_hosts:
+            nodes = random.sample(nodes, min(self.__srv_max_hosts, len(nodes)))
         return results, nodes
 
     def get_hosts(self):

pymongo/topology_description.py

@@ -15,6 +15,7 @@
 """Represent a deployment of MongoDB servers."""
 
 from collections import namedtuple
+from random import sample
 
 from pymongo import common
 from pymongo.errors import ConfigurationError
@@ -223,6 +224,10 @@ def common_wire_version(self):
     def heartbeat_frequency(self):
         return self._topology_settings.heartbeat_frequency
 
+    @property
+    def srv_max_hosts(self):
+        return self._topology_settings._srv_max_hosts
+
     def apply_selector(self, selector, address, custom_selector=None):
 
         def apply_local_threshold(selection):
@@ -446,16 +451,23 @@ def _updated_topology_description_srv_polling(topology_description, seedlist):
     if set(sds.keys()) == set(seedlist):
         return topology_description
 
-    # Add SDs corresponding to servers recently added to the SRV record.
-    for address in seedlist:
-        if address not in sds:
-            sds[address] = ServerDescription(address)
 
     # Remove SDs corresponding to servers no longer part of the SRV record.
     for address in list(sds.keys()):
         if address not in seedlist:
             sds.pop(address)
 
+    if topology_description.srv_max_hosts != 0:
+        new_hosts = set(seedlist) - set(sds.keys())
+        n_to_add = topology_description.srv_max_hosts - len(sds)
+        if n_to_add > 0:
+            seedlist = sample(new_hosts, min(n_to_add, len(new_hosts)))
+        else:
+            seedlist = []
+    # Add SDs corresponding to servers recently added to the SRV record.
+    for address in seedlist:
+        if address not in sds:
+            sds[address] = ServerDescription(address)
     return TopologyDescription(
         topology_description.topology_type,
         sds,

pymongo/uri_parser.py

@@ -393,7 +393,8 @@ def _check_options(nodes, options):
 
 
 def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
-              normalize=True, connect_timeout=None, srv_service_name=None):
+              normalize=True, connect_timeout=None, srv_service_name=None,
+              srv_max_hosts=None):
     """Parse and validate a MongoDB URI.
 
     Returns a dict of the form::
@@ -494,10 +495,8 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
 
         if opts:
             options.update(split_options(opts, validate, warn, normalize))
-
     if srv_service_name is None:
         srv_service_name = options.get("srvServiceName", SRV_SERVICE_NAME)
-
     if '@' in host_part:
         userinfo, _, hosts = host_part.rpartition('@')
         user, passwd = parse_userinfo(userinfo)
@@ -510,7 +509,7 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
 
     hosts = unquote_plus(hosts)
     fqdn = None
-
+    srv_max_hosts = srv_max_hosts or options.get("srvMaxHosts")
     if is_srv:
         if options.get('directConnection'):
             raise ConfigurationError(
@@ -529,7 +528,8 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
         # Use the connection timeout. connectTimeoutMS passed as a keyword
         # argument overrides the same option passed in the connection string.
         connect_timeout = connect_timeout or options.get("connectTimeoutMS")
-        dns_resolver = _SrvResolver(fqdn, connect_timeout, srv_service_name)
+        dns_resolver = _SrvResolver(fqdn, connect_timeout, srv_service_name,
+                                    srv_max_hosts)
         nodes = dns_resolver.get_hosts()
         dns_options = dns_resolver.get_options()
         if dns_options:
@@ -542,11 +542,19 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
             for opt, val in parsed_dns_options.items():
                 if opt not in options:
                     options[opt] = val
+        if options.get("loadBalanced") and srv_max_hosts:
+            raise InvalidURI(
+                "You cannot specify loadBalanced with srvMaxHosts")
+        if options.get("replicaSet") and srv_max_hosts:
+            raise InvalidURI("You cannot specify replicaSet with srvMaxHosts")
         if "tls" not in options and "ssl" not in options:
             options["tls"] = True if validate else 'true'
     elif not is_srv and options.get("srvServiceName") is not None:
         raise ConfigurationError("The srvServiceName option is only allowed "
                                  "with 'mongodb+srv://' URIs")
+    elif not is_srv and srv_max_hosts:
+        raise ConfigurationError("The srvMaxHosts option is only allowed "
+                                 "with 'mongodb+srv://' URIs")
     else:
         nodes = split_hosts(hosts, default_port=default_port)
 

test/srv_seedlist/load-balanced/srvMaxHosts-conflicts_with_loadBalanced-true-txt.json

@@ -0,0 +1,7 @@
+{
+  "uri": "mongodb+srv://test20.test.build.10gen.cc/?srvMaxHosts=1",
+  "seeds": [],
+  "hosts": [],
+  "error": true,
+  "comment": "Should fail because positive integer for srvMaxHosts conflicts with loadBalanced=true (TXT)"
+}

test/srv_seedlist/load-balanced/srvMaxHosts-conflicts_with_loadBalanced-true.json

@@ -0,0 +1,7 @@
+{
+  "uri": "mongodb+srv://test3.test.build.10gen.cc/?loadBalanced=true&srvMaxHosts=1",
+  "seeds": [],
+  "hosts": [],
+  "error": true,
+  "comment": "Should fail because positive integer for srvMaxHosts conflicts with loadBalanced=true"
+}

test/srv_seedlist/load-balanced/srvMaxHosts-zero-txt.json

@@ -0,0 +1,14 @@
+{
+  "uri": "mongodb+srv://test20.test.build.10gen.cc/?srvMaxHosts=0",
+  "seeds": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "hosts": [
+    "localhost.test.build.10gen.cc:27017"
+  ],
+  "options": {
+    "loadBalanced": true,
+    "srvMaxHosts": 0,
+    "ssl": true
+  }
+}