Support loadBalanced in handshake request and serviceId in reply for load balanced clusters (#743)

When connection mode is load balanced

* Add loadBalanced:true in handshake request
* Use serviceId in handshake reply

JAVA-4218

Author: jyemin

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

@@ -25,6 +25,7 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.connection.ServerDescription.getDefaultMaxDocumentSize;
@@ -274,6 +275,9 @@ public boolean equals(final Object o) {
 
         ConnectionDescription that = (ConnectionDescription) o;
 
+        if (maxWireVersion != that.maxWireVersion) {
+            return false;
+        }
         if (maxBatchCount != that.maxBatchCount) {
             return false;
         }
@@ -283,31 +287,32 @@ public boolean equals(final Object o) {
         if (maxMessageSize != that.maxMessageSize) {
             return false;
         }
-        if (!connectionId.equals(that.connectionId)) {
+        if (!Objects.equals(serviceId, that.serviceId)) {
             return false;
         }
-        if (serverType != that.serverType) {
+        if (!connectionId.equals(that.connectionId)) {
             return false;
         }
-        if (maxWireVersion != that.maxWireVersion) {
+        if (serverType != that.serverType) {
             return false;
         }
         if (!compressors.equals(that.compressors)) {
             return false;
         }
-
-        return true;
+        return Objects.equals(saslSupportedMechanisms, that.saslSupportedMechanisms);
     }
 
     @Override
     public int hashCode() {
         int result = connectionId.hashCode();
-        result = 31 * result + maxBatchCount;
+        result = 31 * result + maxWireVersion;
         result = 31 * result + serverType.hashCode();
         result = 31 * result + maxBatchCount;
         result = 31 * result + maxDocumentSize;
         result = 31 * result + maxMessageSize;
         result = 31 * result + compressors.hashCode();
+        result = 31 * result + (serviceId != null ? serviceId.hashCode() : 0);
+        result = 31 * result + (saslSupportedMechanisms != null ? saslSupportedMechanisms.hashCode() : 0);
         return result;
     }
 
@@ -321,6 +326,7 @@ public String toString() {
                 + ", maxDocumentSize=" + maxDocumentSize
                 + ", maxMessageSize=" + maxMessageSize
                 + ", compressors=" + compressors
+                + ", serviceId=" + serviceId
                 + '}';
     }
 }

driver-core/src/main/com/mongodb/internal/connection/DescriptionHelper.java

@@ -42,6 +42,7 @@
 import java.util.Map;
 import java.util.Set;
 
+import static com.mongodb.assertions.Assertions.assertNotNull;
 import static com.mongodb.connection.ConnectionDescription.getDefaultMaxMessageSize;
 import static com.mongodb.connection.ConnectionDescription.getDefaultMaxWriteBatchSize;
 import static com.mongodb.connection.ServerConnectionState.CONNECTED;
@@ -77,11 +78,13 @@ static ConnectionDescription createConnectionDescription(final ClusterConnection
             connectionDescription = connectionDescription.withConnectionId(newConnectionId);
         }
         if (clusterConnectionMode == ClusterConnectionMode.LOAD_BALANCED) {
-            if (manufactureServiceId) {
+            ObjectId serviceId = getServiceId(isMasterResult);
+            if (serviceId != null) {
+                connectionDescription = connectionDescription.withServiceId(serviceId);
+            } else if (manufactureServiceId) {
                 TopologyVersion topologyVersion = getTopologyVersion(isMasterResult);
-                if (topologyVersion != null) {
-                    connectionDescription = connectionDescription.withServiceId(topologyVersion.getProcessId());
-                }
+                assertNotNull(topologyVersion);
+                connectionDescription = connectionDescription.withServiceId(topologyVersion.getProcessId());
             } else {
                 throw new MongoClientException("Driver attempted to initialize in load balancing mode, but the server does not support "
                         + "this mode");
@@ -144,6 +147,11 @@ private static TopologyVersion getTopologyVersion(final BsonDocument isMasterRes
                 ? new TopologyVersion(isMasterResult.getDocument("topologyVersion")) : null;
     }
 
+    private static ObjectId getServiceId(final BsonDocument isMasterResult) {
+        return isMasterResult.containsKey("serviceId") && isMasterResult.get("serviceId").isObjectId()
+                ? isMasterResult.getObjectId("serviceId").getValue() : null;
+    }
+
     private static int getMaxMessageSizeBytes(final BsonDocument isMasterResult) {
         return isMasterResult.getInt32("maxMessageSizeBytes", new BsonInt32(getDefaultMaxMessageSize())).getValue();
     }

driver-core/src/main/com/mongodb/internal/connection/InternalStreamConnectionInitializer.java

@@ -161,6 +161,9 @@ private BsonDocument createIsMasterCommand(final Authenticator authenticator, fi
         if (clientMetadataDocument != null) {
             isMasterCommandDocument.append("client", clientMetadataDocument);
         }
+        if (clusterConnectionMode == ClusterConnectionMode.LOAD_BALANCED) {
+            isMasterCommandDocument.append("loadBalanced", BsonBoolean.TRUE);
+        }
         if (!requestedCompressors.isEmpty()) {
             BsonArray compressors = new BsonArray(this.requestedCompressors.size());
             for (MongoCompressor cur : this.requestedCompressors) {

driver-core/src/test/unit/com/mongodb/internal/connection/DescriptionHelperSpecification.groovy

@@ -16,6 +16,7 @@
 
 package com.mongodb.internal.connection
 
+import com.mongodb.MongoClientException
 import com.mongodb.ServerAddress
 import com.mongodb.Tag
 import com.mongodb.TagSet
@@ -81,6 +82,51 @@ class DescriptionHelperSpecification extends Specification {
                         .withConnectionId(connectionId.withServerValue(1004))
     }
 
+
+    def 'connection description should reflect ismaster result from load balancer'() {
+        given:
+        def connectionId = new ConnectionId(new ServerId(new ClusterId(), serverAddress))
+        ObjectId serviceId = new ObjectId()
+
+        expect:
+        createConnectionDescription(ClusterConnectionMode.LOAD_BALANCED, connectionId,
+                parse("""{
+                        ismaster : true,
+                        msg : "isdbgrid",
+                        maxBsonObjectSize : 16777216,
+                        maxMessageSizeBytes : 48000000,
+                        maxWriteBatchSize : 1000,
+                        localTime : ISODate("2015-03-04T23:55:18.505Z"),
+                        maxWireVersion : 13,
+                        minWireVersion : 0,
+                        connectionId : 1004,
+                        serviceId: {\$oid : "${serviceId.toHexString()}"},
+                        ok : 1
+                        }""")) ==
+                new ConnectionDescription(connectionId, 13, ServerType.SHARD_ROUTER, 1000, 16777216, 48000000, [])
+                        .withConnectionId(connectionId.withServerValue(1004))
+                        .withServiceId(serviceId)
+
+        when:
+        createConnectionDescription(ClusterConnectionMode.LOAD_BALANCED, connectionId,
+                parse('''{
+                        ismaster : true,
+                        msg : "isdbgrid",
+                        maxBsonObjectSize : 16777216,
+                        maxMessageSizeBytes : 48000000,
+                        maxWriteBatchSize : 1000,
+                        localTime : ISODate("2015-03-04T23:55:18.505Z"),
+                        maxWireVersion : 13,
+                        minWireVersion : 0,
+                        connectionId : 1004,
+                        ok : 1
+                        }'''))
+
+        then:
+        def e = thrown(MongoClientException)
+        e.getMessage() == 'Driver attempted to initialize in load balancing mode, but the server does not support this mode'
+    }
+
     def 'connection description should reflect ismaster result with compressors'() {
         def connectionId = new ConnectionId(new ServerId(new ClusterId(), serverAddress))
         expect: