Skip to content

Commit 9fe35e4

Browse files
stIncMalejyemin
stIncMale
and
jyemin
committed
Create DualMessageSequences abstraction
This makes `CommandMessage` generic for any command with two sequences. The new abstraction adds: * The sequence identifiers for each sequence. * A field name validator for both sequences. * A `List<BsonElement>`` for any extra elements required by the splitting logic, so that `txnNumber` doesn't have to be treated specially. Make `SplittablePayload` extend `OpMsgSequence`. This brings SplittablePayload closer in design to `DualMessageSequences`, reducing a potential source of confusion for future readers. JAVA-5529 Co-authored-by: Jeff Yemin <[email protected]>
1 parent ea0633e commit 9fe35e4

File tree

14 files changed

+269
-232
lines changed

14 files changed

+269
-232
lines changed

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

+37-31
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616

1717
package com.mongodb.internal.connection;
1818

19-
import com.mongodb.internal.operation.ClientBulkWriteOperation.ClientBulkWriteCommand;
20-
import com.mongodb.internal.operation.ClientBulkWriteOperation.ClientBulkWriteCommand.OpsAndNsInfo.WritersProviderAndLimitsChecker;
19+
import com.mongodb.internal.connection.DualMessageSequences.EncodeDocumentsResult;
20+
import com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker;
2121
import com.mongodb.internal.validator.NoOpFieldNameValidator;
2222
import com.mongodb.lang.Nullable;
2323
import org.bson.BsonBinaryWriter;
@@ -38,6 +38,9 @@
3838
import java.util.List;
3939

4040
import static com.mongodb.assertions.Assertions.assertTrue;
41+
import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
42+
import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_NOT_REACHED;
43+
import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_REACHED;
4144
import static com.mongodb.internal.connection.MessageSettings.DOCUMENT_HEADROOM_SIZE;
4245
import static java.lang.String.format;
4346
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
@@ -98,46 +101,49 @@ static void writePayload(final BsonWriter writer, final BsonOutput bsonOutput, f
98101
}
99102

100103
/**
101-
* @return See {@link ClientBulkWriteCommand.OpsAndNsInfo#encode(WritersProviderAndLimitsChecker)}.
104+
* @return See {@link DualMessageSequences#encodeDocuments(WritersProviderAndLimitsChecker)}.
102105
*/
103-
static ClientBulkWriteCommand.OpsAndNsInfo.EncodeResult writeOpsAndNsInfo(
104-
final ClientBulkWriteCommand.OpsAndNsInfo opsAndNsInfo,
106+
static EncodeDocumentsResult writeDocumentsOfDualMessageSequences(
107+
final DualMessageSequences dualMessageSequences,
105108
final int commandDocumentSizeInBytes,
106-
final BsonOutput opsOut,
107-
final BsonOutput nsInfoOut,
109+
final BsonOutput firstOutput,
110+
final BsonOutput secondOutput,
108111
final MessageSettings messageSettings,
109112
final boolean validateDocumentSizeLimits) {
110-
BinaryOpsBsonWriters opsWriters = new BinaryOpsBsonWriters(
111-
opsOut,
112-
opsAndNsInfo.getFieldNameValidator(),
113+
BinaryOrdinaryAndStoredBsonWriters firstWriters = new BinaryOrdinaryAndStoredBsonWriters(
114+
firstOutput,
115+
dualMessageSequences.getFirstFieldNameValidator(),
116+
validateDocumentSizeLimits ? messageSettings : null);
117+
BinaryOrdinaryAndStoredBsonWriters secondWriters = new BinaryOrdinaryAndStoredBsonWriters(
118+
secondOutput,
119+
dualMessageSequences.getSecondFieldNameValidator(),
113120
validateDocumentSizeLimits ? messageSettings : null);
114-
BsonBinaryWriter nsInfoWriter = new BsonBinaryWriter(nsInfoOut, NoOpFieldNameValidator.INSTANCE);
115121
// the size of operation-agnostic command fields (a.k.a. extra elements) is counted towards `messageOverheadInBytes`
116122
int messageOverheadInBytes = 1000;
117-
int maxOpsAndNsInfoSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
118-
int opsStart = opsOut.getPosition();
119-
int nsInfoStart = nsInfoOut.getPosition();
123+
int maxSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
124+
int firstStart = firstOutput.getPosition();
125+
int secondStart = secondOutput.getPosition();
120126
int maxBatchCount = messageSettings.getMaxBatchCount();
121-
return opsAndNsInfo.encode(write -> {
122-
int opsBeforeWritePosition = opsOut.getPosition();
123-
int nsInfoBeforeWritePosition = nsInfoOut.getPosition();
124-
int batchCountAfterWrite = write.doAndGetBatchCount(opsWriters, nsInfoWriter);
127+
return dualMessageSequences.encodeDocuments(write -> {
128+
int firstBeforeWritePosition = firstOutput.getPosition();
129+
int secondBeforeWritePosition = secondOutput.getPosition();
130+
int batchCountAfterWrite = write.doAndGetBatchCount(firstWriters, secondWriters);
125131
assertTrue(batchCountAfterWrite <= maxBatchCount);
126-
int opsAndNsInfoSizeInBytes =
127-
opsOut.getPosition() - opsStart
128-
+ nsInfoOut.getPosition() - nsInfoStart;
129-
if (opsAndNsInfoSizeInBytes < maxOpsAndNsInfoSizeInBytes && batchCountAfterWrite < maxBatchCount) {
130-
return WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_NOT_REACHED;
131-
} else if (opsAndNsInfoSizeInBytes > maxOpsAndNsInfoSizeInBytes) {
132-
opsOut.truncateToPosition(opsBeforeWritePosition);
133-
nsInfoOut.truncateToPosition(nsInfoBeforeWritePosition);
132+
int writtenSizeInBytes =
133+
firstOutput.getPosition() - firstStart
134+
+ secondOutput.getPosition() - secondStart;
135+
if (writtenSizeInBytes < maxSizeInBytes && batchCountAfterWrite < maxBatchCount) {
136+
return OK_LIMIT_NOT_REACHED;
137+
} else if (writtenSizeInBytes > maxSizeInBytes) {
138+
firstOutput.truncateToPosition(firstBeforeWritePosition);
139+
secondOutput.truncateToPosition(secondBeforeWritePosition);
134140
if (batchCountAfterWrite == 1) {
135-
// we have failed to write a single model
141+
// we have failed to write a single document
136142
throw createBsonMaximumSizeExceededException(messageSettings.getMaxDocumentSize());
137143
}
138-
return WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
144+
return FAIL_LIMIT_EXCEEDED;
139145
} else {
140-
return WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_REACHED;
146+
return OK_LIMIT_REACHED;
141147
}
142148
});
143149
}
@@ -233,14 +239,14 @@ private static boolean exceedsLimits(final MessageSettings settings, final int m
233239
private BsonWriterHelper() {
234240
}
235241

236-
private static final class BinaryOpsBsonWriters implements WritersProviderAndLimitsChecker.OpsBsonWriters {
242+
private static final class BinaryOrdinaryAndStoredBsonWriters implements WritersProviderAndLimitsChecker.OrdinaryAndStoredBsonWriters {
237243
private final BsonBinaryWriter writer;
238244
private final BsonWriter storedDocumentWriter;
239245

240246
/**
241247
* @param messageSettings Non-{@code null} iff the document size limits must be validated.
242248
*/
243-
BinaryOpsBsonWriters(
249+
BinaryOrdinaryAndStoredBsonWriters(
244250
final BsonOutput out,
245251
final FieldNameValidator validator,
246252
@Nullable final MessageSettings messageSettings) {

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

+27-28
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import com.mongodb.connection.ClusterConnectionMode;
2525
import com.mongodb.internal.TimeoutContext;
2626
import com.mongodb.internal.connection.MessageSequences.EmptyMessageSequences;
27-
import com.mongodb.internal.operation.ClientBulkWriteOperation.ClientBulkWriteCommand;
2827
import com.mongodb.internal.session.SessionContext;
2928
import com.mongodb.lang.Nullable;
3029
import org.bson.BsonArray;
@@ -57,7 +56,7 @@
5756
import static com.mongodb.connection.ServerType.STANDALONE;
5857
import static com.mongodb.internal.connection.BsonWriterHelper.appendElementsToDocument;
5958
import static com.mongodb.internal.connection.BsonWriterHelper.backpatchLength;
60-
import static com.mongodb.internal.connection.BsonWriterHelper.writeOpsAndNsInfo;
59+
import static com.mongodb.internal.connection.BsonWriterHelper.writeDocumentsOfDualMessageSequences;
6160
import static com.mongodb.internal.connection.BsonWriterHelper.writePayload;
6261
import static com.mongodb.internal.connection.ByteBufBsonDocument.createList;
6362
import static com.mongodb.internal.connection.ByteBufBsonDocument.createOne;
@@ -81,11 +80,11 @@ public final class CommandMessage extends RequestMessage {
8180
private final MessageSequences sequences;
8281
private final boolean responseExpected;
8382
/**
84-
* {@code null} iff either {@link #sequences} is not of the {@link ClientBulkWriteCommand.OpsAndNsInfo} type,
83+
* {@code null} iff either {@link #sequences} is not of the {@link DualMessageSequences} type,
8584
* or it is of that type, but it has not been {@linkplain #encodeMessageBodyWithMetadata(ByteBufferBsonOutput, OperationContext) encoded}.
8685
*/
8786
@Nullable
88-
private Boolean opsAndNsInfoRequireResponse;
87+
private Boolean dualMessageSequencesRequireResponse;
8988
private final ClusterConnectionMode clusterConnectionMode;
9089
private final ServerApi serverApi;
9190

@@ -122,7 +121,7 @@ public final class CommandMessage extends RequestMessage {
122121
this.commandFieldNameValidator = commandFieldNameValidator;
123122
this.readPreference = readPreference;
124123
this.responseExpected = responseExpected;
125-
opsAndNsInfoRequireResponse = null;
124+
dualMessageSequencesRequireResponse = null;
126125
this.exhaustAllowed = exhaustAllowed;
127126
this.sequences = sequences;
128127
this.clusterConnectionMode = notNull("clusterConnectionMode", clusterConnectionMode);
@@ -206,12 +205,11 @@ boolean isResponseExpected() {
206205
if (responseExpected) {
207206
return true;
208207
} else {
209-
if (sequences instanceof ValidatableSplittablePayload) {
210-
ValidatableSplittablePayload validatableSplittablePayload = (ValidatableSplittablePayload) sequences;
211-
SplittablePayload payload = validatableSplittablePayload.getSplittablePayload();
208+
if (sequences instanceof SplittablePayload) {
209+
SplittablePayload payload = (SplittablePayload) sequences;
212210
return payload.isOrdered() && payload.hasAnotherSplit();
213-
} else if (sequences instanceof ClientBulkWriteCommand.OpsAndNsInfo) {
214-
return assertNotNull(opsAndNsInfoRequireResponse);
211+
} else if (sequences instanceof DualMessageSequences) {
212+
return assertNotNull(dualMessageSequencesRequireResponse);
215213
} else if (!(sequences instanceof EmptyMessageSequences)) {
216214
fail(sequences.toString());
217215
}
@@ -233,38 +231,34 @@ protected EncodingMetadata encodeMessageBodyWithMetadata(final ByteBufferBsonOut
233231
bsonOutput.writeByte(0); // payload type
234232
commandStartPosition = bsonOutput.getPosition();
235233
ArrayList<BsonElement> extraElements = getExtraElements(operationContext);
236-
// `OpsAndNsInfo` requires validation only if no response is expected, otherwise we must rely on the server validation
237-
boolean validateDocumentSizeLimits = !(sequences instanceof ClientBulkWriteCommand.OpsAndNsInfo) || !responseExpected;
234+
// `DualMessageSequences` requires validation only if no response is expected, otherwise we must rely on the server validation
235+
boolean validateDocumentSizeLimits = !(sequences instanceof DualMessageSequences) || !responseExpected;
238236

239237
int commandDocumentSizeInBytes = writeDocument(command, bsonOutput, commandFieldNameValidator, validateDocumentSizeLimits);
240-
if (sequences instanceof ValidatableSplittablePayload) {
238+
if (sequences instanceof SplittablePayload) {
241239
appendElementsToDocument(bsonOutput, commandStartPosition, extraElements);
242-
ValidatableSplittablePayload validatableSplittablePayload = (ValidatableSplittablePayload) sequences;
243-
SplittablePayload payload = validatableSplittablePayload.getSplittablePayload();
240+
SplittablePayload payload = (SplittablePayload) sequences;
244241
writeOpMsgSectionWithPayloadType1(bsonOutput, payload.getPayloadName(), () -> {
245242
writePayload(
246-
new BsonBinaryWriter(bsonOutput, validatableSplittablePayload.getFieldNameValidator()),
243+
new BsonBinaryWriter(bsonOutput, payload.getFieldNameValidator()),
247244
bsonOutput, getSettings(), messageStartPosition, payload, getSettings().getMaxDocumentSize()
248245
);
249246
return null;
250247
});
251-
} else if (sequences instanceof ClientBulkWriteCommand.OpsAndNsInfo) {
252-
ClientBulkWriteCommand.OpsAndNsInfo opsAndNsInfo = (ClientBulkWriteCommand.OpsAndNsInfo) sequences;
248+
} else if (sequences instanceof DualMessageSequences) {
249+
DualMessageSequences dualMessageSequences = (DualMessageSequences) sequences;
253250
try (ByteBufferBsonOutput.Branch bsonOutputBranch2 = bsonOutput.branch();
254251
ByteBufferBsonOutput.Branch bsonOutputBranch1 = bsonOutput.branch()) {
255-
ClientBulkWriteCommand.OpsAndNsInfo.EncodeResult opsAndNsInfoEncodeResult = writeOpMsgSectionWithPayloadType1(
256-
bsonOutputBranch1, "ops", () ->
257-
writeOpMsgSectionWithPayloadType1(bsonOutputBranch2, "nsInfo", () ->
258-
writeOpsAndNsInfo(
259-
opsAndNsInfo, commandDocumentSizeInBytes, bsonOutputBranch1,
252+
DualMessageSequences.EncodeDocumentsResult encodeDocumentsResult = writeOpMsgSectionWithPayloadType1(
253+
bsonOutputBranch1, dualMessageSequences.getFirstSequenceId(), () ->
254+
writeOpMsgSectionWithPayloadType1(bsonOutputBranch2, dualMessageSequences.getSecondSequenceId(), () ->
255+
writeDocumentsOfDualMessageSequences(
256+
dualMessageSequences, commandDocumentSizeInBytes, bsonOutputBranch1,
260257
bsonOutputBranch2, getSettings(), validateDocumentSizeLimits)
261258
)
262259
);
263-
opsAndNsInfoRequireResponse = opsAndNsInfoEncodeResult.isServerResponseRequired();
264-
Long txnNumber = opsAndNsInfoEncodeResult.getTxnNumber();
265-
if (txnNumber != null) {
266-
extraElements.add(new BsonElement(TXN_NUMBER_KEY, new BsonInt64(txnNumber)));
267-
}
260+
dualMessageSequencesRequireResponse = encodeDocumentsResult.isServerResponseRequired();
261+
extraElements.addAll(encodeDocumentsResult.getExtraElements());
268262
appendElementsToDocument(bsonOutput, commandStartPosition, extraElements);
269263
}
270264
} else if (sequences instanceof EmptyMessageSequences) {
@@ -391,6 +385,11 @@ private void addReadConcernDocument(final List<BsonElement> extraElements, final
391385
}
392386
}
393387

388+
/**
389+
* @param sequenceId The identifier of the sequence contained in the {@code OP_MSG} section to be written.
390+
* @param writeDocumentsAction The action that writes the documents of the sequence.
391+
* @see <a href="https://github.com/mongodb/specifications/blob/master/source/message/OP_MSG.md">OP_MSG</a>
392+
*/
394393
private <R> R writeOpMsgSectionWithPayloadType1(
395394
final ByteBufferBsonOutput bsonOutput,
396395
final String sequenceId,

0 commit comments

Comments
 (0)