RUST-980 Run load balancer tests on evergreen, and update existing tests (#477)

Author: abr-egn

.evergreen/config.yml

@@ -365,6 +365,8 @@ functions:
           ${PREPARE_SHELL}
           export MONGODB_URI="${MONGODB_URI}"
           export SSL="${SSL}"
+          export SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}"
+          export MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}"
           . .evergreen/generate-uri.sh
 
           SINGLE_THREAD=${SINGLE_THREAD} \
@@ -603,6 +605,24 @@ functions:
           ${PREPARE_SHELL}
           echo '{"results": [{ "status": "FAIL", "test_file": "Build", "log_raw": "No test-results.json found was created"  } ]}' > ${PROJECT_DIRECTORY}/test-results.json
 
+  "start load balancer":
+    - command: shell.exec
+      params:
+        script: |
+          ${PREPARE_SHELL}
+          export MONGODB_URI="${MONGODB_URI}"
+          ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh start
+    - command: expansions.update
+      params:
+        file: lb-expansion.yml
+  
+  "stop load balancer":
+    - command: shell.exec
+      params:
+        script: |
+          ${PREPARE_SHELL}
+          ${DRIVERS_TOOLS}/.evergreen/run-load-balancer.sh stop
+
 pre:
   - func: "fetch source"
   - func: "prepare resources"
@@ -613,6 +633,7 @@ pre:
   - func: "install dependencies"
 
 post:
+  - func: "stop load balancer"
   - func: "stop mongo orchestration"
   - func: "upload-mo-artifacts"
   - func: "cleanup"
@@ -815,6 +836,17 @@ tasks:
           MONGODB_VERSION: "latest"
           TOPOLOGY: "sharded_cluster"
       - func: "run tests"
+  
+  - name: "test-latest-load_balancer"
+    tags: ["latest", "load_balancer"]
+    commands:
+      - func: "bootstrap mongo-orchestration"
+        vars:
+          MONGODB_VERSION: "latest"
+          TOPOLOGY: "sharded_cluster"
+      - func: "start load balancer"
+      - func: "run tests"
+
 
   - name: "test-latest-aws-auth"
     # "latest" explicitly left off to keep this out of the generic matrix
@@ -1369,6 +1401,13 @@ buildvariants:
         async-runtime: "*"
       then:
         remove_tasks: [".3.6", ".4.0"]
+    # haproxy isn't installed on windows / ubuntu-arm
+    - if:
+        os: ["ubuntu-18.04-arm64", "windows-64-vs2017"]
+        auth-and-tls: "*"
+        async-runtime: "*"
+      then:
+        remove_tasks: ".load_balancer"
 - matrix_name: "x509-auth"
   matrix_spec:
     os:

.evergreen/generate-uri.sh

@@ -10,13 +10,25 @@ DRIVERS_TOOLS_X509=`echo "$DRIVERS_TOOLS_X509" | sed 's/\//%2F/g'`
 CA_FILE="${DRIVERS_TOOLS_X509}%2Fca.pem"
 CERT_FILE="${DRIVERS_TOOLS_X509}%2Fclient.pem"
 
-if [[ "$MONGODB_URI" == *"?"* ]]; then
-    export MONGODB_URI="${MONGODB_URI}&"
-else
-    export MONGODB_URI="${MONGODB_URI}/?"
-fi
-
+update_uri() {
+    local ORIG_URI=$1
+    if [[ "$ORIG_URI" == "" ]]; then
+        return
+    fi
+    # The rustls library requires a domain name.
+    ORIG_URI=$(echo "$ORIG_URI" | sed s/127.0.0.1/localhost/)
+    if [[ "$ORIG_URI" == *"?"* ]]; then
+        ORIG_URI="${ORIG_URI}&"
+    else
+        ORIG_URI="${ORIG_URI}/?"
+    fi
+    echo "${ORIG_URI}tls=true&tlsCAFile=${CA_FILE}&tlsCertificateKeyFile=${CERT_FILE}&tlsAllowInvalidCertificates=true"
+}
 
-export MONGODB_URI="${MONGODB_URI}tls=true&tlsCAFile=${CA_FILE}&tlsCertificateKeyFile=${CERT_FILE}&tlsAllowInvalidCertificates=true"
+export MONGODB_URI="$(update_uri ${MONGODB_URI})"
+export SINGLE_MONGOS_LB_URI="$(update_uri ${SINGLE_MONGOS_LB_URI})"
+export MULTI_MONGOS_LB_URI="$(update_uri ${MULTI_MONGOS_LB_URI})"
 
 echo "MONGODB_URI: ${MONGODB_URI}"
+echo "SINGLE_MONGOS_LB_URI: ${SINGLE_MONGOS_LB_URI}"
+echo "MULTI_MONGOS_LB_URI: ${MULTI_MONGOS_LB_URI}"

src/client/options/mod.rs

