Support load-balanced mode (#698)

Support load-balanced mode in the synchronous driver.  Reactive streams driver support to follow.

This is an implementation of the 1.0.0 version of the load balancer support specification published to https://github.com/mongodb/specifications/blob/master/source/load-balancers/load-balancers.rst.

JAVA-4078, JAVA-4079, JAVA-4080, JAVA-4081, JAVA-4082, JAVA-4083

Author: jyemin

.evergreen/.evg.yml

@@ -240,6 +240,22 @@ functions:
         script: |
           DRIVERS_TOOLS="${DRIVERS_TOOLS}" sh ${DRIVERS_TOOLS}/.evergreen/atlas_data_lake/run-mongohouse-local.sh
 
+  "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
+
   "run tests":
     - command: shell.exec
       type: test
@@ -262,6 +278,18 @@ functions:
           REQUIRE_API_VERSION=${REQUIRE_API_VERSION} \
           .evergreen/run-tests.sh
 
+  "run load-balancer tests":
+    - command: shell.exec
+      type: test
+      params:
+        working_dir: "src"
+        script: |
+          ${PREPARE_SHELL}
+          AUTH="${AUTH}" SSL="${SSL}" JDK="${JDK}" \
+          SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}" \
+          MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}" \
+          .evergreen/run-load-balancer-tests.sh
+
   "run reactive streams tck tests":
     - command: shell.exec
       type: test
@@ -695,6 +723,7 @@ post:
   - func: "upload working dir logs"
   - func: "upload test results"
   - func: "cleanup"
+  - func: "stop load-balancer"
 
 tasks:
 
@@ -711,6 +740,12 @@ tasks:
         - func: "bootstrap mongo-orchestration"
         - func: "run tests"
 
+    - name: load-balancer-test
+      commands:
+        - func: "bootstrap mongo-orchestration"
+        - func: "run load-balancer"
+        - func: "run load-balancer tests"
+
     - name: "scala-tests"
       commands:
         - func: "bootstrap mongo-orchestration"
@@ -1425,6 +1460,12 @@ buildvariants:
   tasks:
     - name: "test"
 
+- matrix_name: "tests-load-balancer-secure"
+  matrix_spec: { auth: "auth", ssl: "ssl", jdk: ["jdk11"], version: ["latest"], topology: "sharded-cluster", os: "ubuntu" }
+  display_name: "Load Balancer ${version} ${auth} ${ssl} ${jdk} ${os}"
+  tasks:
+    - name: "load-balancer-test"
+
 - matrix_name: "tests-slow"
   matrix_spec: { auth: "noauth", ssl: "*", jdk: "jdk8", version: ["4.2"], topology: "*", os: "linux" }
   display_name: "Slow: ${version} ${topology} ${ssl} ${jdk} ${os} "

