PYTHON-1675 SRV polling for mongos discovery

Author: prashantmital

doc/changelog.rst

@@ -39,6 +39,8 @@ Version 3.9 adds support for MongoDB 4.2. Highlights include:
   enabled by default. See the :class:`~pymongo.mongo_client.MongoClient`
   documentation for details.
 - Support zstandard for wire protocol compression.
+- Support for periodically polling DNS SRV records to update the mongos proxy
+  list without having to change client configuration.
 
   Now that supported operations are retried automatically and transparently,
   users should consider adjusting any custom retry logic to prevent

pymongo/common.py

@@ -73,6 +73,9 @@
 # Spec requires at least 500ms between ismaster calls.
 MIN_HEARTBEAT_INTERVAL = 0.5
 
+# Spec requires at least 60s between SRV rescans.
+MIN_SRV_RESCAN_INTERVAL = 60
+
 # Default connectTimeout in seconds.
 CONNECT_TIMEOUT = 20.0
 

pymongo/mongo_client.py

@@ -585,6 +585,7 @@ def __init__(
         password = None
         dbase = None
         opts = {}
+        fqdn = None
         for entity in host:
             if "://" in entity:
                 res = uri_parser.parse_uri(
@@ -594,6 +595,7 @@ def __init__(
                 password = res["password"] or password
                 dbase = res["database"] or dbase
                 opts = res["options"]
+                fqdn = res["fqdn"]
             else:
                 seeds.update(uri_parser.split_hosts(entity, port))
         if not seeds:
@@ -673,7 +675,8 @@ def __init__(
             local_threshold_ms=options.local_threshold_ms,
             server_selection_timeout=options.server_selection_timeout,
             server_selector=options.server_selector,
-            heartbeat_frequency=options.heartbeat_frequency)
+            heartbeat_frequency=options.heartbeat_frequency,
+            fqdn=fqdn)
 
         self._topology = Topology(self._topology_settings)
         if connect:

pymongo/monitor.py

@@ -18,13 +18,42 @@
 
 from pymongo import common, periodic_executor
 from pymongo.errors import OperationFailure
-from pymongo.server_type import SERVER_TYPE
 from pymongo.monotonic import time as _time
 from pymongo.read_preferences import MovingAverage
 from pymongo.server_description import ServerDescription
+from pymongo.server_type import SERVER_TYPE
+from pymongo.srv_resolver import _SrvResolver
+
+
+class MonitorBase(object):
+    def __init__(self, *args, **kwargs):
+        """Override this method to create an executor."""
+        raise NotImplementedError
+
+    def open(self):
+        """Start monitoring, or restart after a fork.
+
+        Multiple calls have no effect.
+        """
+        self._executor.open()
+
+    def close(self):
+        """Close and stop monitoring.
+
+        open() restarts the monitor after closing.
+        """
+        self._executor.close()
+
+    def join(self, timeout=None):
+        """Wait for the monitor to stop."""
+        self._executor.join(timeout)
+
+    def request_check(self):
+        """If the monitor is sleeping, wake it soon."""
+        self._executor.wake()
 
 
-class Monitor(object):
+class Monitor(MonitorBase):
     def __init__(
             self,
             server_description,
@@ -68,31 +97,13 @@ def target():
         self_ref = weakref.ref(self, executor.close)
         self._topology = weakref.proxy(topology, executor.close)
 
-    def open(self):
-        """Start monitoring, or restart after a fork.
-
-        Multiple calls have no effect.
-        """
-        self._executor.open()
-
     def close(self):
-        """Close and stop monitoring.
-
-        open() restarts the monitor after closing.
-        """
-        self._executor.close()
+        super(Monitor, self).close()
 
         # Increment the pool_id and maybe close the socket. If the executor
         # thread has the socket checked out, it will be closed when checked in.
         self._pool.reset()
 
-    def join(self, timeout=None):
-        self._executor.join(timeout)
-
-    def request_check(self):
-        """If the monitor is sleeping, wake and check the server soon."""
-        self._executor.wake()
-
     def _run(self):
         try:
             self._server_description = self._check_with_retry()
@@ -182,3 +193,66 @@ def _check_with_socket(self, sock_info):
             self._topology.receive_cluster_time(
                 exc.details.get('$clusterTime'))
             raise
+
+
+class SrvMonitor(MonitorBase):
+    def __init__(self, topology, topology_settings):
+        """Class to poll SRV records on a background thread.
+
+        Pass a Topology and a TopologySettings.
+
+        The Topology is weakly referenced.
+        """
+        self._settings = topology_settings
+        self._fqdn = self._settings.fqdn
+
+        # We strongly reference the executor and it weakly references us via
+        # this closure. When the monitor is freed, stop the executor soon.
+        def target():
+            monitor = self_ref()
+            if monitor is None:
+                return False  # Stop the executor.
+            SrvMonitor._run(monitor)
+            return True
+
+        executor = periodic_executor.PeriodicExecutor(
+            interval=common.MIN_SRV_RESCAN_INTERVAL,
+            min_interval=self._settings.heartbeat_frequency,
+            target=target,
+            name="pymongo_srv_polling_thread")
+
+        self._executor = executor
+
+        # Avoid cycles. When self or topology is freed, stop executor soon.
+        self_ref = weakref.ref(self, executor.close)
+        self._topology = weakref.proxy(topology, executor.close)
+
+    def _run(self):
+        try:
+            self._seedlist = self._get_seedlist()
+            self._topology.on_srv_update(self._seedlist)
+        except ReferenceError:
+            # Topology was garbage-collected.
+            self.close()
+
+    def _get_seedlist(self):
+        """Poll SRV records for a seedlist.
+
+        Returns a list of ServerDescriptions.
+        """
+        try:
+            seedlist, ttl = _SrvResolver(self._fqdn).get_hosts_and_min_ttl()
+            if len(seedlist) == 0:
+                # As per the spec: this should be treated as a failure.
+                raise Exception
+        except Exception:
+            # As per the spec, upon encountering an error:
+            # - An error must not be raised
+            # - SRV records must be rescanned every heartbeatFrequencyMS
+            # - Topology must be left unchanged
+            self.request_check()
+            return self._seedlist
+        else:
+            self._executor.update_interval(
+                max(ttl, common.MIN_SRV_RESCAN_INTERVAL))
+            return seedlist

pymongo/periodic_executor.py

@@ -102,6 +102,9 @@ def wake(self):
         """Execute the target function soon."""
         self._event = True
 
+    def update_interval(self, new_interval):
+        self._interval = new_interval
+
     def __should_stop(self):
         with self._lock:
             if self._stopped:

pymongo/settings.py

@@ -20,9 +20,9 @@
 from pymongo import common, monitor, pool
 from pymongo.common import LOCAL_THRESHOLD_MS, SERVER_SELECTION_TIMEOUT
 from pymongo.errors import ConfigurationError
-from pymongo.topology_description import TOPOLOGY_TYPE
 from pymongo.pool import PoolOptions
 from pymongo.server_description import ServerDescription
+from pymongo.topology_description import TOPOLOGY_TYPE
 
 
 class TopologySettings(object):
@@ -36,7 +36,8 @@ def __init__(self,
                  local_threshold_ms=LOCAL_THRESHOLD_MS,
                  server_selection_timeout=SERVER_SELECTION_TIMEOUT,
                  heartbeat_frequency=common.HEARTBEAT_FREQUENCY,
-                 server_selector=None):
+                 server_selector=None,
+                 fqdn=None):
         """Represent MongoClient's configuration.
 
         Take a list of (host, port) pairs and optional replica set name.
@@ -55,6 +56,7 @@ def __init__(self,
         self._local_threshold_ms = local_threshold_ms
         self._server_selection_timeout = server_selection_timeout
         self._server_selector = server_selector
+        self._fqdn = fqdn
         self._heartbeat_frequency = heartbeat_frequency
         self._direct = (len(self._seeds) == 1 and not replica_set_name)
         self._topology_id = ObjectId()
@@ -100,6 +102,10 @@ def server_selector(self):
     def heartbeat_frequency(self):
         return self._heartbeat_frequency
 
+    @property
+    def fqdn(self):
+        return self._fqdn
+
     @property
     def direct(self):
         """Connect directly to a single server, or use a set of servers?

pymongo/srv_resolver.py

@@ -0,0 +1,103 @@
+# Copyright 2019-present 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.
+
+"""Support for resolving hosts and options from mongodb+srv:// URIs."""
+
+try:
+    from dns import resolver
+    _HAVE_DNSPYTHON = True
+except ImportError:
+    _HAVE_DNSPYTHON = False
+
+from bson.py3compat import PY3
+
+from pymongo.errors import ConfigurationError
+
+
+if PY3:
+    # dnspython can return bytes or str from various parts
+    # of its API depending on version. We always want str.
+    def maybe_decode(text):
+        if isinstance(text, bytes):
+            return text.decode()
+        return text
+else:
+    def maybe_decode(text):
+        return text
+
+
+class _SrvResolver(object):
+    def __init__(self, fqdn):
+        self.__fqdn = fqdn
+
+        # Validate the fully qualified domain name.
+        try:
+            self.__plist = self.__fqdn.split(".")[1:]
+        except Exception:
+            raise ConfigurationError("Invalid URI host")
+        self.__slen = len(self.__plist)
+        if self.__slen < 2:
+            raise ConfigurationError("Invalid URI host")
+
+    def get_options(self):
+        try:
+            results = resolver.query(self.__fqdn, 'TXT')
+        except (resolver.NoAnswer, resolver.NXDOMAIN):
+            # No TXT records
+            return None
+        except Exception as exc:
+            raise ConfigurationError(str(exc))
+        if len(results) > 1:
+            raise ConfigurationError('Only one TXT record is supported')
+        return (
+            b'&'.join([b''.join(res.strings) for res in results])).decode(
+            'utf-8')
+
+    def _resolve_uri(self, encapsulate_errors):
+        try:
+            results = resolver.query('_mongodb._tcp.' + self.__fqdn, 'SRV')
+        except Exception as exc:
+            if not encapsulate_errors:
+                # Raise the original error.
+                raise
+            # Else, raise all errors as ConfigurationError.
+            raise ConfigurationError(str(exc))
+        return results
+
+    def _get_srv_response_and_hosts(self, encapsulate_errors):
+        results = self._resolve_uri(encapsulate_errors)
+
+        # Construct address tuples
+        nodes = [
+            (maybe_decode(res.target.to_text(omit_final_dot=True)), res.port)
+            for res in results]
+
+        # Validate hosts
+        for node in nodes:
+            try:
+                nlist = node[0].split(".")[1:][-self.__slen:]
+            except Exception:
+                raise ConfigurationError("Invalid SRV host")
+            if self.__plist != nlist:
+                raise ConfigurationError("Invalid SRV host")
+
+        return results, nodes
+
+    def get_hosts(self):
+        _, nodes = self._get_srv_response_and_hosts(True)
+        return nodes
+
+    def get_hosts_and_min_ttl(self):
+        results, nodes = self._get_srv_response_and_hosts(False)
+        return nodes, results.rrset.ttl