JAVA-2636: Move logic from InternalStreamConnection to CommandMessage for getting the BsonDocument to use for command monitoring.

Take advantage of the observation that a command with a payload is likely to be small (minus the payload)
so it's a reasonable approach to simple convert the BSON back into a BsonDocument and then add the payload as a BsonArray
initialized with a List<ByteBufBsonDocument> from the OP_MSG payload.

Author: jyemin

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

@@ -41,7 +41,7 @@ class ByteBufBsonDocument extends AbstractByteBufBsonDocument {
 
     private final transient ByteBuf byteBuf;
 
-    static List<ByteBufBsonDocument> create(final ResponseBuffers responseBuffers) {
+    static List<ByteBufBsonDocument> createList(final ResponseBuffers responseBuffers) {
         int numDocuments = responseBuffers.getReplyHeader().getNumberReturned();
         ByteBuf documentsBuffer = responseBuffers.getBodyByteBuffer();
         documentsBuffer.order(ByteOrder.LITTLE_ENDIAN);
@@ -57,11 +57,7 @@ static List<ByteBufBsonDocument> create(final ResponseBuffers responseBuffers) {
         return documents;
     }
 
-    static ByteBufBsonDocument createOne(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
-        return create(bsonOutput, startPosition).get(0);
-    }
-
-    static List<ByteBufBsonDocument> create(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
+    static List<ByteBufBsonDocument> createList(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
         List<ByteBuf> duplicateByteBuffers = bsonOutput.getByteBuffers();
         CompositeByteBuf outputByteBuf = new CompositeByteBuf(duplicateByteBuffers);
         outputByteBuf.position(startPosition);
@@ -82,6 +78,20 @@ static List<ByteBufBsonDocument> create(final ByteBufferBsonOutput bsonOutput, f
         return documents;
     }
 
+    static ByteBufBsonDocument createOne(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
+        List<ByteBuf> duplicateByteBuffers = bsonOutput.getByteBuffers();
+        CompositeByteBuf outputByteBuf = new CompositeByteBuf(duplicateByteBuffers);
+        outputByteBuf.position(startPosition);
+        int documentSizeInBytes = outputByteBuf.getInt();
+        ByteBuf slice = outputByteBuf.duplicate();
+        slice.position(startPosition);
+        slice.limit(startPosition + documentSizeInBytes);
+        for (ByteBuf byteBuffer : duplicateByteBuffers) {
+            byteBuffer.release();
+        }
+        return new ByteBufBsonDocument(slice);
+    }
+
     @Override
     @SuppressWarnings("deprecation")
     public String toJson() {
@@ -130,6 +140,10 @@ public BsonDocument clone() {
         return new RawBsonDocument(clonedBytes);
     }
 
+    int getSizeInBytes() {
+        return byteBuf.getInt(byteBuf.position());
+    }
+
     BsonDocument toBsonDocument() {
         ByteBuf duplicateByteBuf = byteBuf.duplicate();
         BsonBinaryReader bsonReader = new BsonBinaryReader(new ByteBufferBsonInput(duplicateByteBuf));

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

@@ -20,6 +20,7 @@
 import com.mongodb.MongoNamespace;
 import com.mongodb.ReadPreference;
 import com.mongodb.internal.validator.MappedFieldNameValidator;
+import org.bson.BsonArray;
 import org.bson.BsonBinaryWriter;
 import org.bson.BsonDocument;
 import org.bson.BsonElement;
@@ -29,6 +30,7 @@
 import org.bson.codecs.EncoderContext;
 import org.bson.io.BsonOutput;
 
+import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -68,6 +70,31 @@ final class CommandMessage extends RequestMessage {
         this.readPreference = readPreference;
     }
 
+    BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
+        ByteBufBsonDocument byteBufBsonDocument = ByteBufBsonDocument.createOne(bsonOutput,
+                getEncodingMetadata().getFirstDocumentPosition());
+        BsonDocument commandBsonDocument;
+
+        if (useOpMsg() && containsPayload()) {
+            commandBsonDocument = byteBufBsonDocument.toBsonDocument();
+
+            int payloadStartPosition = getEncodingMetadata().getFirstDocumentPosition()
+                    + byteBufBsonDocument.getSizeInBytes()
+                    + 1 // payload type
+                    + 4 // payload size
+                    + payload.getPayloadName().getBytes(Charset.forName("UTF-8")).length + 1;  // null-terminated UTF-8 payload name
+            commandBsonDocument.append(payload.getPayloadName(),
+                    new BsonArray(ByteBufBsonDocument.createList(bsonOutput, payloadStartPosition)));
+        } else {
+            commandBsonDocument = byteBufBsonDocument;
+        }
+
+        if (commandBsonDocument.containsKey("$query")) {
+            commandBsonDocument = commandBsonDocument.getDocument("$query");
+        }
+        return commandBsonDocument;
+    }
+
     boolean containsPayload() {
         return payload != null;
     }

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

@@ -179,7 +179,7 @@ private BsonDocument asGetMoreCommandResponseDocument(final QueryResult<T> query
         List<ByteBufBsonDocument> rawResultDocuments = Collections.emptyList();
         if (responseBuffers.getReplyHeader().getNumberReturned() != 0) {
             responseBuffers.reset();
-            rawResultDocuments = ByteBufBsonDocument.create(responseBuffers);
+            rawResultDocuments = ByteBufBsonDocument.createList(responseBuffers);
         }
 
         BsonDocument cursorDocument = new BsonDocument("id",

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

@@ -83,7 +83,7 @@ public void onResult(final WriteConcernResult result, final Throwable t) {
     @Override
     protected BsonDocument getAsWriteCommand(final ByteBufferBsonOutput bsonOutput, final int firstDocumentPosition) {
         return getBaseCommandDocument("insert")
-               .append("documents", new BsonArray(ByteBufBsonDocument.create(bsonOutput, firstDocumentPosition)));
+               .append("documents", new BsonArray(ByteBufBsonDocument.createList(bsonOutput, firstDocumentPosition)));
 
     }
 

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

@@ -675,28 +675,13 @@ public String getCommandName() {
         public void sendStartedEvent(final ByteBufferBsonOutput bsonOutput) {
             if ((commandListener != null || sendCompressor != null) && opened()) {
 
-                AbstractByteBufBsonDocument byteBufBsonDocument;
-                if (message.containsPayload()) {
-                    byteBufBsonDocument = ByteBufPayloadBsonDocument.create(bsonOutput,
-                            message.getEncodingMetadata().getFirstDocumentPosition());
-                } else {
-                    byteBufBsonDocument = ByteBufBsonDocument.createOne(bsonOutput,
-                            message.getEncodingMetadata().getFirstDocumentPosition());
-                }
-
-                BsonDocument commandBsonDocument = byteBufBsonDocument;
-                if (byteBufBsonDocument.containsKey("$query")) {
-                    commandBsonDocument = byteBufBsonDocument.getDocument("$query");
-                    commandName = commandBsonDocument.keySet().iterator().next();
-                } else {
-                    commandName = byteBufBsonDocument.getFirstKey();
-                }
+                BsonDocument commandDocument = message.getCommandDocument(bsonOutput);
+                commandName = commandDocument.getFirstKey();
 
                 BsonDocument commandDocumentForEvent = (SECURITY_SENSITIVE_COMMANDS.contains(commandName))
-                                                               ? new BsonDocument() : commandBsonDocument;
+                                                               ? new BsonDocument() : commandDocument;
                 sendCommandStartedEvent(message, new MongoNamespace(message.getCollectionName()).getDatabaseName(), commandName,
                         commandDocumentForEvent, getDescription(), commandListener);
-
             }
         }
 

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

@@ -413,7 +413,7 @@ private BsonDocument asFindCommandDocument(final ByteBufferBsonOutput bsonOutput
 
         boolean isExplain = false;
 
-        List<ByteBufBsonDocument> documents = ByteBufBsonDocument.create(bsonOutput, firstDocumentPosition);
+        List<ByteBufBsonDocument> documents = ByteBufBsonDocument.createList(bsonOutput, firstDocumentPosition);
 
         ByteBufBsonDocument rawQueryDocument = documents.get(0);
         for (Map.Entry<String, BsonValue> cur : rawQueryDocument.entrySet()) {
@@ -474,7 +474,7 @@ private BsonDocument asFindCommandResponseDocument(final ResponseBuffers respons
         List<ByteBufBsonDocument> rawResultDocuments = Collections.emptyList();
         if (responseBuffers.getReplyHeader().getNumberReturned() > 0) {
             responseBuffers.reset();
-            rawResultDocuments = ByteBufBsonDocument.create(responseBuffers);
+            rawResultDocuments = ByteBufBsonDocument.createList(responseBuffers);
         }
 
         if (isExplain) {

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

@@ -82,7 +82,7 @@ public void onResult(final WriteConcernResult result, final Throwable t) {
 
     @Override
     protected BsonDocument getAsWriteCommand(final ByteBufferBsonOutput bsonOutput, final int firstDocumentPosition) {
-        List<ByteBufBsonDocument> documents = ByteBufBsonDocument.create(bsonOutput, firstDocumentPosition);
+        List<ByteBufBsonDocument> documents = ByteBufBsonDocument.createList(bsonOutput, firstDocumentPosition);
         BsonDocument updateDocument = new BsonDocument("q", documents.get(0)).append("u", documents.get(1));
         if (updateRequest.isMulti()) {
             updateDocument.append("multi", BsonBoolean.TRUE);