Skip to content

Commit 40e3744

Browse files
committed
Support any number of Document Sequence Sections in CommandMessage#getCommandDocument
JAVA-5536
1 parent 7d80cef commit 40e3744

File tree

2 files changed

+55
-48
lines changed

2 files changed

+55
-48
lines changed

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

Lines changed: 12 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -53,38 +53,26 @@ final class ByteBufBsonDocument extends BsonDocument {
5353

5454
private final transient ByteBuf byteBuf;
5555

56-
static List<ByteBufBsonDocument> createList(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
57-
List<ByteBuf> duplicateByteBuffers = bsonOutput.getByteBuffers();
58-
CompositeByteBuf outputByteBuf = new CompositeByteBuf(duplicateByteBuffers);
59-
outputByteBuf.position(startPosition);
56+
// Create a list of ByteBufBsonDocument from a ByteBuf positioned at the start of the first document of an OP_MSG Section
57+
// of type Kind 1: Document Sequence
58+
// The provided ByteBuf will be positioned at the end of the section upon normal completion of the method
59+
static List<ByteBufBsonDocument> createList(final ByteBuf outputByteBuf) {
6060
List<ByteBufBsonDocument> documents = new ArrayList<>();
61-
int curDocumentStartPosition = startPosition;
6261
while (outputByteBuf.hasRemaining()) {
63-
int documentSizeInBytes = outputByteBuf.getInt();
64-
ByteBuf slice = outputByteBuf.duplicate();
65-
slice.position(curDocumentStartPosition);
66-
slice.limit(curDocumentStartPosition + documentSizeInBytes);
67-
documents.add(new ByteBufBsonDocument(slice));
68-
curDocumentStartPosition += documentSizeInBytes;
69-
outputByteBuf.position(outputByteBuf.position() + documentSizeInBytes - 4);
70-
}
71-
for (ByteBuf byteBuffer : duplicateByteBuffers) {
72-
byteBuffer.release();
62+
ByteBufBsonDocument curDocument = createOne(outputByteBuf);
63+
documents.add(curDocument);
7364
}
7465
return documents;
7566
}
7667

77-
static ByteBufBsonDocument createOne(final ByteBufferBsonOutput bsonOutput, final int startPosition) {
78-
List<ByteBuf> duplicateByteBuffers = bsonOutput.getByteBuffers();
79-
CompositeByteBuf outputByteBuf = new CompositeByteBuf(duplicateByteBuffers);
80-
outputByteBuf.position(startPosition);
68+
// Create a ByteBufBsonDocument from a ByteBuf positioned at the start of a BSON document.
69+
// The provided ByteBuf will be positioned at the end of the section upon normal completion of the method
70+
static ByteBufBsonDocument createOne(final ByteBuf outputByteBuf) {
8171
int documentSizeInBytes = outputByteBuf.getInt();
8272
ByteBuf slice = outputByteBuf.duplicate();
83-
slice.position(startPosition);
84-
slice.limit(startPosition + documentSizeInBytes);
85-
for (ByteBuf byteBuffer : duplicateByteBuffers) {
86-
byteBuffer.release();
87-
}
73+
slice.position(slice.position() - 4);
74+
slice.limit(slice.position() + documentSizeInBytes);
75+
outputByteBuf.position(outputByteBuf.position() + documentSizeInBytes - 4);
8876
return new ByteBufBsonDocument(slice);
8977
}
9078

@@ -138,10 +126,6 @@ <T> T findInDocument(final Finder<T> finder) {
138126
return finder.notFound();
139127
}
140128

141-
int getSizeInBytes() {
142-
return byteBuf.getInt(byteBuf.position());
143-
}
144-
145129
BsonDocument toBaseBsonDocument() {
146130
ByteBuf duplicateByteBuf = byteBuf.duplicate();
147131
try (BsonBinaryReader bsonReader = new BsonBinaryReader(new ByteBufferBsonInput(duplicateByteBuf))) {

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

Lines changed: 43 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,14 @@
3030
import org.bson.BsonElement;
3131
import org.bson.BsonInt64;
3232
import org.bson.BsonString;
33+
import org.bson.ByteBuf;
3334
import org.bson.FieldNameValidator;
3435
import org.bson.io.BsonOutput;
3536

3637
import java.nio.charset.StandardCharsets;
3738
import java.util.ArrayList;
3839
import java.util.List;
40+
import java.util.stream.IntStream;
3941

4042
import static com.mongodb.ReadPreference.primary;
4143
import static com.mongodb.ReadPreference.primaryPreferred;
@@ -46,6 +48,8 @@
4648
import static com.mongodb.connection.ServerType.SHARD_ROUTER;
4749
import static com.mongodb.connection.ServerType.STANDALONE;
4850
import static com.mongodb.internal.connection.BsonWriterHelper.writePayload;
51+
import static com.mongodb.internal.connection.ByteBufBsonDocument.createList;
52+
import static com.mongodb.internal.connection.ByteBufBsonDocument.createOne;
4953
import static com.mongodb.internal.connection.ReadConcernHelper.getReadConcernDocument;
5054
import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_TWO_WIRE_VERSION;
5155
import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_ZERO_WIRE_VERSION;
@@ -107,30 +111,49 @@ public final class CommandMessage extends RequestMessage {
107111
this.serverApi = serverApi;
108112
}
109113

114+
// Create a BsonDocument representing the logical document encoded by an OP_MSG.
115+
// The returned document will contain all the fields, including ones represented by
116+
// OP_MSG Sections of type Kind 1: Document Sequence
110117
BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
111-
ByteBufBsonDocument byteBufBsonDocument = ByteBufBsonDocument.createOne(bsonOutput,
112-
getEncodingMetadata().getFirstDocumentPosition());
113-
BsonDocument commandBsonDocument;
114-
115-
if (containsPayload()) {
116-
commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
117-
118-
int payloadStartPosition = getEncodingMetadata().getFirstDocumentPosition()
119-
+ byteBufBsonDocument.getSizeInBytes()
120-
+ 1 // payload type
121-
+ 4 // payload size
122-
+ payload.getPayloadName().getBytes(StandardCharsets.UTF_8).length + 1; // null-terminated UTF-8 payload name
123-
commandBsonDocument.append(payload.getPayloadName(),
124-
new BsonArray(ByteBufBsonDocument.createList(bsonOutput, payloadStartPosition)));
125-
} else {
126-
commandBsonDocument = byteBufBsonDocument;
118+
List<ByteBuf> duplicateByteBuffers = bsonOutput.getByteBuffers();
119+
try {
120+
CompositeByteBuf outputByteBuf = new CompositeByteBuf(duplicateByteBuffers);
121+
outputByteBuf.position(getEncodingMetadata().getFirstDocumentPosition());
122+
ByteBufBsonDocument byteBufBsonDocument = createOne(outputByteBuf);
123+
124+
if (outputByteBuf.hasRemaining()) {
125+
BsonDocument commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
126+
127+
while (outputByteBuf.hasRemaining()) {
128+
outputByteBuf.position(outputByteBuf.position() + 1 /* payload type */ + 4 /* payload size */);
129+
String payloadName = getPayloadName(outputByteBuf);
130+
assertFalse(payloadName.contains("."));
131+
commandBsonDocument.append(payloadName, new BsonArray(createList(outputByteBuf)));
132+
}
133+
return commandBsonDocument;
134+
} else {
135+
return byteBufBsonDocument;
136+
}
137+
} finally {
138+
duplicateByteBuffers.forEach(ByteBuf::release);
127139
}
128-
129-
return commandBsonDocument;
130140
}
131141

132-
boolean containsPayload() {
133-
return payload != null;
142+
// Get the field name from a ByteBuf positioned at the start of the document sequence identifier of an OP_MSG Section of type
143+
// Kind 1: Document Sequence.
144+
// Upon normal completion of the method, the ByteBuf will be positioned at the start of the first BSON object in the sequence.
145+
private String getPayloadName(final ByteBuf outputByteBuf) {
146+
List<Byte> payloadNameBytes = new ArrayList<>();
147+
byte curByte = outputByteBuf.get();
148+
while (curByte != 0) {
149+
payloadNameBytes.add(curByte);
150+
curByte = outputByteBuf.get();
151+
}
152+
153+
// Convert List<Byte> to byte[]
154+
byte[] byteArray = new byte[payloadNameBytes.size()];
155+
IntStream.range(0, payloadNameBytes.size()).forEachOrdered(i -> byteArray[i] = payloadNameBytes.get(i));
156+
return new String(byteArray, StandardCharsets.UTF_8);
134157
}
135158

136159
boolean isResponseExpected() {

0 commit comments

Comments
 (0)