.evergreen/run-load-balancer-tests.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+set -o xtrace  # Write all commands first to stderr
+set -o errexit # Exit the script with error if any of the commands fail
+
+# Supported/used environment variables:
+#       AUTH                        Set to enable authentication. Values are: "auth" / "noauth" (default)
+#       SSL                         Set to enable SSL. Values are "ssl" / "nossl" (default)
+#       JDK                         Set the version of java to be used.  Java versions can be set from the java toolchain /opt/java
+#       SINGLE_MONGOS_LB_URI        Set the URI pointing to a load balancer configured with a single mongos server
+#       MULTI_MONGOS_LB_URI         Set the URI pointing to a load balancer configured with multiple mongos servers
+
+AUTH=${AUTH:-noauth}
+SSL=${SSL:-nossl}
+MONGODB_URI=${MONGODB_URI:-}
+JDK=${JDK:-jdk8}
+
+export JAVA_HOME="/opt/java/jdk11"
+
+############################################
+#            Main Program                  #
+############################################
+
+if [ "$SSL" != "nossl" ]; then
+  # We generate the keystore and truststore on every run with the certs in the drivers-tools repo
+  if [ ! -f client.pkc ]; then
+    openssl pkcs12 -CAfile ${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem -export -in ${DRIVERS_TOOLS}/.evergreen/x509gen/client.pem -out client.pkc -password pass:bithere
+  fi
+
+  cp ${JAVA_HOME}/lib/security/cacerts mongo-truststore
+  ${JAVA_HOME}/bin/keytool -importcert -trustcacerts -file ${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem -keystore mongo-truststore -storepass changeit -storetype JKS -noprompt
+
+  # We add extra gradle arguments for SSL
+  GRADLE_EXTRA_VARS="-Pssl.enabled=true -Pssl.keyStoreType=pkcs12 -Pssl.keyStore=$(pwd)/client.pkc -Pssl.keyStorePassword=bithere -Pssl.trustStoreType=jks -Pssl.trustStore=$(pwd)/mongo-truststore -Pssl.trustStorePassword=changeit"
+  SINGLE_MONGOS_LB_URI="${SINGLE_MONGOS_LB_URI}&ssl=true&sslInvalidHostNameAllowed=true"
+  MULTI_MONGOS_LB_URI="${MULTI_MONGOS_LB_URI}&ssl=true&sslInvalidHostNameAllowed=true"
+fi
+
+echo "Running $AUTH tests over $SSL and connecting to $SINGLE_MONGOS_LB_URI"
+
+echo "Running tests with ${JDK}"
+./gradlew -version
+
+./gradlew -PjdkHome=/opt/java/${JDK} \
+  -Dorg.mongodb.test.uri=${SINGLE_MONGOS_LB_URI} \
+  -Dorg.mongodb.test.transaction.uri=${MULTI_MONGOS_LB_URI} \
+  ${GRADLE_EXTRA_VARS} --stacktrace --info --continue driver-sync:test --tests LoadBalancerTest

driver-core/src/main/com/mongodb/ConnectionString.java

@@ -228,6 +228,7 @@
  * {@link MongoClientSettings#getUuidRepresentation()} for documentation of semantics of this parameter.  Defaults to "javaLegacy", but
  * will change to "unspecified" in the next major release.</li>
  * <li>{@code directConnection=true|false}. If true the driver will set the connection to be a direct connection to the host.</li>
+ * <li>{@code loadBalanced=true|false}. If true the driver will assume that it's connecting to MongoDB through a load balancer.</li>
  * </ul>
  *
  * @mongodb.driver.manual reference/connection-string Connection String Format
@@ -237,7 +238,8 @@ public class ConnectionString {
 
     private static final String MONGODB_PREFIX = "mongodb://";
     private static final String MONGODB_SRV_PREFIX = "mongodb+srv://";
-    private static final Set<String> ALLOWED_OPTIONS_IN_TXT_RECORD = new HashSet<String>(asList("authsource", "replicaset"));
+    private static final Set<String> ALLOWED_OPTIONS_IN_TXT_RECORD =
+            new HashSet<String>(asList("authsource", "replicaset", "loadbalanced"));
     private static final String UTF_8 = "UTF-8";
 
     private static final Logger LOGGER = Loggers.getLogger("uri");
@@ -250,6 +252,7 @@ public class ConnectionString {
     private final String connectionString;
 
     private Boolean directConnection;
+    private Boolean loadBalanced;
     private ReadPreference readPreference;
     private WriteConcern writeConcern;
     private Boolean retryWrites;
@@ -401,6 +404,18 @@ public ConnectionString(final String connectionString) {
             }
         }
 
+        if (loadBalanced != null && loadBalanced) {
+            if (directConnection != null) {
+                throw new IllegalArgumentException("directConnection can not be specified with loadBalanced=true");
+            }
+            if (requiredReplicaSetName != null) {
+                throw new IllegalArgumentException("replicaSet can not be specified with loadBalanced=true");
+            }
+            if (hosts.size() > 1) {
+                throw new IllegalArgumentException("Only one host can be specified with loadBalanced=true");
+            }
+        }
+
         credential = createCredentials(combinedOptionsMaps, userName, password);
         warnOnUnsupportedOptions(combinedOptionsMaps);
     }
@@ -448,6 +463,7 @@ public ConnectionString(final String connectionString) {
         GENERAL_OPTIONS_KEYS.add("uuidrepresentation");
 
         GENERAL_OPTIONS_KEYS.add("directconnection");
+        GENERAL_OPTIONS_KEYS.add("loadbalanced");
 
         COMPRESSOR_KEYS.add("compressors");
         COMPRESSOR_KEYS.add("zlibcompressionlevel");
@@ -551,6 +567,8 @@ private void translateOptions(final Map<String, List<String>> optionsMap) {
                 uuidRepresentation = createUuidRepresentation(value);
             } else if (key.equals("directconnection")) {
                 directConnection = parseBoolean(value, "directconnection");
+            } else if (key.equals("loadbalanced")) {
+                loadBalanced = parseBoolean(value, "loadbalanced");
             }
         }
 
@@ -1119,6 +1137,17 @@ public Boolean isDirectConnection() {
         return directConnection;
     }
 
+    /**
+     * Indicates if the connection is through a load balancer.
+     *
+     * @return true if a load-balanced connection
+     * @since 4.3
+     */
+    @Nullable
+    public Boolean isLoadBalanced() {
+        return loadBalanced;
+    }
+
     /**
      * Get the unparsed connection string.
      *

driver-core/src/main/com/mongodb/ReadPreference.java

@@ -135,6 +135,8 @@ public final List<ServerDescription> choose(final ClusterDescription clusterDesc
             case SHARDED:
             case STANDALONE:
                 return chooseForNonReplicaSet(clusterDescription);
+            case LOAD_BALANCED:
+                return clusterDescription.getServerDescriptions();
             case UNKNOWN:
                 return Collections.emptyList();
             default:

driver-core/src/main/com/mongodb/assertions/Assertions.java

@@ -177,6 +177,13 @@ public static boolean assertFalse(final boolean value) throws AssertionError {
         return false;
     }
 
+    /**
+     * @throws AssertionError Always
+     */
+    public static void fail() throws AssertionError {
+        throw new AssertionError();
+    }
+
     private Assertions() {
     }
 }

driver-core/src/main/com/mongodb/connection/ClusterConnectionMode.java

@@ -30,5 +30,12 @@ public enum ClusterConnectionMode {
     /**
      * Connect to multiple servers in a cluster (either a replica set or multiple mongos servers)
      */
-    MULTIPLE
+    MULTIPLE,
+
+    /**
+     * Connect to one or more mongos servers via a load balancer
+     *
+     * @since 4.3
+     */
+    LOAD_BALANCED
 }

driver-core/src/main/com/mongodb/connection/ClusterSettings.java

@@ -265,8 +265,16 @@ public Builder addClusterListener(final ClusterListener clusterListener) {
          */
         public Builder applyConnectionString(final ConnectionString connectionString) {
             Boolean directConnection = connectionString.isDirectConnection();
-
-            if (connectionString.isSrvProtocol()) {
+            Boolean loadBalanced = connectionString.isLoadBalanced();
+
+            if (loadBalanced != null && loadBalanced) {
+                mode(ClusterConnectionMode.LOAD_BALANCED);
+                if (connectionString.isSrvProtocol()) {
+                    srvHost(connectionString.getHosts().get(0));
+                } else {
+                    hosts(singletonList(createServerAddress(connectionString.getHosts().get(0))));
+                }
+            } else if (connectionString.isSrvProtocol()) {
                 mode(ClusterConnectionMode.MULTIPLE);
                 srvHost(connectionString.getHosts().get(0));
             } else if ((directConnection != null && directConnection)
@@ -553,13 +561,17 @@ private ClusterSettings(final Builder builder) {
             }
         }
 
+        if (builder.mode == ClusterConnectionMode.LOAD_BALANCED && builder.srvHost == null && builder.hosts.size() != 1) {
+            throw new IllegalArgumentException("Multiple hosts cannot be specified when in load balancing mode");
+        }
+
         srvHost = builder.srvHost;
         hosts = builder.hosts;
         if (srvHost != null) {
             if (builder.mode == ClusterConnectionMode.SINGLE) {
                 throw new IllegalArgumentException("An SRV host name was provided but the connection mode is not MULTIPLE");
             }
-            mode = ClusterConnectionMode.MULTIPLE;
+            mode = builder.mode != null ? builder.mode : ClusterConnectionMode.MULTIPLE;
         } else {
             if (builder.mode == ClusterConnectionMode.SINGLE && builder.hosts.size() > 1) {
                 throw new IllegalArgumentException("Can not directly connect to more than one server");

driver-core/src/main/com/mongodb/connection/ClusterType.java

@@ -37,6 +37,12 @@ public enum ClusterType {
      */
     SHARDED,
 
+    /**
+     * A load-balanced cluster, connected via a single load balancer
+     *
+     * @since 4.3
+     */
+    LOAD_BALANCED,
     /**
      * The cluster type is not yet known.
      */

driver-core/src/main/com/mongodb/connection/ConnectionDescription.java

@@ -18,7 +18,9 @@
 
 import com.mongodb.ServerAddress;
 import com.mongodb.annotations.Immutable;
+import com.mongodb.lang.Nullable;
 import org.bson.BsonArray;
+import org.bson.types.ObjectId;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -34,6 +36,7 @@
  */
 @Immutable
 public class ConnectionDescription {
+    @Nullable private final ObjectId serviceId;
     private final ConnectionId connectionId;
     private final int maxWireVersion;
     private final ServerType serverType;
@@ -90,6 +93,28 @@ public ConnectionDescription(final ConnectionId connectionId, final int maxWireV
     public ConnectionDescription(final ConnectionId connectionId, final int maxWireVersion,
                                  final ServerType serverType, final int maxBatchCount, final int maxDocumentSize,
                                  final int maxMessageSize, final List<String> compressors, final BsonArray saslSupportedMechanisms) {
+        this(null, connectionId, maxWireVersion, serverType, maxBatchCount, maxDocumentSize, maxMessageSize, compressors,
+                saslSupportedMechanisms);
+    }
+
+    /**
+     * Construct an instance.
+     *
+     * @param serviceId       the service id, which may be null
+     * @param connectionId    the connection id
+     * @param maxWireVersion  the max wire version
+     * @param serverType      the server type
+     * @param maxBatchCount   the max batch count
+     * @param maxDocumentSize the max document size in bytes
+     * @param maxMessageSize  the max message size in bytes
+     * @param compressors     the available compressors on the connection
+     * @param saslSupportedMechanisms the supported SASL mechanisms
+     * @since 4.3
+     */
+    public ConnectionDescription(@Nullable final ObjectId serviceId, final ConnectionId connectionId, final int maxWireVersion,
+                                 final ServerType serverType, final int maxBatchCount, final int maxDocumentSize,
+                                 final int maxMessageSize, final List<String> compressors, final BsonArray saslSupportedMechanisms) {
+        this.serviceId = serviceId;
         this.connectionId = connectionId;
         this.serverType = serverType;
         this.maxBatchCount = maxBatchCount;
@@ -99,7 +124,6 @@ public ConnectionDescription(final ConnectionId connectionId, final int maxWireV
         this.compressors = notNull("compressors", Collections.unmodifiableList(new ArrayList<String>(compressors)));
         this.saslSupportedMechanisms = saslSupportedMechanisms;
     }
-
     /**
      * Creates a new connection description with the set connection id
      *
@@ -109,8 +133,21 @@ public ConnectionDescription(final ConnectionId connectionId, final int maxWireV
      */
     public ConnectionDescription withConnectionId(final ConnectionId connectionId) {
         notNull("connectionId", connectionId);
-        return new ConnectionDescription(connectionId, maxWireVersion, serverType, maxBatchCount, maxDocumentSize, maxMessageSize,
-                compressors, saslSupportedMechanisms);
+        return new ConnectionDescription(serviceId, connectionId, maxWireVersion, serverType, maxBatchCount, maxDocumentSize,
+                maxMessageSize, compressors, saslSupportedMechanisms);
+    }
+
+    /**
+     * Creates a new connection description with the given service id
+     *
+     * @param serviceId the service id
+     * @return the new connection description
+     * @since 4.3
+     */
+    public ConnectionDescription withServiceId(final ObjectId serviceId) {
+        notNull("serviceId", serviceId);
+        return new ConnectionDescription(serviceId, connectionId, maxWireVersion, serverType, maxBatchCount, maxDocumentSize,
+                maxMessageSize, compressors, saslSupportedMechanisms);
     }
 
     /**
@@ -131,6 +168,17 @@ public ConnectionId getConnectionId() {
         return connectionId;
     }
 
+    /**
+     * Gets the id of the service this connection is to
+     *
+     * @return the service id, which may be null
+     * @since 4.3
+     */
+    @Nullable
+    public ObjectId getServiceId() {
+        return serviceId;
+    }
+
     /**
      * The latest version of the wire protocol that this MongoDB server is capable of using to communicate with clients.
      *