Skip to content

Commit 14fc2fa

Browse files
authored
Support any number of Document Sequence Sections in CommandMessage#getCommandDocument (#1456)
JAVA-5536
1 parent a461dba commit 14fc2fa

File tree

2 files changed

+90
-49
lines changed

2 files changed

+90
-49
lines changed

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

Lines changed: 18 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -53,38 +53,31 @@ 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+
/**
57+
* Create a list of ByteBufBsonDocument from a buffer positioned at the start of the first document of an OP_MSG Section
58+
* of type Document Sequence (Kind 1).
59+
* <p>
60+
* The provided buffer will be positioned at the end of the section upon normal completion of the method
61+
*/
62+
static List<ByteBufBsonDocument> createList(final ByteBuf outputByteBuf) {
6063
List<ByteBufBsonDocument> documents = new ArrayList<>();
61-
int curDocumentStartPosition = startPosition;
6264
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();
65+
ByteBufBsonDocument curDocument = createOne(outputByteBuf);
66+
documents.add(curDocument);
7367
}
7468
return documents;
7569
}
7670

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);
71+
/**
72+
* Create a ByteBufBsonDocument from a buffer positioned at the start of a BSON document.
73+
* The provided buffer will be positioned at the end of the document upon normal completion of the method
74+
*/
75+
static ByteBufBsonDocument createOne(final ByteBuf outputByteBuf) {
76+
int documentStart = outputByteBuf.position();
8177
int documentSizeInBytes = outputByteBuf.getInt();
82-
ByteBuf slice = outputByteBuf.duplicate();
83-
slice.position(startPosition);
84-
slice.limit(startPosition + documentSizeInBytes);
85-
for (ByteBuf byteBuffer : duplicateByteBuffers) {
86-
byteBuffer.release();
87-
}
78+
int documentEnd = documentStart + documentSizeInBytes;
79+
ByteBuf slice = outputByteBuf.duplicate().position(documentStart).limit(documentEnd);
80+
outputByteBuf.position(documentEnd);
8881
return new ByteBufBsonDocument(slice);
8982
}
9083

@@ -138,10 +131,6 @@ <T> T findInDocument(final Finder<T> finder) {
138131
return finder.notFound();
139132
}
140133

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

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

Lines changed: 72 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package com.mongodb.internal.connection;
1818

1919
import com.mongodb.MongoClientException;
20+
import com.mongodb.MongoInternalException;
2021
import com.mongodb.MongoNamespace;
2122
import com.mongodb.ReadPreference;
2223
import com.mongodb.ServerApi;
@@ -31,9 +32,12 @@
3132
import org.bson.BsonElement;
3233
import org.bson.BsonInt64;
3334
import org.bson.BsonString;
35+
import org.bson.ByteBuf;
3436
import org.bson.FieldNameValidator;
3537
import org.bson.io.BsonOutput;
3638

39+
import java.io.ByteArrayOutputStream;
40+
import java.io.UnsupportedEncodingException;
3741
import java.nio.charset.StandardCharsets;
3842
import java.util.ArrayList;
3943
import java.util.List;
@@ -47,6 +51,8 @@
4751
import static com.mongodb.connection.ServerType.SHARD_ROUTER;
4852
import static com.mongodb.connection.ServerType.STANDALONE;
4953
import static com.mongodb.internal.connection.BsonWriterHelper.writePayload;
54+
import static com.mongodb.internal.connection.ByteBufBsonDocument.createList;
55+
import static com.mongodb.internal.connection.ByteBufBsonDocument.createOne;
5056
import static com.mongodb.internal.connection.ReadConcernHelper.getReadConcernDocument;
5157
import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_TWO_WIRE_VERSION;
5258
import static com.mongodb.internal.operation.ServerVersionHelper.FOUR_DOT_ZERO_WIRE_VERSION;
@@ -108,30 +114,76 @@ public final class CommandMessage extends RequestMessage {
108114
this.serverApi = serverApi;
109115
}
110116

