JAVA-2737: Always set slaveOk for direct connections to replica set members

jyemin · jyemin · commit cb5e5862bc17 · 2018-01-18T13:36:21.000-05:00
* for OP_QUERY (for MongoDB versions &lt; 3.6), that means settings the slaveOk wire protocol bit.
* for OP_MSG (for MongoDB versions &gt;= 3.6), that means including a $readPreference field with mode of "primaryPreferred".
diff --git a/driver-core/src/main/com/mongodb/connection/CommandMessage.java b/driver-core/src/main/com/mongodb/connection/CommandMessage.java
@@ -38,8 +38,12 @@
 import java.util.Map;
 
 import static com.mongodb.ReadPreference.primary;
+import static com.mongodb.ReadPreference.primaryPreferred;
 import static com.mongodb.assertions.Assertions.isTrue;
 import static com.mongodb.connection.BsonWriterHelper.writePayload;
+import static com.mongodb.connection.ClusterConnectionMode.MULTIPLE;
+import static com.mongodb.connection.ClusterConnectionMode.SINGLE;
+import static com.mongodb.connection.ServerType.SHARD_ROUTER;
 
 /**
  * A command message that uses OP_MSG or OP_QUERY to send the command.
@@ -52,15 +56,18 @@ final class CommandMessage extends RequestMessage {
     private final SplittablePayload payload;
     private final FieldNameValidator payloadFieldNameValidator;
     private final boolean responseExpected;
+    private final ClusterConnectionMode clusterConnectionMode;
 
     CommandMessage(final MongoNamespace namespace, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
                    final ReadPreference readPreference, final MessageSettings settings) {
-        this(namespace, command, commandFieldNameValidator, readPreference, settings, true, null, null);
+        this(namespace, command, commandFieldNameValidator, readPreference, settings, true, null, null,
+                MULTIPLE);
     }
 
     CommandMessage(final MongoNamespace namespace, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
                    final ReadPreference readPreference, final MessageSettings settings, final boolean responseExpected,
-                   final SplittablePayload payload, final FieldNameValidator payloadFieldNameValidator) {
+                   final SplittablePayload payload, final FieldNameValidator payloadFieldNameValidator,
+                   final ClusterConnectionMode clusterConnectionMode) {
         super(namespace.getFullName(), getOpCode(settings), settings);
         this.namespace = namespace;
         this.command = command;
@@ -69,6 +76,7 @@ final class CommandMessage extends RequestMessage {
         this.responseExpected = responseExpected;
         this.payload = payload;
         this.payloadFieldNameValidator = payloadFieldNameValidator;
+        this.clusterConnectionMode = clusterConnectionMode;
     }
 
     BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
@@ -102,7 +110,7 @@ boolean containsPayload() {
 
     boolean isResponseExpected() {
         isTrue("The message must be encoded before determining if a response is expected", getEncodingMetadata() != null);
-        return calculateIsResponseExpected();
+        return !useOpMsg() || requireOpMsgResponse();
     }
 
     ReadPreference getReadPreference() {
@@ -134,9 +142,9 @@ protected EncodingMetadata encodeMessageBodyWithMetadata(final BsonOutput bsonOu
             }
 
             // Write the flag bits
-            bsonOutput.writeInt32(flagPosition, getFlagBits());
+            bsonOutput.writeInt32(flagPosition, getOpMsgFlagBits());
         } else {
-            bsonOutput.writeInt32(0);
+            bsonOutput.writeInt32(getOpQueryFlagBits());
             bsonOutput.writeCString(namespace.getFullName());
             bsonOutput.writeInt32(0);
             bsonOutput.writeInt32(-1);
@@ -167,20 +175,44 @@ private void addDocumentWithPayload(final BsonOutput bsonOutput, final int messa
         getCodec(commandToEncode).encode(bsonWriter, commandToEncode, EncoderContext.builder().build());
     }
 
-    private int getFlagBits() {
-        if (calculateIsResponseExpected()) {
+    private int getOpMsgFlagBits() {
+        return getOpMsgResponseExpectedFlagBit();
+    }
+
+    private int getOpMsgResponseExpectedFlagBit() {
+        if (requireOpMsgResponse()) {
             return 0;
         } else {
             return 1 << 1;
         }
     }
 
-    private boolean calculateIsResponseExpected() {
-        // If there is another message in the payload require that the response is acknowledged
-        if (!responseExpected && useOpMsg() && payload != null && payload.hasAnotherSplit()) {
+    private boolean requireOpMsgResponse() {
+        if (responseExpected) {
             return true;
+        } else {
+            return payload != null && payload.hasAnotherSplit();
         }
-        return responseExpected;
+    }
+
+    private int getOpQueryFlagBits() {
+        return getOpQuerySlaveOkFlagBit();
+    }
+
+    private int getOpQuerySlaveOkFlagBit() {
+        if (isSlaveOk()) {
+            return 1 << 2;
+        } else {
+            return 0;
+        }
+    }
+
+    private boolean isSlaveOk() {
+        return readPreference.isSlaveOk() || isDirectConnectionToNonShardRouter();
+    }
+
+    private boolean isDirectConnectionToNonShardRouter() {
+        return clusterConnectionMode == SINGLE && getSettings().getServerType() != SHARD_ROUTER;
     }
 
     private boolean useOpMsg() {
@@ -206,10 +238,13 @@ private List<BsonElement> getExtraElements(final SessionContext sessionContext)
         }
         if (!isDefaultReadPreference(getReadPreference())) {
             extraElements.add(new BsonElement("$readPreference", getReadPreference().toDocument()));
+        } else if (isDirectConnectionToNonShardRouter()) {
+                extraElements.add(new BsonElement("$readPreference", primaryPreferred().toDocument()));
         }
         return extraElements;
     }
 
+    @SuppressWarnings("BooleanMethodIsAlwaysInverted")
     private boolean isDefaultReadPreference(final ReadPreference readPreference) {
         return readPreference.equals(primary());
     }
diff --git a/driver-core/src/main/com/mongodb/connection/CommandProtocolImpl.java b/driver-core/src/main/com/mongodb/connection/CommandProtocolImpl.java
@@ -41,16 +41,19 @@ class CommandProtocolImpl<T> implements CommandProtocol<T> {
     private final FieldNameValidator payloadFieldNameValidator;
     private final Decoder<T> commandResultDecoder;
     private final boolean responseExpected;
+    private final ClusterConnectionMode clusterConnectionMode;
     private SessionContext sessionContext;
 
     CommandProtocolImpl(final String database, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
                         final ReadPreference readPreference, final Decoder<T> commandResultDecoder) {
-        this(database, command, commandFieldNameValidator, readPreference, commandResultDecoder, true, null, null);
+        this(database, command, commandFieldNameValidator, readPreference, commandResultDecoder, true, null, null,
+                ClusterConnectionMode.MULTIPLE);
     }
 
     CommandProtocolImpl(final String database, final BsonDocument command, final FieldNameValidator commandFieldNameValidator,
                         final ReadPreference readPreference, final Decoder<T> commandResultDecoder, final boolean responseExpected,
-                        final SplittablePayload payload, final FieldNameValidator payloadFieldNameValidator) {
+                        final SplittablePayload payload, final FieldNameValidator payloadFieldNameValidator,
+                        final ClusterConnectionMode clusterConnectionMode) {
         notNull("database", database);
         this.namespace = new MongoNamespace(notNull("database", database), MongoNamespace.COMMAND_COLLECTION_NAME);
         this.command = notNull("command", command);
@@ -60,6 +63,7 @@ class CommandProtocolImpl<T> implements CommandProtocol<T> {
         this.responseExpected = responseExpected;
         this.payload = payload;
         this.payloadFieldNameValidator = payloadFieldNameValidator;
+        this.clusterConnectionMode = notNull("clusterConnectionMode", clusterConnectionMode);
 
         isTrueArgument("payloadFieldNameValidator cannot be null if there is a payload.",
                 payload == null || payloadFieldNameValidator != null);
@@ -112,7 +116,8 @@ public CommandProtocolImpl<T> sessionContext(final SessionContext sessionContext
 
     private CommandMessage getCommandMessage(final InternalConnection connection) {
         return new CommandMessage(namespace, command, commandFieldNameValidator, readPreference,
-                    getMessageSettings(connection.getDescription()), responseExpected, payload, payloadFieldNameValidator);
+                    getMessageSettings(connection.getDescription()), responseExpected, payload,
+                payloadFieldNameValidator, clusterConnectionMode);
     }
 
     private String getCommandName() {
diff --git a/driver-core/src/main/com/mongodb/connection/DefaultServerConnection.java b/driver-core/src/main/com/mongodb/connection/DefaultServerConnection.java
@@ -124,7 +124,7 @@ public <T> T command(final String database, final BsonDocument command, final Fi
                          final boolean responseExpected, final SplittablePayload payload,
                          final FieldNameValidator payloadFieldNameValidator) {
         return executeProtocol(new CommandProtocolImpl<T>(database, command, commandFieldNameValidator, readPreference,
-                commandResultDecoder, responseExpected, payload, payloadFieldNameValidator), sessionContext);
+                commandResultDecoder, responseExpected, payload, payloadFieldNameValidator, clusterConnectionMode), sessionContext);
     }
 
     @Override
@@ -149,7 +149,8 @@ public <T> void commandAsync(final String database, final BsonDocument command,
                                  final SessionContext sessionContext, final boolean responseExpected, final SplittablePayload payload,
                                  final FieldNameValidator payloadFieldNameValidator, final SingleResultCallback<T> callback) {
         executeProtocolAsync(new CommandProtocolImpl<T>(database, command, commandFieldNameValidator, readPreference,
-                commandResultDecoder, responseExpected, payload,  payloadFieldNameValidator), sessionContext, callback);
+                commandResultDecoder, responseExpected, payload,  payloadFieldNameValidator, clusterConnectionMode),
+                sessionContext, callback);
     }
 
     @Override
diff --git a/driver-core/src/main/com/mongodb/connection/MessageSettings.java b/driver-core/src/main/com/mongodb/connection/MessageSettings.java
@@ -34,6 +34,7 @@ final class MessageSettings {
     private final int maxMessageSize;
     private final int maxBatchCount;
     private final ServerVersion serverVersion;
+    private final ServerType serverType;
 
     /**
      * Gets the builder
@@ -53,6 +54,7 @@ public static final class Builder {
         private int maxMessageSize = DEFAULT_MAX_MESSAGE_SIZE;
         private int maxBatchCount = DEFAULT_MAX_BATCH_COUNT;
         private ServerVersion serverVersion;
+        private ServerType serverType;
 
         /**
          * Build it.
@@ -100,6 +102,11 @@ public Builder serverVersion(final ServerVersion serverVersion) {
             this.serverVersion = serverVersion;
             return this;
         }
+
+        public Builder serverType(final ServerType serverType) {
+            this.serverType = serverType;
+            return this;
+        }
     }
 
     /**
@@ -133,10 +140,15 @@ public ServerVersion getServerVersion() {
         return serverVersion;
     }
 
+    public ServerType getServerType() {
+        return serverType;
+    }
+
     private MessageSettings(final Builder builder) {
         this.maxDocumentSize = builder.maxDocumentSize;
         this.maxMessageSize = builder.maxMessageSize;
         this.maxBatchCount = builder.maxBatchCount;
         this.serverVersion = builder.serverVersion;
+        this.serverType = builder.serverType;
     }
 }
diff --git a/driver-core/src/main/com/mongodb/connection/ProtocolHelper.java b/driver-core/src/main/com/mongodb/connection/ProtocolHelper.java
@@ -182,11 +182,12 @@ static MongoException getQueryFailureException(final BsonDocument errorDocument,
 
     static MessageSettings getMessageSettings(final ConnectionDescription connectionDescription) {
         return MessageSettings.builder()
-                              .maxDocumentSize(connectionDescription.getMaxDocumentSize())
-                              .maxMessageSize(connectionDescription.getMaxMessageSize())
-                              .maxBatchCount(connectionDescription.getMaxBatchCount())
-                              .serverVersion(connectionDescription.getServerVersion())
-                              .build();
+                .maxDocumentSize(connectionDescription.getMaxDocumentSize())
+                .maxMessageSize(connectionDescription.getMaxMessageSize())
+                .maxBatchCount(connectionDescription.getMaxBatchCount())
+                .serverVersion(connectionDescription.getServerVersion())
+                .serverType(connectionDescription.getServerType())
+                .build();
     }
 
     static void encodeMessage(final RequestMessage message, final BsonOutput bsonOutput) {
diff --git a/driver-core/src/main/com/mongodb/connection/ReplyHeader.java b/driver-core/src/main/com/mongodb/connection/ReplyHeader.java
@@ -50,6 +50,7 @@ final class ReplyHeader {
     private final long cursorId;
     private final int startingFrom;
     private final int numberReturned;
+    private final int opMsgFlagBits;
 
     ReplyHeader(final ByteBuf header, final MessageHeader messageHeader) {
         this(messageHeader.getMessageLength(), messageHeader.getOpCode(), messageHeader, header);
@@ -70,8 +71,8 @@ private ReplyHeader(final int messageLength, final int opCode, final MessageHead
             startingFrom = 0;
             numberReturned = 1;
 
-            header.getInt();  // ignore flag bits
-            header.get();     // ignore payload type
+            opMsgFlagBits = header.getInt();
+            header.get();  // ignore payload type
         } else if (opCode == OP_REPLY.getValue()) {
             if (messageLength < TOTAL_REPLY_HEADER_LENGTH) {
                 throw new MongoInternalException(format("The reply message length %d is less than the mimimum message length %d",
@@ -82,6 +83,7 @@ private ReplyHeader(final int messageLength, final int opCode, final MessageHead
             cursorId = header.getLong();
             startingFrom = header.getInt();
             numberReturned = header.getInt();
+            opMsgFlagBits = 0;
 
             if (numberReturned < 0) {
                 throw new MongoInternalException(format("The reply message number of returned documents, %d, is less than 0",
@@ -187,4 +189,9 @@ public boolean isCursorNotFound() {
     public boolean isQueryFailure() {
         return (responseFlags & QUERY_FAILURE_RESPONSE_FLAG) == QUERY_FAILURE_RESPONSE_FLAG;
     }
+
+    // for unit testing
+    int getOpMsgFlagBits() {
+        return opMsgFlagBits;
+    }
 }
diff --git a/driver-core/src/test/unit/com/mongodb/connection/CommandMessageSpecification.groovy b/driver-core/src/test/unit/com/mongodb/connection/CommandMessageSpecification.groovy
diff --git a/driver-core/src/test/unit/com/mongodb/connection/DefaultServerConnectionSpecification.groovy b/driver-core/src/test/unit/com/mongodb/connection/DefaultServerConnectionSpecification.groovy
