PYTHON-2676 Add load balancer tests in EVG

.evergreen/config.yml

@@ -411,6 +411,11 @@ functions:
           if [ -n "${SETDEFAULTENCODING}" ]; then
             export SETDEFAULTENCODING="${SETDEFAULTENCODING}"
           fi
+          if [ -n "${test_loadbalancer}" ]; then
+            export TEST_LOADBALANCER=1
+            export SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}"
+            export MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}"
+          fi
 
           PYTHON_BINARY=${PYTHON_BINARY} \
             GREEN_FRAMEWORK=${GREEN_FRAMEWORK} \
@@ -788,6 +793,22 @@ functions:
           -v \
           --fault revoked
 
+  "run load-balancer":
+    - command: shell.exec
+      params:
+        script: |
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} MONGODB_URI=${MONGODB_URI} bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh start
+    - command: expansions.update
+      params:
+        file: lb-expansion.yml
+
+  "stop load-balancer":
+    - command: shell.exec
+      params:
+        script: |
+          cd ${DRIVERS_TOOLS}/.evergreen
+          DRIVERS_TOOLS=${DRIVERS_TOOLS} bash ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh stop
+
   "teardown_docker":
     - command: shell.exec
       params:
@@ -1537,6 +1558,13 @@ tasks:
         - func: "run aws auth test with aws EC2 credentials"
         - func: "run aws ECS auth test"
 
+    - name: load-balancer-test
+      commands:
+        - func: "bootstrap mongo-orchestration"
+          vars:
+            TOPOLOGY: "sharded_cluster"
+        - func: "run load-balancer"
+        - func: "run tests"
 # }}}
     - name: "coverage-report"
       tags: ["coverage"]
@@ -1941,6 +1969,16 @@ axes:
         variables:
           ORCHESTRATION_FILE: "versioned-api-testing.json"
 
+  # Run load balancer tests?
+  - id: loadbalancer
+    display_name: "Load Balancer"
+    values:
+      - id: "enabled"
+        display_name: "Load Balancer"
+        variables:
+          test_loadbalancer: true
+        batchtime: 10080  # 7 days
+
 buildvariants:
 - matrix_name: "tests-all"
   matrix_spec:
@@ -2463,6 +2501,17 @@ buildvariants:
     - name: "aws-auth-test-4.4"
     - name: "aws-auth-test-latest"
 
+- matrix_name: "load-balancer"
+  matrix_spec:
+    platform: ubuntu-18.04
+    mongodb-version: ["latest"]
+    auth-ssl: "*"
+    python-version: ["3.6", "3.9"]
+    loadbalancer: "*"
+  display_name: "Load Balancer ${platform} ${python-version} ${mongodb-version} ${auth-ssl}"
+  tasks:
+    - name: "load-balancer-test"
+
 - matrix_name: "Release"
   matrix_spec:
     platform: [ubuntu-20.04, windows-64-vsMulti-small, macos-1014]

.evergreen/run-tests.sh

@@ -51,6 +51,11 @@ fi
 if [ "$SSL" != "nossl" ]; then
     export CLIENT_PEM="$DRIVERS_TOOLS/.evergreen/x509gen/client.pem"
     export CA_PEM="$DRIVERS_TOOLS/.evergreen/x509gen/ca.pem"
+
+    if [ -n "$TEST_LOADBALANCER" ]; then
+        export SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}&tls=true"
+        export MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}&tls=true"
+    fi
 fi
 
 # For createvirtualenv.
@@ -191,7 +196,12 @@ if [ -z "$GREEN_FRAMEWORK" ]; then
         # causing this script to exit.
         $PYTHON -c "from bson import _cbson; from pymongo import _cmessage"
     fi
-    $PYTHON $COVERAGE_ARGS setup.py $C_EXTENSIONS test $TEST_ARGS $OUTPUT
+
+    if [ -n "$TEST_LOADBALANCER" ]; then
+        $PYTHON -m xmlrunner discover -s test/load_balancer -v --locals -o $XUNIT_DIR
+    else
+        $PYTHON $COVERAGE_ARGS setup.py $C_EXTENSIONS test $TEST_ARGS $OUTPUT
+    fi
 else
     # --no_ext has to come before "test" so there is no way to toggle extensions here.
     $PYTHON green_framework_test.py $GREEN_FRAMEWORK $OUTPUT

pymongo/client_session.py

@@ -660,7 +660,7 @@ def abort_transaction(self):
             pass
         finally:
             self._transaction.state = _TxnState.ABORTED
