PYTHON-1677 Connections survive primary stepdown

Author: prashantmital

doc/changelog.rst

@@ -65,6 +65,8 @@ Version 3.9 adds support for MongoDB 4.2. Highlights include:
   the buffer protocol.
 - Resume tokens can now be accessed from a ``ChangeStream`` cursor using the
   :attr:`~pymongo.change_stream.ChangeStream.resume_token` attribute.
+- Connections now survive primary step-down. Applications should expect less
+  socket connection turnover during replica set elections.
 
 .. _URI options specification: https://github.com/mongodb/specifications/blob/master/source/uri-options/uri-options.rst
 

pymongo/helpers.py

@@ -29,17 +29,21 @@
                             WriteConcernError,
                             WTimeoutError)
 
-# From the Server Discovery and Monitoring spec, the "not master" error codes
-# are combined with the "node is recovering" error codes.
+# From the SDAM spec, the "node is shutting down" codes.
+_SHUTDOWN_CODES = frozenset([
+    11600,  # InterruptedAtShutdown
+    91,     # ShutdownInProgress
+])
+# From the SDAM spec, the "not master" error codes are combined with the
+# "node is recovering" error codes (of which the "node is shutting down"
+# errors are a subset).
 _NOT_MASTER_CODES = frozenset([
     10107,  # NotMaster
     13435,  # NotMasterNoSlaveOk
-    11600,  # InterruptedAtShutdown
     11602,  # InterruptedDueToReplStateChange
     13436,  # NotMasterOrSecondary
     189,    # PrimarySteppedDown
-    91,     # ShutdownInProgress
-])
+]) | _SHUTDOWN_CODES
 # From the retryable writes spec.
 _RETRYABLE_ERROR_CODES = _NOT_MASTER_CODES | frozenset([
     7,     # HostNotFound

pymongo/mongo_client.py

@@ -1195,9 +1195,11 @@ def _get_topology(self):
 
     @contextlib.contextmanager
     def _get_socket(self, server, session, exhaust=False):
-        with self._reset_on_error(server.description.address, session):
-            with server.get_socket(self.__all_credentials,
-                                   checkout=exhaust) as sock_info:
+        with _MongoClientErrorHandler(
+                self, server.description.address, session) as err_handler:
+            with server.get_socket(
+                    self.__all_credentials, checkout=exhaust) as sock_info:
+                err_handler.contribute_socket(sock_info)
                 yield sock_info
 
     def _select_server(self, server_selector, session, address=None):
@@ -1289,8 +1291,10 @@ def _run_operation_with_response(self, operation, unpack_res,
             server = self._select_server(
                 operation.read_preference, operation.session, address=address)
 
-            with self._reset_on_error(server.description.address,
-                                      operation.session):
+            with _MongoClientErrorHandler(
+                    self, server.description.address,
+                    operation.session) as err_handler:
+                err_handler.contribute_socket(operation.exhaust_mgr.sock)
                 return server.run_operation_with_response(
                     operation.exhaust_mgr.sock,
                     operation,
@@ -1314,49 +1318,6 @@ def _cmd(session, server, sock_info, slave_ok):
             retryable=isinstance(operation, message._Query),
             exhaust=exhaust)
 
-    @contextlib.contextmanager
-    def _reset_on_error(self, server_address, session):
-        """On "not master" or "node is recovering" errors reset the server
-        according to the SDAM spec.
-
-        Unpin the session on transient transaction errors.
-        """
-        try:
-            try:
-                yield
-            except PyMongoError as exc:
-                if session and exc.has_error_label(
-                        "TransientTransactionError"):
-                    session._unpin_mongos()
-                raise
-        except NetworkTimeout:
-            # The socket has been closed. Don't reset the server.
-            # Server Discovery And Monitoring Spec: "When an application
-            # operation fails because of any network error besides a socket
-            # timeout...."
-            if session:
-                session._server_session.mark_dirty()
-            raise
-        except NotMasterError:
-            # "When the client sees a "not master" error it MUST replace the
-            # server's description with type Unknown. It MUST request an
-            # immediate check of the server."
-            self._reset_server_and_request_check(server_address)
-            raise
-        except ConnectionFailure:
-            # "Client MUST replace the server's description with type Unknown
-            # ... MUST NOT request an immediate check of the server."
-            self.__reset_server(server_address)
-            if session:
-                session._server_session.mark_dirty()
-            raise
-        except OperationFailure as exc:
-            if exc.code in helpers._RETRYABLE_ERROR_CODES:
-                # Do not request an immediate check since the server is likely
-                # shutting down.
-                self.__reset_server(server_address)
-            raise
-
     def _retry_with_session(self, retryable, func, session, bulk):
         """Execute an operation with at most one consecutive retries
 
@@ -1494,7 +1455,7 @@ def _retryable_write(self, retryable, func, session):
         with self._tmp_session(session) as s:
             return self._retry_with_session(retryable, func, s, None)
 
-    def __reset_server(self, address):
+    def _reset_server(self, address):
         """Clear our connection pool for a server and mark it Unknown."""
         self._topology.reset_server(address)
 
@@ -2158,3 +2119,69 @@ def __next__(self):
         raise TypeError("'MongoClient' object is not iterable")
 
     next = __next__
+
+
+class _MongoClientErrorHandler(object):
+    """Error handler for MongoClient."""
+    __slots__ = ('_client', '_server_address', '_session', '_max_wire_version')
+
+    def __init__(self, client, server_address, session):
+        self._client = client
+        self._server_address = server_address
+        self._session = session
+        self._max_wire_version = None
+
+    def contribute_socket(self, sock_info):
+        """Provide socket information to the error handler."""
+        # Currently, we only extract the max_wire_version information.
+        self._max_wire_version = sock_info.max_wire_version
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            return
+
+        if issubclass(exc_type, PyMongoError):
+            if self._session and exc_val.has_error_label(
+                    "TransientTransactionError"):
+                self._session._unpin_mongos()
+
+        if issubclass(exc_type, NetworkTimeout):
+            # The socket has been closed. Don't reset the server.
+            # Server Discovery And Monitoring Spec: "When an application
+            # operation fails because of any network error besides a socket
+            # timeout...."
+            if self._session:
+                self._session._server_session.mark_dirty()
+        elif issubclass(exc_type, NotMasterError):
+            # As per the SDAM spec if:
+            #   - the server sees a "not master" error, and
+            #   - the server is not shutting down, and
+            #   - the server version is >= 4.2, then
+            # we keep the existing connection pool, but mark the server type
+            # as Unknown and request an immediate check of the server.
+            # Otherwise, we clear the connection pool, mark the server as
+            # Unknown and request an immediate check of the server.
+            err_code = exc_val.details.get('code', -1)
+            is_shutting_down = err_code in helpers._SHUTDOWN_CODES
+            if (is_shutting_down or (self._max_wire_version is None) or
+                    (self._max_wire_version <= 7)):
+                # Clear the pool, mark server Unknown and request check.
+                self._client._reset_server_and_request_check(
+                    self._server_address)
+            else:
+                self._client._topology.mark_server_unknown_and_request_check(
+                    self._server_address)
+        elif issubclass(exc_type, ConnectionFailure):
+            # "Client MUST replace the server's description with type Unknown
+            # ... MUST NOT request an immediate check of the server."
+            self._client._reset_server(self._server_address)
+            if self._session:
+                self._session._server_session.mark_dirty()
+        elif issubclass(exc_type, OperationFailure):
+            # Do not request an immediate check since the server is likely
+            # shutting down.
+            if exc_val.code in helpers._RETRYABLE_ERROR_CODES:
+                self._client._reset_server(self._server_address)

pymongo/pool.py

@@ -651,7 +651,8 @@ def _raise_if_not_writable(self, unacknowledged):
         """
         if unacknowledged and not self.is_writable:
             # Write won't succeed, bail as if we'd received a not master error.
-            raise NotMasterError("not master")
+            raise NotMasterError("not master", {
+                "ok": 0, "errmsg": "not master", "code": 10107})
 
     def legacy_write(self, request_id, msg, max_doc_size, with_last_error):
         """Send OP_INSERT, etc., optionally returning response as a dict.
@@ -768,6 +769,9 @@ def send_cluster_time(self, command, session, client):
     def update_last_checkin_time(self):
         self.last_checkin_time = _time()
 
+    def update_is_writable(self, is_writable):
+        self.is_writable = is_writable
+
     def idle_time_seconds(self):
         """Seconds since this socket was last checked into its pool."""
         return _time() - self.last_checkin_time
@@ -958,6 +962,8 @@ def __init__(self, address, options, handshake=True):
         # Monotonically increasing connection ID required for CMAP Events.
         self.next_connection_id = 1
         self.closed = False
+        # Track whether the sockets in this pool are writeable or not.
+        self.is_writable = None
 
         # Keep track of resets, so we notice sockets created before the most
         # recent reset and close them.
@@ -1012,6 +1018,15 @@ def _reset(self, close):
             for sock_info in sockets:
                 sock_info.close_socket(ConnectionClosedReason.STALE)
 
+    def update_is_writable(self, is_writable):
+        """Updates the is_writable attribute on all sockets currently in the
+        Pool.
+        """
+        self.is_writable = is_writable
+        with self.lock:
+            for socket in self.sockets:
+                socket.update_is_writable(self.is_writable)
+
     def reset(self):
         self._reset(close=False)
 
@@ -1075,6 +1090,7 @@ def connect(self):
         sock_info = SocketInfo(sock, self, self.address, conn_id)
         if self.handshake:
             sock_info.ismaster(self.opts.metadata, None)
+            self.is_writable = sock_info.is_writable
 
         return sock_info
 
@@ -1194,6 +1210,7 @@ def return_socket(self, sock_info, publish_checkin=True):
                 sock_info.close_socket(ConnectionClosedReason.STALE)
             elif not sock_info.closed:
                 sock_info.update_last_checkin_time()
+                sock_info.update_is_writable(self.is_writable)
                 with self.lock:
                     self.sockets.appendleft(sock_info)
 

pymongo/topology.py

@@ -409,12 +409,18 @@ def reset_server(self, address):
         Do *not* request an immediate check.
         """
         with self._lock:
-            self._reset_server(address)
+            self._reset_server(address, reset_pool=True)
 
     def reset_server_and_request_check(self, address):
         """Clear our pool for a server, mark it Unknown, and check it soon."""
         with self._lock:
-            self._reset_server(address)
+            self._reset_server(address, reset_pool=True)
+            self._request_check(address)
+
+    def mark_server_unknown_and_request_check(self, address):
+        """Mark a server Unknown, and check it soon."""
+        with self._lock:
+            self._reset_server(address, reset_pool=False)
             self._request_check(address)
 
     def update_pool(self):
@@ -523,16 +529,17 @@ def _ensure_opened(self):
         for server in itervalues(self._servers):
             server.open()
 
-    def _reset_server(self, address):
-        """Clear our pool for a server and mark it Unknown.
+    def _reset_server(self, address, reset_pool):
+        """Mark a server Unknown and optionally reset it's pool.
 
         Hold the lock when calling this. Does *not* request an immediate check.
         """
         server = self._servers.get(address)
 
         # "server" is None if another thread removed it from the topology.
         if server:
-            server.reset()
+            if reset_pool:
+                server.reset()
 
             # Mark this server Unknown.
             self._description = self._description.reset_server(address)
@@ -578,7 +585,14 @@ def _update_servers(self):
                 self._servers[address] = server
                 server.open()
             else:
+                # Cache old is_writable value.
+                was_writable = self._servers[address].description.is_writable
+                # Update server description.
                 self._servers[address].description = sd
+                # Update is_writable value of the pool, if it changed.
+                if was_writable != sd.is_writable:
+                    self._servers[address].pool.update_is_writable(
+                        sd.is_writable)
 
         for address, server in list(self._servers.items()):
             if not self._description.has_server(address):

test/__init__.py

@@ -159,7 +159,6 @@ def __init__(self):
         """Create a client and grab essential information from the server."""
         self.connection_attempts = []
         self.connected = False
-        self.ismaster = {}
         self.w = None
         self.nodes = set()
         self.replica_set_name = None
@@ -184,6 +183,10 @@ def __init__(self):
         if COMPRESSORS:
             self.default_client_options["compressors"] = COMPRESSORS
 
+    @property
+    def ismaster(self):
+        return self.client.admin.command('isMaster')
+
     def _connect(self, host, port, **kwargs):
         # Jython takes a long time to connect.
         if sys.platform.startswith('java'):
@@ -253,7 +256,7 @@ def _init_client(self):
                 self.cmd_line = self.client.admin.command('getCmdLineOpts')
 
             self.server_status = self.client.admin.command('serverStatus')
-            self.ismaster = ismaster = self.client.admin.command('isMaster')
+            ismaster = self.ismaster
             self.sessions_enabled = 'logicalSessionTimeoutMinutes' in ismaster
 
             if 'setName' in ismaster:
@@ -276,18 +279,17 @@ def _init_client(self):
                         **self.default_client_options)
 
                 # Get the authoritative ismaster result from the primary.
-                self.ismaster = self.client.admin.command('ismaster')
+                ismaster = self.ismaster
                 nodes = [partition_node(node.lower())
-                         for node in self.ismaster.get('hosts', [])]
+                         for node in ismaster.get('hosts', [])]
                 nodes.extend([partition_node(node.lower())
-                              for node in self.ismaster.get('passives', [])])
+                              for node in ismaster.get('passives', [])])
                 nodes.extend([partition_node(node.lower())
-                              for node in self.ismaster.get('arbiters', [])])
+                              for node in ismaster.get('arbiters', [])])
                 self.nodes = set(nodes)
             else:
-                self.ismaster = ismaster
                 self.nodes = set([(host, port)])
-            self.w = len(self.ismaster.get("hosts", [])) or 1
+            self.w = len(ismaster.get("hosts", [])) or 1
             self.version = Version.from_client(self.client)
 
             if 'enableTestCommands=1' in self.cmd_line['argv']: