PYTHON-1163 update topology from handshake reply

Author: ajdavis

pymongo/pool.py

@@ -217,14 +217,15 @@ class PoolOptions(object):
                  '__connect_timeout', '__socket_timeout',
                  '__wait_queue_timeout', '__wait_queue_multiple',
                  '__ssl_context', '__ssl_match_hostname', '__socket_keepalive',
-                 '__event_listeners', '__appname', '__metadata')
+                 '__event_listeners', '__appname', '__metadata',
+                 '__handshake_callback')
 
     def __init__(self, max_pool_size=100, min_pool_size=0,
                  max_idle_time_ms=None, connect_timeout=None,
                  socket_timeout=None, wait_queue_timeout=None,
                  wait_queue_multiple=None, ssl_context=None,
                  ssl_match_hostname=True, socket_keepalive=False,
-                 event_listeners=None, appname=None):
+                 event_listeners=None, appname=None, handshake_callback=None):
 
         self.__max_pool_size = max_pool_size
         self.__min_pool_size = min_pool_size
@@ -242,6 +243,27 @@ def __init__(self, max_pool_size=100, min_pool_size=0,
         if appname:
             self.__metadata['application'] = {'name': appname}
 
+        self.__handshake_callback = handshake_callback
+
+    def with_options(self, **kwargs):
+        options = {
+            'max_pool_size': self.max_pool_size,
+            'min_pool_size': self.min_pool_size,
+            'max_idle_time_ms': self.max_idle_time_ms,
+            'connect_timeout': self.connect_timeout,
+            'socket_timeout': self.socket_timeout,
+            'wait_queue_timeout': self.wait_queue_timeout,
+            'wait_queue_multiple': self.wait_queue_multiple,
+            'ssl_context': self.ssl_context,
+            'ssl_match_hostname': self.ssl_match_hostname,
+            'socket_keepalive': self.socket_keepalive,
+            'event_listeners': self.event_listeners,
+            'appname': self.appname,
+            'handshake_callback': self.handshake_callback}
+
+        options.update(kwargs)
+        return PoolOptions(**options)
+
     @property
     def max_pool_size(self):
         """The maximum allowable number of concurrent connections to each
@@ -335,6 +357,11 @@ def metadata(self):
         """
         return self.__metadata.copy()
 
+    @property
+    def handshake_callback(self):
+        """Receives an ismaster reply and updates the topology."""
+        return self.__handshake_callback
+
 
 class SocketInfo(object):
     """Store a socket with some metadata.
@@ -746,6 +773,8 @@ def connect(self):
                     ('ismaster', 1),
                     ('client', self.opts.metadata)
                 ])
+
+                start = _time()
                 ismaster = IsMaster(
                     command(sock,
                             'admin',
@@ -754,13 +783,20 @@ def connect(self):
                             False,
                             ReadPreference.PRIMARY,
                             DEFAULT_CODEC_OPTIONS))
+
+                # Can raise ConnectionFailure.
+                self._handshake_callback(ismaster, _time() - start)
             else:
                 ismaster = None
             return SocketInfo(sock, self, ismaster, self.address)
         except socket.error as error:
             if sock is not None:
                 sock.close()
             _raise_connection_failure(self.address, error)
+        except:
+            if sock is not None:
+                sock.close()
+            raise
 
     @contextlib.contextmanager
     def get_socket(self, all_credentials, checkout=False):
@@ -889,6 +925,14 @@ def _check(self, sock_info):
         else:
             return self.connect()
 
+    def _handshake_callback(self, ismaster, round_trip_time):
+        callback = self.opts.handshake_callback
+        if callback:
+            kept = callback(self.address, ismaster, round_trip_time)
+            if not kept:
+                _raise_connection_failure(
+                    self.address, "server removed from topology")
+
     def _raise_wait_queue_timeout(self):
         raise ConnectionFailure(
             'Timed out waiting for socket from pool with max_size %r and'

pymongo/server_description.py

@@ -38,21 +38,25 @@ class ServerDescription(object):
       - `ismaster`: Optional IsMaster instance
       - `round_trip_time`: Optional float
       - `error`: Optional, the last error attempting to connect to the server
+      - `from_handshake`: Optional, whether this is from expanding a connection
+        pool, rather than from background monitoring.
     """
 
     __slots__ = (
         '_address', '_server_type', '_all_hosts', '_tags', '_replica_set_name',
         '_primary', '_max_bson_size', '_max_message_size',
         '_max_write_batch_size', '_min_wire_version', '_max_wire_version',
         '_round_trip_time', '_me', '_is_writable', '_is_readable', '_error',
-        '_set_version', '_election_id', '_last_write_date', '_last_update_time')
+        '_set_version', '_election_id', '_last_write_date', '_last_update_time',
+        '_from_handshake')
 
     def __init__(
             self,
             address,
             ismaster=None,
             round_trip_time=None,
-            error=None):
+            error=None,
+            from_handshake=False):
         self._address = address
         if not ismaster:
             ismaster = IsMaster({})
