Skip to content

Commit b8734d1

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 b8734d1

File tree

14 files changed

+264
-231
lines changed

14 files changed

+264
-231
lines changed

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

+34-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,46 @@ 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(),
113116
validateDocumentSizeLimits ? messageSettings : null);
114-
BsonBinaryWriter nsInfoWriter = new BsonBinaryWriter(nsInfoOut, NoOpFieldNameValidator.INSTANCE);
117+
BsonBinaryWriter secondWriter = new BsonBinaryWriter(secondOutput, dualMessageSequences.getSecondFieldNameValidator());
115118
// the size of operation-agnostic command fields (a.k.a. extra elements) is counted towards `messageOverheadInBytes`
116119
int messageOverheadInBytes = 1000;
117-
int maxOpsAndNsInfoSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
118-
int opsStart = opsOut.getPosition();
119-
int nsInfoStart = nsInfoOut.getPosition();
120+
int maxSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
121+
int firstStart = firstOutput.getPosition();
122+
int secondStart = secondOutput.getPosition();
120123
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);
124+
return dualMessageSequences.encodeDocuments(write -> {
125+
int firstBeforeWritePosition = firstOutput.getPosition();
126+
int secondBeforeWritePosition = secondOutput.getPosition();
127+
int batchCountAfterWrite = write.doAndGetBatchCount(firstWriters, secondWriter);
125128
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);
129+
int writtenSizeInBytes =
130+
firstOutput.getPosition() - firstStart
131+
+ secondOutput.getPosition() - secondStart;
132+
if (writtenSizeInBytes < maxSizeInBytes && batchCountAfterWrite < maxBatchCount) {
133+
return OK_LIMIT_NOT_REACHED;
134+
} else if (writtenSizeInBytes > maxSizeInBytes) {
135+
firstOutput.truncateToPosition(firstBeforeWritePosition);
136+
secondOutput.truncateToPosition(secondBeforeWritePosition);
134137
if (batchCountAfterWrite == 1) {
135-
// we have failed to write a single model
138+
// we have failed to write a single document
136139
throw createBsonMaximumSizeExceededException(messageSettings.getMaxDocumentSize());
137140
}
138-
return WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
141+
return FAIL_LIMIT_EXCEEDED;
139142
} else {
140-
return WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_REACHED;
143+
return OK_LIMIT_REACHED;
141144
}
142145
});
143146
}
@@ -233,14 +236,14 @@ private static boolean exceedsLimits(final MessageSettings settings, final int m
233236
private BsonWriterHelper() {
234237
}
235238

236-
private static final class BinaryOpsBsonWriters implements WritersProviderAndLimitsChecker.OpsBsonWriters {
239+
private static final class BinaryOrdinaryAndStoredBsonWriters implements WritersProviderAndLimitsChecker.OrdinaryAndStoredBsonWriters {
237240
private final BsonBinaryWriter writer;
238241
private final BsonWriter storedDocumentWriter;
239242

240243
/**
241244
* @param messageSettings Non-{@code null} iff the document size limits must be validated.
242245
*/
243-
BinaryOpsBsonWriters(
246+
BinaryOrdinaryAndStoredBsonWriters(
244247
final BsonOutput out,
245248
final FieldNameValidator validator,
246249
@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,

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

+134
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.mongodb.internal.connection;
18+
19+
import org.bson.BsonElement;
20+
import org.bson.BsonWriter;
21+
import org.bson.FieldNameValidator;
22+
import org.bson.io.BsonOutput;
23+
24+
import java.util.List;
25+
26+
/**
27+
* Two sequences that may either be coupled or independent.
28+
* <p>
29+
* This class is not part of the public API and may be removed or changed at any time.</p>
30+
*/
31+
public abstract class DualMessageSequences extends MessageSequences {
32+
33+
private final String firstSequenceId;
34+
private final FieldNameValidator firstFieldNameValidator;
35+
private final String secondSequenceId;
36+
private final FieldNameValidator secondFieldNameValidator;
37+
38+
protected DualMessageSequences(
39+
final String firstSequenceId,
40+
final FieldNameValidator firstFieldNameValidator,
41+
final String secondSequenceId,
42+
final FieldNameValidator secondFieldNameValidator) {
43+
this.firstSequenceId = firstSequenceId;
44+
this.firstFieldNameValidator = firstFieldNameValidator;
45+
this.secondSequenceId = secondSequenceId;
46+
this.secondFieldNameValidator = secondFieldNameValidator;
47+
}
48+
49+
FieldNameValidator getFirstFieldNameValidator() {
50+
return firstFieldNameValidator;
51+
}
52+
53+
FieldNameValidator getSecondFieldNameValidator() {
54+
return secondFieldNameValidator;
55+
}
56+
57+
String getFirstSequenceId() {
58+
return firstSequenceId;
59+
}
60+
61+
String getSecondSequenceId() {
62+
return secondSequenceId;
63+
}
64+
65+
protected abstract EncodeDocumentsResult encodeDocuments(WritersProviderAndLimitsChecker writersProviderAndLimitsChecker);
66+
67+
/**
68+
* @see #tryWrite(WriteAction)
69+
*/
70+
public interface WritersProviderAndLimitsChecker {
71+
/**
72+
* Provides writers to the specified {@link WriteAction},
73+
* {@linkplain WriteAction#doAndGetBatchCount(OrdinaryAndStoredBsonWriters, BsonWriter) executes} it,
74+
* checks the {@linkplain MessageSettings limits}.
75+
* <p>
76+
* May be called multiple times per {@link #encodeDocuments(WritersProviderAndLimitsChecker)}.</p>
77+
*/
78+
WriteResult tryWrite(WriteAction write);
79+
80+
/**
81+
* @see #doAndGetBatchCount(OrdinaryAndStoredBsonWriters, BsonWriter)
82+
*/
83+
interface WriteAction {
84+
/**
85+
* Writes documents to the sequences using the provided writers.
86+
*
87+
* @return The resulting batch count since the beginning of {@link #encodeDocuments(WritersProviderAndLimitsChecker)}.
88+
* It is generally allowed to be greater than {@link MessageSettings#getMaxBatchCount()}.
89+
*/
90+
// VAKOTODO pass OrdinaryAndStoredBsonWriters for both first and second?
91+
int doAndGetBatchCount(OrdinaryAndStoredBsonWriters firstWriter, BsonWriter secondWriter);
92+
}
93+
94+
interface OrdinaryAndStoredBsonWriters {
95+
BsonWriter getWriter();
96+
97+
/**
98+
* A {@link BsonWriter} to use for writing documents that are intended to be stored in a database.
99+
* Must write to the same {@linkplain BsonOutput output} as {@link #getWriter()} does.
100+
*/
101+
BsonWriter getStoredDocumentWriter();
102+
}
103+
104+
enum WriteResult {
105+
FAIL_LIMIT_EXCEEDED,
106+
OK_LIMIT_REACHED,
107+
OK_LIMIT_NOT_REACHED
108+
}
109+
}
110+
111+
public static final class EncodeDocumentsResult {
112+
private final boolean serverResponseRequired;
113+
private final List<BsonElement> extraElements;
114+
115+
/**
116+
* @param extraElements See {@link #getExtraElements()}.
117+
*/
118+
public EncodeDocumentsResult(final boolean serverResponseRequired, final List<BsonElement> extraElements) {
119+
this.serverResponseRequired = serverResponseRequired;
120+
this.extraElements = extraElements;
121+
}
122+
123+
boolean isServerResponseRequired() {
124+
return serverResponseRequired;
125+
}
126+
127+
/**
128+
* {@linkplain BsonElement Key/value pairs} to be added to the document contained in the {@code OP_MSG} section with payload type 0.
129+
*/
130+
List<BsonElement> getExtraElements() {
131+
return extraElements;
132+
}
133+
}
134+
}

0 commit comments

Comments
 (0)