PYTHON-1438 Mark a server unknown when connection handshake fails with a network timeout error (#461)

Author: ShaneHarvey

pymongo/mongo_client.py

@@ -2258,7 +2258,7 @@ def _add_retryable_write_error(exc, max_wire_version):
 class _MongoClientErrorHandler(object):
     """Handle errors raised when executing an operation."""
     __slots__ = ('client', 'server_address', 'session', 'max_wire_version',
-                 'sock_generation')
+                 'sock_generation', 'completed_handshake')
 
     def __init__(self, client, server, session):
         self.client = client
@@ -2270,11 +2270,13 @@ def __init__(self, client, server, session):
         # completes then the error's generation number is the generation
         # of the pool at the time the connection attempt was started."
         self.sock_generation = server.pool.generation
+        self.completed_handshake = False
 
     def contribute_socket(self, sock_info):
         """Provide socket information to the error handler."""
         self.max_wire_version = sock_info.max_wire_version
         self.sock_generation = sock_info.generation
+        self.completed_handshake = True
 
     def __enter__(self):
         return self
@@ -2295,5 +2297,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     self.session._unpin_mongos()
 
         err_ctx = _ErrorContext(
-            exc_val, self.max_wire_version, self.sock_generation)
+            exc_val, self.max_wire_version, self.sock_generation,
+            self.completed_handshake)
         self.client._topology.handle_error(self.server_address, err_ctx)

pymongo/topology.py

@@ -567,7 +567,8 @@ def _handle_error(self, address, err_ctx):
         server = self._servers[address]
         error = err_ctx.error
         exc_type = type(error)
-        if issubclass(exc_type, NetworkTimeout):
+        if (issubclass(exc_type, NetworkTimeout) and
+                err_ctx.completed_handshake):
             # 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
@@ -750,10 +751,12 @@ def __repr__(self):
 
 class _ErrorContext(object):
     """An error with context for SDAM error handling."""
-    def __init__(self, error, max_wire_version, sock_generation):
+    def __init__(self, error, max_wire_version, sock_generation,
+                 completed_handshake):
         self.error = error
         self.max_wire_version = max_wire_version
         self.sock_generation = sock_generation
+        self.completed_handshake = completed_handshake
 
 
 def _is_stale_error_topology_version(current_tv, error_tv):

test/test_discovery_and_monitoring.py

@@ -119,11 +119,16 @@ def got_app_error(topology, app_error):
             raise AssertionError('unknown error type: %s' % (error_type,))
         assert False
     except (AutoReconnect, NotMasterError, OperationFailure) as e:
-        if when == 'beforeHandshakeCompletes' and error_type == 'timeout':
-            raise unittest.SkipTest('PYTHON-2211')
+        if when == 'beforeHandshakeCompletes':
+            completed_handshake = False
+        elif when == 'afterHandshakeCompletes':
+            completed_handshake = True
+        else:
+            assert False, 'Unknown when field %s' % (when,)
 
         topology.handle_error(
-            server_address, _ErrorContext(e, max_wire_version, generation))
+            server_address, _ErrorContext(e, max_wire_version, generation,
+                                          completed_handshake))
 
 
 def get_type(topology, hostname):

test/test_topology.py

@@ -419,7 +419,7 @@ def test_handle_error(self):
             'setName': 'rs',
             'hosts': ['a', 'b']})
 
-        errctx = _ErrorContext(AutoReconnect('mock'), 0, 0)
+        errctx = _ErrorContext(AutoReconnect('mock'), 0, 0, True)
         t.handle_error(('a', 27017), errctx)
         self.assertEqual(SERVER_TYPE.Unknown, get_type(t, 'a'))
         self.assertEqual(SERVER_TYPE.RSSecondary, get_type(t, 'b'))
@@ -480,7 +480,7 @@ def test_handle_error_removed_server(self):
         t = create_mock_topology(replica_set_name='rs')
 
         # No error resetting a server not in the TopologyDescription.
-        errctx = _ErrorContext(AutoReconnect('mock'), 0, 0)
+        errctx = _ErrorContext(AutoReconnect('mock'), 0, 0, True)
         t.handle_error(('b', 27017), errctx)
 
         # Server was *not* added as type Unknown.