-            self._unpin_mongos()
+            self._unpin()
 
     def _finish_transaction_with_retry(self, command_name):
         """Run commit or abort with one retry after any retryable error.
@@ -779,13 +779,13 @@ def _pinned_address(self):
             return self._transaction.pinned_address
         return None
 
-    def _pin_mongos(self, server):
-        """Pin this session to the given mongos Server."""
+    def _pin(self, server):
+        """Pin this session to the given Server."""
         self._transaction.sharded = True
         self._transaction.pinned_address = server.description.address
 
-    def _unpin_mongos(self):
-        """Unpin this session from any pinned mongos address."""
+    def _unpin(self):
+        """Unpin this session from any pinned Server."""
         self._transaction.pinned_address = None
 
     def _txn_read_preference(self):
@@ -906,9 +906,11 @@ def get_server_session(self, session_timeout_minutes):
         return _ServerSession(self.generation)
 
     def return_server_session(self, server_session, session_timeout_minutes):
-        self._clear_stale(session_timeout_minutes)
-        if not server_session.timed_out(session_timeout_minutes):
-            self.return_server_session_no_lock(server_session)
+        if session_timeout_minutes is not None:
+            self._clear_stale(session_timeout_minutes)
+            if server_session.timed_out(session_timeout_minutes):
+                return
+        self.return_server_session_no_lock(server_session)
 
     def return_server_session_no_lock(self, server_session):
         # Discard sessions from an old pool to avoid duplicate sessions in the

pymongo/mongo_client.py

@@ -1197,15 +1197,16 @@ def _select_server(self, server_selector, session, address=None):
                 server = topology.select_server(server_selector)
                 # Pin this session to the selected server if it's performing a
                 # sharded transaction.
-                if server.description.mongos and (session and
-                                                  session.in_transaction):
-                    session._pin_mongos(server)
+                if (server.description.server_type in (
+                        SERVER_TYPE.Mongos, SERVER_TYPE.LoadBalancer)
+                        and session and session.in_transaction):
+                    session._pin(server)
             return server
         except PyMongoError as exc:
             # Server selection errors in a transaction are transient.
             if session and session.in_transaction:
                 exc._add_error_label("TransientTransactionError")
-                session._unpin_mongos()
+                session._unpin()
             raise
 
     def _socket_for_writes(self, session):
@@ -1350,7 +1351,7 @@ def is_retrying():
                 _add_retryable_write_error(exc, max_wire_version)
                 retryable_error = exc.has_error_label("RetryableWriteError")
                 if retryable_error:
-                    session._unpin_mongos()
+                    session._unpin()
                 if is_retrying() or not retryable_error:
                     raise
                 if bulk:
@@ -1965,7 +1966,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', 'completed_handshake')
+                 'sock_generation', 'completed_handshake', 'service_id')
 
     def __init__(self, client, server, session):
         self.client = client
@@ -1978,11 +1979,13 @@ def __init__(self, client, server, session):
         # of the pool at the time the connection attempt was started."
         self.sock_generation = server.pool.generation
         self.completed_handshake = False
+        self.service_id = None
 
     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.service_id = sock_info.service_id
         self.completed_handshake = True
 
     def __enter__(self):
@@ -2001,9 +2004,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             if issubclass(exc_type, PyMongoError):
                 if (exc_val.has_error_label("TransientTransactionError") or
                         exc_val.has_error_label("RetryableWriteError")):
-                    self.session._unpin_mongos()
+                    self.session._unpin()
 
         err_ctx = _ErrorContext(
             exc_val, self.max_wire_version, self.sock_generation,
-            self.completed_handshake)
+            self.completed_handshake, self.service_id)
         self.client._topology.handle_error(self.server_address, err_ctx)

pymongo/monitoring.py

@@ -512,13 +512,16 @@ def register(listener):
 class _CommandEvent(object):
     """Base class for command events."""
 
-    __slots__ = ("__cmd_name", "__rqst_id", "__conn_id", "__op_id")
+    __slots__ = ("__cmd_name", "__rqst_id", "__conn_id", "__op_id",
+                 "__service_id")
 
-    def __init__(self, command_name, request_id, connection_id, operation_id):
+    def __init__(self, command_name, request_id, connection_id, operation_id,
+                 service_id=None):
         self.__cmd_name = command_name
         self.__rqst_id = request_id
         self.__conn_id = connection_id
         self.__op_id = operation_id
+        self.__service_id = service_id
 
     @property
     def command_name(self):
@@ -535,6 +538,14 @@ def connection_id(self):
         """The address (host, port) of the server this command was sent to."""
         return self.__conn_id
 
+    @property
+    def service_id(self):
+        """The service_id this command was sent to, or ``None``.
+
+        .. versionadded:: 3.12
+        """
+        return self.__service_id
+
     @property
     def operation_id(self):
         """An id for this series of events or None."""
@@ -551,15 +562,17 @@ class CommandStartedEvent(_CommandEvent):
       - `connection_id`: The address (host, port) of the server this command
         was sent to.
       - `operation_id`: An optional identifier for a series of related events.
+      - `service_id`: The service_id this command was sent to, or ``None``.
     """
     __slots__ = ("__cmd", "__db")
 
-    def __init__(self, command, database_name, *args):
+    def __init__(self, command, database_name, *args, service_id=None):
         if not command:
             raise ValueError("%r is not a valid command" % (command,))
         # Command name must be first key.
         command_name = next(iter(command))