@@ -75,6 +79,7 @@ def __init__(
         self._me = ismaster.me
         self._last_update_time = _time()
         self._error = error
+        self._from_handshake = from_handshake  # For tests.
 
         if ismaster.last_write_date:
             # Convert from datetime to seconds.

pymongo/topology.py

@@ -28,13 +28,13 @@
 
 from pymongo import common
 from pymongo import periodic_executor
-from pymongo.pool import PoolOptions
 from pymongo.topology_description import (updated_topology_description,
                                           TOPOLOGY_TYPE,
                                           TopologyDescription)
 from pymongo.errors import ServerSelectionTimeoutError
 from pymongo.monotonic import time as _time
 from pymongo.server import Server
+from pymongo.server_description import ServerDescription
 from pymongo.server_selectors import (any_server_selector,
                                       arbiter_server_selector,
                                       secondary_server_selector,
@@ -237,7 +237,10 @@ def select_server_by_address(self, address,
                                   address)
 
     def on_change(self, server_description):
-        """Process a new ServerDescription after an ismaster call completes."""
+        """Process a new ServerDescription after an ismaster call completes.
+
+        Returns False if the server was removed from the topology.
+        """
         # We do no I/O holding the lock.
         with self._lock:
             # Any monitored server was definitely in the topology description
@@ -266,6 +269,9 @@ def on_change(self, server_description):
 
                 # Wake waiters in select_servers().
                 self._condition.notify_all()
+                return self._description.has_server(server_description.address)
+            else:
+                return False
 
     def get_server_by_address(self, address):
         """Get a Server or None.
@@ -337,9 +343,13 @@ def reset_server_and_request_check(self, address):
 
     def update_pool(self):
         # Remove any stale sockets and add new sockets if pool is too small.
+        # Avoid locking around network I/O, or deadlocking when a new connection
+        # opens and calls Topology.on_change() with the ismaster reply.
         with self._lock:
-            for server in self._servers.values():
-                server._pool.remove_stale_sockets()
+            pools = [server._pool for server in self._servers.values()]
+
+        for pool in pools:
+            pool.remove_stale_sockets()
 
     def close(self):
         """Clear pools and terminate monitors. Topology reopens on demand."""
@@ -448,22 +458,36 @@ def _update_servers(self):
                 self._servers.pop(address)
 
     def _create_pool_for_server(self, address):
-        return self._settings.pool_class(address, self._settings.pool_options)
+        # Server Discovery And Monitoring Spec: When a client calls ismaster
+        # to handshake a new connection for application operations, use the
+        # ismaster reply to update the topology.
+        ref = weakref.proxy(self)
 
-    def _create_pool_for_monitor(self, address):
-        options = self._settings.pool_options
+        def handshake_callback(address, ismaster, round_trip_time):
+            sd = ServerDescription(address, ismaster, round_trip_time,
+                                   from_handshake=True)
+
+            try:
+                # Return False if server was removed from topology.
+                return ref.on_change(sd)
+            except ReferenceError:
+                return True
+
+        server_pool_options = self._settings.pool_options.with_options(
+            handshake_callback=handshake_callback)
 
+        return self._settings.pool_class(address, server_pool_options)
+
+    def _create_pool_for_monitor(self, address):
         # According to the Server Discovery And Monitoring Spec, monitors use
         # connect_timeout for both connect_timeout and socket_timeout. The
         # pool only has one socket so maxPoolSize and so on aren't needed.
-        monitor_pool_options = PoolOptions(
-            connect_timeout=options.connect_timeout,
-            socket_timeout=options.connect_timeout,
-            ssl_context=options.ssl_context,
-            ssl_match_hostname=options.ssl_match_hostname,
+        opts = self._settings.pool_options
+        monitor_pool_options = opts.with_options(
+            socket_timeout=opts.connect_timeout,
             socket_keepalive=True,
-            event_listeners=options.event_listeners,
-            appname=options.appname)
+            max_pool_size=None,
+            min_pool_size=None)
 
         return self._settings.pool_class(address, monitor_pool_options,
                                          handshake=False)

test/pymongo_mocks.py

@@ -58,6 +58,10 @@ def get_socket(self, all_credentials, checkout=False):
             sock_info.mock_port = self.mock_port
             yield sock_info
 
+    def _handshake_callback(self, ismaster, round_trip_time):
+        # Don't mock how PyMongo updates topology from ismaster reply.
+        return True
+
 
 class MockMonitor(Monitor):
     def __init__(

test/test_topology.py

@@ -22,6 +22,7 @@
 
 from bson.py3compat import imap
 from pymongo import common
+from pymongo import monitoring
 from pymongo.read_preferences import ReadPreference, Secondary
 from pymongo.server_type import SERVER_TYPE
 from pymongo.topology import Topology
@@ -37,7 +38,7 @@
                                       writable_server_selector)
 from pymongo.settings import TopologySettings
 from test import client_knobs, unittest
-from test.utils import wait_until
+from test.utils import rs_or_single_client, wait_until
 
 
 class MockSocketInfo(object):
@@ -297,6 +298,23 @@ def new_average():
             if tries > 10:
                 self.fail("Didn't ever calculate correct new average")
 
+    def test_update_from_handshake(self):
+        class ServerHandshakes(monitoring.ServerListener, list):
+            def opened(self, e):
+                pass
+
+            def description_changed(self, e):
+                if e.new_description._from_handshake:
+                    self.append(e)
+
+            def closed(self, e):
+                pass
+
+        handshakes = ServerHandshakes()
+        client = rs_or_single_client(event_listeners=[handshakes])
+        client.admin.command('ping')
+        wait_until(lambda: handshakes, 'record handshakes')
+
 
 class TestMultiServerTopology(TopologyTest):
     def test_readable_writable(self):