117+
/**
118+
* Create a BsonDocument representing the logical document encoded by an OP_MSG.
119+
* <p>
120+
* The returned document will contain all the fields from the Body (Kind 0) Section, as well as all fields represented by
121+
* OP_MSG Document Sequence (Kind 1) Sections.
122+
*/
111123
BsonDocument getCommandDocument(final ByteBufferBsonOutput bsonOutput) {
112-
ByteBufBsonDocument byteBufBsonDocument = ByteBufBsonDocument.createOne(bsonOutput,
113-
getEncodingMetadata().getFirstDocumentPosition());
114-
BsonDocument commandBsonDocument;
115-
116-
if (containsPayload()) {
117-
commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
118-
119-
int payloadStartPosition = getEncodingMetadata().getFirstDocumentPosition()
120-
+ byteBufBsonDocument.getSizeInBytes()
121-
+ 1 // payload type
122-
+ 4 // payload size
123-
+ payload.getPayloadName().getBytes(StandardCharsets.UTF_8).length + 1; // null-terminated UTF-8 payload name
124-
commandBsonDocument.append(payload.getPayloadName(),
125-
new BsonArray(ByteBufBsonDocument.createList(bsonOutput, payloadStartPosition)));
126-
} else {
127-
commandBsonDocument = byteBufBsonDocument;
124+
List<ByteBuf> byteBuffers = bsonOutput.getByteBuffers();
125+
try {
126+
CompositeByteBuf byteBuf = new CompositeByteBuf(byteBuffers);
127+
try {
128+
byteBuf.position(getEncodingMetadata().getFirstDocumentPosition());
129+
ByteBufBsonDocument byteBufBsonDocument = createOne(byteBuf);
130+
131+
// If true, it means there is at least one Kind 1:Document Sequence in the OP_MSG
132+
if (byteBuf.hasRemaining()) {
133+
BsonDocument commandBsonDocument = byteBufBsonDocument.toBaseBsonDocument();
134+
135+
// Each loop iteration processes one Document Sequence
136+
// When there are no more bytes remaining, there are no more Document Sequences
137+
while (byteBuf.hasRemaining()) {
138+
// skip reading the payload type, we know it is 1
139+
byteBuf.position(byteBuf.position() + 1);
140+
int sequenceStart = byteBuf.position();
141+
int sequenceSizeInBytes = byteBuf.getInt();
142+
int sectionEnd = sequenceStart + sequenceSizeInBytes;
143+
144+
String fieldName = getSequenceIdentifier(byteBuf);
145+
// If this assertion fires, it means that the driver has started using document sequences for nested fields. If
146+
// so, this method will need to change in order to append the value to the correct nested document.
147+
assertFalse(fieldName.contains("."));
148+
149+
ByteBuf documentsByteBufSlice = byteBuf.duplicate().limit(sectionEnd);
150+
try {
151+
commandBsonDocument.append(fieldName, new BsonArray(createList(documentsByteBufSlice)));
152+
} finally {
153+
documentsByteBufSlice.release();
154+
}
155+
byteBuf.position(sectionEnd);
156+
}
157+
return commandBsonDocument;
158+
} else {
159+
return byteBufBsonDocument;
160+
}
161+
} finally {
162+
byteBuf.release();
163+
}
164+
} finally {
165+
byteBuffers.forEach(ByteBuf::release);
128166
}
129-
130-
return commandBsonDocument;
131167
}
132168

133-
boolean containsPayload() {
134-
return payload != null;
169+
/**
170+
* Get the field name from a buffer positioned at the start of the document sequence identifier of an OP_MSG Section of type
171+
* Document Sequence (Kind 1).
172+
* <p>
173+
* Upon normal completion of the method, the buffer will be positioned at the start of the first BSON object in the sequence.
174+
*/
175+
private String getSequenceIdentifier(final ByteBuf byteBuf) {
176+
ByteArrayOutputStream sequenceIdentifierBytes = new ByteArrayOutputStream();
177+
byte curByte = byteBuf.get();
178+
while (curByte != 0) {
179+
sequenceIdentifierBytes.write(curByte);
180+
curByte = byteBuf.get();
181+
}
182+
try {
183+
return sequenceIdentifierBytes.toString(StandardCharsets.UTF_8.name());
184+
} catch (UnsupportedEncodingException e) {
185+
throw new MongoInternalException("Unexpected exception", e);
186+
}
135187
}
136188

137189
boolean isResponseExpected() {

0 commit comments

Comments
 (0)