@@ -569,6 +569,9 @@ pub(crate) struct TestOptions {
 
     /// Mock response for `SrvPollingMonitor::lookup_hosts`.
     pub(crate) mock_lookup_hosts: Option<Result<LookupHosts>>,
+
+    /// Mock the `serviceId` response for a load-balanced hello.
+    pub(crate) mock_service_id: bool,
 }
 
 fn default_hosts() -> Vec<ServerAddress> {

src/cmap/conn/mod.rs

@@ -313,6 +313,7 @@ impl Connection {
         let (tx, rx) = mpsc::channel(1);
         self.pinned_sender = Some(tx);
         Ok(PinnedConnectionHandle {
+            id: self.id,
             receiver: Arc::new(Mutex::new(rx)),
         })
     }
@@ -369,7 +370,8 @@ impl Drop for Connection {
         if let Some(pool_manager) = self.pool_manager.take() {
             let mut dropped_connection = self.take();
             let result = if let Some(sender) = self.pinned_sender.as_mut() {
-                // Preserve the timestamp for pinned connections.
+                // Preserve the pool manager and timestamp for pinned connections.
+                dropped_connection.pool_manager = Some(pool_manager.clone());
                 dropped_connection.ready_and_available_time = self.ready_and_available_time;
                 match sender.try_send(dropped_connection) {
                     Ok(()) => Ok(()),
@@ -385,7 +387,11 @@ impl Drop for Connection {
                     Err(mpsc::error::TrySendError::Full(mut conn)) => {
                         // Panic in debug mode
                         if cfg!(debug_assertions) {
-                            panic!("buffer full when attempting to return a pinned connection")
+                            panic!(
+                                "buffer full when attempting to return a pinned connection (id = \
+                                 {})",
+                                conn.id
+                            );
                         }
                         // TODO RUST-230 log an error in non-debug mode.
                         conn.pinned_sender = None;
@@ -409,6 +415,7 @@ impl Drop for Connection {
 /// normal pool via this handle.
 #[derive(Debug)]
 pub(crate) struct PinnedConnectionHandle {
+    id: u32,
     receiver: Arc<Mutex<mpsc::Receiver<Connection>>>,
 }
 
@@ -418,6 +425,7 @@ impl PinnedConnectionHandle {
     /// normal borrow.
     pub(crate) fn replicate(&self) -> Self {
         Self {
+            id: self.id,
             receiver: self.receiver.clone(),
         }
     }
@@ -426,10 +434,12 @@ impl PinnedConnectionHandle {
     /// connection has been unpinned.
     pub(crate) async fn take_connection(&self) -> Result<Connection> {
         let mut receiver = self.receiver.lock().await;
-        receiver
-            .recv()
-            .await
-            .ok_or_else(|| Error::internal("cannot take connection after unpin"))
+        receiver.recv().await.ok_or_else(|| {
+            Error::internal(format!(
+                "cannot take connection after unpin (id={})",
+                self.id
+            ))
+        })
     }
 
     /// Return the pinned connection to the normal connection pool.

src/cmap/establish/handshake/mod.rs

@@ -143,6 +143,8 @@ pub(crate) struct Handshaker {
     /// given the same pool options, so it can be created at the time the Handshaker is created.
     command: Command,
     credential: Option<Credential>,
+    #[cfg(test)]
+    mock_service_id: bool,
 }
 
 impl Handshaker {
@@ -154,6 +156,9 @@ impl Handshaker {
         let mut command =
             is_master_command(options.as_ref().and_then(|opts| opts.server_api.as_ref()));
 
+        #[cfg(test)]
+        let mut mock_service_id = false;
+
         if let Some(options) = options {
             if let Some(app_name) = options.app_name {
                 metadata.application = Some(AppMetadata { name: app_name });
@@ -185,13 +190,20 @@ impl Handshaker {
             if options.load_balanced {
                 command.body.insert("loadBalanced", true);
             }
+
+            #[cfg(test)]
+            {
+                mock_service_id = options.mock_service_id;
+            }
         }
 
         command.body.insert("client", metadata);
 
         Self {
             command,
             credential,
+            #[cfg(test)]
+            mock_service_id,
         }
     }
 
@@ -207,6 +219,24 @@ impl Handshaker {
         let client_first = set_speculative_auth_info(&mut command.body, self.credential.as_ref())?;
 
         let mut is_master_reply = run_is_master(conn, command, topology, handler).await?;
+        // TODO PM-2369 Remove serviceId mocking when it's returned by the server.
+        #[cfg(test)]
+        {
+            if self.command.body.contains_key("loadBalanced")
+                && is_master_reply.command_response.service_id.is_none()
+                && self.mock_service_id
+            {
+                is_master_reply.command_response.service_id = Some(
+                    is_master_reply
+                        .command_response
+                        .topology_version
+                        .as_ref()
+                        .unwrap()
+                        .get_object_id("processId")
+                        .unwrap(),
+                );
+            }
+        }
         if self.command.body.contains_key("loadBalanced")
             && is_master_reply.command_response.service_id.is_none()
         {
@@ -256,6 +286,8 @@ pub(crate) struct HandshakerOptions {
     driver_info: Option<DriverInfo>,
     server_api: Option<ServerApi>,
     load_balanced: bool,
+    #[cfg(test)]
+    mock_service_id: bool,
 }
 
 impl From<ConnectionPoolOptions> for HandshakerOptions {
@@ -266,6 +298,8 @@ impl From<ConnectionPoolOptions> for HandshakerOptions {
             driver_info: options.driver_info,
             server_api: options.server_api,
             load_balanced: options.load_balanced.unwrap_or(false),
+            #[cfg(test)]
+            mock_service_id: options.mock_service_id,
         }
     }
 }
@@ -278,6 +312,8 @@ impl From<ClientOptions> for HandshakerOptions {
             driver_info: options.driver_info,
             server_api: options.server_api,
             load_balanced: options.load_balanced.unwrap_or(false),
+            #[cfg(test)]
+            mock_service_id: options.test_options.map_or(false, |to| to.mock_service_id),
         }
     }
 }

src/cmap/mod.rs

@@ -99,7 +99,7 @@ impl ConnectionPool {
         let (manager, _) = manager::channel();
         let handle = PoolWorkerHandle::new_mocked();
         let (connection_requester, _) = connection_requester::channel(Default::default(), handle);
-        let (_, generation_subscriber) = status::channel();
+        let (_, generation_subscriber) = status::channel(PoolGeneration::normal());
 
         Self {
             address,

src/cmap/options.rs

@@ -89,6 +89,11 @@ pub(crate) struct ConnectionPoolOptions {
 
     /// Whether or not the client is connecting to a MongoDB cluster through a load balancer.
     pub(crate) load_balanced: Option<bool>,
+
+    /// Whether or not to mock the `serviceId` field of a hello through a load balancer.
+    #[cfg(test)]
+    #[serde(skip)]
+    pub(crate) mock_service_id: bool,
 }
 
 impl ConnectionPoolOptions {
@@ -109,6 +114,11 @@ impl ConnectionPoolOptions {
             #[cfg(test)]
             ready: None,
             load_balanced: options.load_balanced,
+            #[cfg(test)]
+            mock_service_id: options
+                .test_options
+                .as_ref()
+                .map_or(false, |to| to.mock_service_id),
         }
     }
 

src/cmap/status.rs

@@ -7,17 +7,9 @@ struct PoolStatus {
     generation: PoolGeneration,
 }
 
-impl Default for PoolStatus {
-    fn default() -> Self {
-        PoolStatus {
-            generation: PoolGeneration::normal(),
-        }
-    }
-}
-
 /// Create a channel for publishing and receiving updates to the pool's generation.
-pub(super) fn channel() -> (PoolGenerationPublisher, PoolGenerationSubscriber) {
-    let (sender, receiver) = tokio::sync::watch::channel(Default::default());
+pub(super) fn channel(init: PoolGeneration) -> (PoolGenerationPublisher, PoolGenerationSubscriber) {
+    let (sender, receiver) = tokio::sync::watch::channel(PoolStatus { generation: init });
     (
         PoolGenerationPublisher { sender },
         PoolGenerationSubscriber { receiver },

src/cmap/test/integration.rs

@@ -78,6 +78,10 @@ async fn concurrent_connections() {
     let _guard = LOCK.run_exclusively().await;
 
     let mut options = CLIENT_OPTIONS.clone();
+    if options.load_balanced.unwrap_or(false) {
+        println!("skipping concurrent_connections test due to load-balanced topology");
+        return;
+    }
     options.direct_connection = Some(true);
     options.hosts.drain(1..);
 
@@ -163,6 +167,12 @@ async fn connection_error_during_establishment() {
     let _guard: RwLockWriteGuard<_> = LOCK.run_exclusively().await;
 
     let mut client_options = CLIENT_OPTIONS.clone();
+    if client_options.load_balanced.unwrap_or(false) {
+        println!(
+            "skipping connection_error_during_establishment test due to load-balanced topology"
+        );
+        return;
+    }
     client_options.heartbeat_freq = Duration::from_secs(300).into(); // high so that monitors dont trip failpoint
     client_options.hosts.drain(1..);
     client_options.direct_connection = Some(true);

src/cmap/test/mod.rs

@@ -425,6 +425,9 @@ async fn cmap_spec_tests() {
         let _guard: RwLockWriteGuard<()> = LOCK.run_exclusively().await;
 
         let mut options = CLIENT_OPTIONS.clone();
+        if options.load_balanced.unwrap_or(false) {
+            return;
+        }
         options.hosts.drain(1..);
         options.direct_connection = Some(true);
         let client = EventClient::with_options(options).await;