-        super(CommandStartedEvent, self).__init__(command_name, *args)
+        super(CommandStartedEvent, self).__init__(
+            command_name, *args, service_id=service_id)
         if command_name.lower() in _SENSITIVE_COMMANDS:
             self.__cmd = {}
         else:
@@ -577,9 +590,12 @@ def database_name(self):
         return self.__db
 
     def __repr__(self):
-        return "<%s %s db: %r, command: %r, operation_id: %s>" % (
-            self.__class__.__name__, self.connection_id, self.database_name,
-            self.command_name, self.operation_id)
+        return (
+            "<%s %s db: %r, command: %r, operation_id: %s, "
+            "service_id: %s>") % (
+                self.__class__.__name__, self.connection_id,
+                self.database_name, self.command_name, self.operation_id,
+                self.service_id)
 
 
 class CommandSucceededEvent(_CommandEvent):
@@ -593,13 +609,15 @@ class CommandSucceededEvent(_CommandEvent):
       - `connection_id`: The address (host, port) of the server this command
         was sent to.
       - `operation_id`: An optional identifier for a series of related events.
+      - `service_id`: The service_id this command was sent to, or ``None``.
     """
     __slots__ = ("__duration_micros", "__reply")
 
     def __init__(self, duration, reply, command_name,
-                 request_id, connection_id, operation_id):
+                 request_id, connection_id, operation_id, service_id=None):
         super(CommandSucceededEvent, self).__init__(
-            command_name, request_id, connection_id, operation_id)
+            command_name, request_id, connection_id, operation_id,
+            service_id=service_id)
         self.__duration_micros = _to_micros(duration)
         if command_name.lower() in _SENSITIVE_COMMANDS:
             self.__reply = {}
@@ -617,9 +635,12 @@ def reply(self):
         return self.__reply
 
     def __repr__(self):
-        return "<%s %s command: %r, operation_id: %s, duration_micros: %s>" % (
-            self.__class__.__name__, self.connection_id,
-            self.command_name, self.operation_id, self.duration_micros)
+        return (
+            "<%s %s command: %r, operation_id: %s, duration_micros: %s, "
+            "service_id: %s>") % (
+                self.__class__.__name__, self.connection_id,
+                self.command_name, self.operation_id, self.duration_micros,
+                self.service_id)
 
 
 class CommandFailedEvent(_CommandEvent):
@@ -633,11 +654,12 @@ class CommandFailedEvent(_CommandEvent):
       - `connection_id`: The address (host, port) of the server this command
         was sent to.
       - `operation_id`: An optional identifier for a series of related events.
+      - `service_id`: The service_id this command was sent to, or ``None``.
     """
     __slots__ = ("__duration_micros", "__failure")
 
-    def __init__(self, duration, failure, *args):
-        super(CommandFailedEvent, self).__init__(*args)
+    def __init__(self, duration, failure, *args, service_id=None):
+        super(CommandFailedEvent, self).__init__(*args, service_id=service_id)
         self.__duration_micros = _to_micros(duration)
         self.__failure = failure
 
@@ -654,9 +676,10 @@ def failure(self):
     def __repr__(self):
         return (
             "<%s %s command: %r, operation_id: %s, duration_micros: %s, "
-            "failure: %r>" % (
+            "failure: %r, service_id: %s>") % (
                 self.__class__.__name__, self.connection_id, self.command_name,
-                self.operation_id, self.duration_micros, self.failure))
+                self.operation_id, self.duration_micros, self.failure,
+                self.service_id)
 
 
 class _PoolEvent(object):
@@ -721,10 +744,29 @@ class PoolClearedEvent(_PoolEvent):
     :Parameters:
      - `address`: The address (host, port) pair of the server this Pool is
        attempting to connect to.
+     - `service_id`: The service_id this command was sent to, or ``None``.
 
     .. versionadded:: 3.9
     """
-    __slots__ = ()
+    __slots__ = ("__service_id",)
+
+    def __init__(self, address, service_id=None):
+        super(PoolClearedEvent, self).__init__(address)
+        self.__service_id = service_id
+
+    @property
+    def service_id(self):
+        """Connections with this service_id are cleared.
+
+        When service_id is ``None``, all connections in the pool are cleared.
+
+        .. versionadded:: 3.12
+        """
+        return self.__service_id
+
+    def __repr__(self):
+        return '%s(%r, %r)' % (
+            self.__class__.__name__, self.address, self.__service_id)
 
 
 class PoolClosedEvent(_PoolEvent):
@@ -1508,10 +1550,10 @@ def publish_pool_ready(self, address):
             except Exception:
                 _handle_exception()
 
-    def publish_pool_cleared(self, address):
+    def publish_pool_cleared(self, address, service_id):
         """Publish a :class:`PoolClearedEvent` to all pool listeners.
         """
-        event = PoolClearedEvent(address)
+        event = PoolClearedEvent(address, service_id)
         for subscriber in self.__cmap_listeners:
             try:
                 subscriber.pool_cleared(event)