PYTHON-3336 Test Failure - test_load_balancer failing (#1000)

Author: blink1073

.pre-commit-config.yaml

@@ -56,7 +56,7 @@ repos:
   rev: 0.11.1
   hooks:
     - id: doc8
-      args: [--max-line-length=200]
+      args: ["--ignore=D001"]  # ignore line length
       stages: [manual]
 
 - repo: https://github.com/sirosen/check-jsonschema

CONTRIBUTING.rst

@@ -70,6 +70,23 @@ branch and submit a `pull request <https://help.github.com/articles/using-pull-r
 You might also use the GitHub `Edit <https://github.com/blog/844-forking-with-the-edit-button>`_
 button.
 
+Running Tests Locally
+---------------------
+- Ensure you have started the appropriate Mongo Server(s).
+- Run ``python setup.py test`` to run all of the tests.
+- Run ``python setup.py test -s test.<mod_name>.<class_name>.<test_name>`` to
+  run specific tests.  You can omit the ``<test_name>`` to test a full class
+  and the ``<class_name>`` to test a full module.  For example:
+  ``python setup.py test -s test.test_change_stream.TestUnifiedChangeStreamsErrors.test_change_stream_errors_on_ElectionInProgress``.
+
+Running Load Balancer Tests Locally
+-----------------------------------
+- Install ``haproxy`` (available as ``brew install haproxy`` on macOS).
+- Clone ``drivers-evergreen-tools``: ``git clone [email protected]:mongodb-labs/drivers-evergreen-tools.git``.
+- Start the servers using ``LOAD_BALANCER=true TOPOLOGY=sharded_cluster AUTH=noauth SSL=nossl MONGODB_VERSION=6.0 DRIVERS_TOOLS=./drivers-evergreen-tools MONGO_ORCHESTRATION_HOME=./drivers-evergreen-tools/.evergreen/orchestration ./drivers-evergreen-tools/.evergreen/run-orchestration.sh``.
+- Start the load balancer using: ``MONGODB_URI='mongodb://localhost:27017,localhost:27018/' .evergreen/run-load-balancer.sh start``.
+- Run the tests using: ``LOADBALANCER=1 TEST_LOADBALANCER=1 SINGLE_MONGOS_LB_URI='mongodb://127.0.0.1:8000/?loadBalanced=true' MULTI_MONGOS_LB_URI='mongodb://127.0.0.1:8001/?loadBalanced=true' MONGODB_URI='mongodb://localhost:27017,localhost:27018/' python setup.py test -s test.test_load_balancer``.
+
 Re-sync Spec Tests
 ------------------
 

pymongo/errors.py

@@ -61,6 +61,15 @@ class ConnectionFailure(PyMongoError):
     """Raised when a connection to the database cannot be made or is lost."""
 
 
+class WaitQueueTimeoutError(ConnectionFailure):
+    """Raised when an operation times out waiting to checkout a connection from the pool.
+
+    Subclass of :exc:`~pymongo.errors.ConnectionFailure`.
+
+    .. versionadded:: 4.2
+    """
+
+
 class AutoReconnect(ConnectionFailure):
     """Raised when a connection to the database is lost and an attempt to
     auto-reconnect will be made.

pymongo/mongo_client.py

@@ -80,6 +80,7 @@
     OperationFailure,
     PyMongoError,
     ServerSelectionTimeoutError,
+    WaitQueueTimeoutError,
 )
 from pymongo.pool import ConnectionClosedReason
 from pymongo.read_preferences import ReadPreference, _ServerMode
@@ -1182,6 +1183,7 @@ def _get_socket(self, server, session):
         with _MongoClientErrorHandler(self, server, session) as err_handler:
             # Reuse the pinned connection, if it exists.
             if in_txn and session._pinned_connection:
+                err_handler.contribute_socket(session._pinned_connection)
                 yield session._pinned_connection
                 return
             with server.get_socket(handler=err_handler) as sock_info:
@@ -2064,9 +2066,11 @@ def _add_retryable_write_error(exc, max_wire_version):
             if code in helpers._RETRYABLE_ERROR_CODES:
                 exc._add_error_label("RetryableWriteError")
 
-    # Connection errors are always retryable except NotPrimaryError which is
+    # Connection errors are always retryable except NotPrimaryError and WaitQueueTimeoutError which is
     # handled above.
-    if isinstance(exc, ConnectionFailure) and not isinstance(exc, NotPrimaryError):
+    if isinstance(exc, ConnectionFailure) and not isinstance(
+        exc, (NotPrimaryError, WaitQueueTimeoutError)
+    ):
         exc._add_error_label("RetryableWriteError")
 
 

pymongo/monitoring.py

@@ -1774,7 +1774,7 @@ def publish_connection_check_out_failed(self, address, reason):
         event = ConnectionCheckOutFailedEvent(address, reason)
         for subscriber in self.__cmap_listeners:
             try:
-                subscriber.connection_check_out_started(event)
+                subscriber.connection_check_out_failed(event)
             except Exception:
                 _handle_exception()
 

pymongo/pool.py

@@ -52,6 +52,7 @@
     NotPrimaryError,
     OperationFailure,
     PyMongoError,
+    WaitQueueTimeoutError,
     _CertificateError,
 )
 from pymongo.hello import Hello, HelloCompat
@@ -1637,7 +1638,7 @@ def _raise_wait_queue_timeout(self) -> NoReturn:
         timeout = _csot.get_timeout() or self.opts.wait_queue_timeout
         if self.opts.load_balanced:
             other_ops = self.active_sockets - self.ncursors - self.ntxns
-            raise ConnectionFailure(
+            raise WaitQueueTimeoutError(
                 "Timeout waiting for connection from the connection pool. "
                 "maxPoolSize: %s, connections in use by cursors: %s, "
                 "connections in use by transactions: %s, connections in use "
@@ -1650,7 +1651,7 @@ def _raise_wait_queue_timeout(self) -> NoReturn:
                     timeout,
                 )
             )
-        raise ConnectionFailure(
+        raise WaitQueueTimeoutError(
             "Timed out while checking out a connection from connection pool. "
             "maxPoolSize: %s, timeout: %s" % (self.opts.max_pool_size, timeout)
         )

test/test_cmap.py

@@ -38,7 +38,12 @@
 
 from bson.objectid import ObjectId
 from bson.son import SON
-from pymongo.errors import ConnectionFailure, OperationFailure, PyMongoError
+from pymongo.errors import (
+    ConnectionFailure,
+    OperationFailure,
+    PyMongoError,
+    WaitQueueTimeoutError,
+)
 from pymongo.monitoring import (
     ConnectionCheckedInEvent,
     ConnectionCheckedOutEvent,
@@ -73,7 +78,7 @@
     "ConnectionPoolClosed": PoolClosedEvent,
     # Error types.
     "PoolClosedError": _PoolClosedError,
-    "WaitQueueTimeoutError": ConnectionFailure,
+    "WaitQueueTimeoutError": WaitQueueTimeoutError,
 }
 
 

test/utils.py

@@ -38,7 +38,20 @@
 from pymongo.cursor import CursorType
 from pymongo.errors import ConfigurationError, OperationFailure
 from pymongo.hello import HelloCompat
-from pymongo.monitoring import _SENSITIVE_COMMANDS
+from pymongo.monitoring import (
+    _SENSITIVE_COMMANDS,
+    ConnectionCheckedInEvent,
+    ConnectionCheckedOutEvent,
+    ConnectionCheckOutFailedEvent,
+    ConnectionCheckOutStartedEvent,
+    ConnectionClosedEvent,
+    ConnectionCreatedEvent,
+    ConnectionReadyEvent,
+    PoolClearedEvent,
+    PoolClosedEvent,
+    PoolCreatedEvent,
+    PoolReadyEvent,
+)
 from pymongo.pool import _CancellationContext, _PoolGeneration
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference
@@ -81,36 +94,47 @@ def wait_for_event(self, event, count):
 
 class CMAPListener(BaseListener, monitoring.ConnectionPoolListener):
     def connection_created(self, event):
+        assert isinstance(event, ConnectionCreatedEvent)
         self.add_event(event)
 
     def connection_ready(self, event):
+        assert isinstance(event, ConnectionReadyEvent)
         self.add_event(event)
 
     def connection_closed(self, event):
+        assert isinstance(event, ConnectionClosedEvent)
         self.add_event(event)
 
     def connection_check_out_started(self, event):
+        assert isinstance(event, ConnectionCheckOutStartedEvent)
         self.add_event(event)
 
     def connection_check_out_failed(self, event):
+        assert isinstance(event, ConnectionCheckOutFailedEvent)
         self.add_event(event)
 
     def connection_checked_out(self, event):
+        assert isinstance(event, ConnectionCheckedOutEvent)
         self.add_event(event)
 
     def connection_checked_in(self, event):
+        assert isinstance(event, ConnectionCheckedInEvent)
         self.add_event(event)
 
     def pool_created(self, event):
+        assert isinstance(event, PoolCreatedEvent)
         self.add_event(event)
 
     def pool_ready(self, event):
+        assert isinstance(event, PoolReadyEvent)
         self.add_event(event)
 
     def pool_cleared(self, event):
+        assert isinstance(event, PoolClearedEvent)
         self.add_event(event)
 
     def pool_closed(self, event):
+        assert isinstance(event, PoolClosedEvent)
         self.add_event(event)