Skip to content

Commit 7d7f3eb

Browse files
jyeminjyemin
committed
Remove OpsBsonWriters
Move the logic for determining whether the document limit needs to be applied, as well as the application of it, into ClientBulkWriteOperation. This required a few small changes, including: * Retain and rename StoredDocumentSizeLimitCheckingBsonBinaryWriter to DocumentSizeLimitCheckingBsonBinaryWriter, and use it in the encoder directly when necessary. It might make sense to move DocumentSizeLimitCheckingBsonBinaryWriter into ClientBulkWriteOperation. In order to do this, the BatchEncoder now has to know whether it should apply the limit checking, and for that it requires both the maxDocumentLength, which is accessible from the ConnectionDescription, and whether the write is acknowledged, which is accessible from the WriteConcer. * Pass BsonBinaryWriter instead of BsonWriter to the encoder so that it can be wrapped if necessary. This is a bit unpleasant but seems ok, since what other kind of BsonWriter would make sense inside of Connection?
1 parent 9f1f9ac commit 7d7f3eb

File tree

4 files changed

+59
-92
lines changed

4 files changed

+59
-92
lines changed

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

Lines changed: 19 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818

1919
import com.mongodb.internal.connection.DualSplittablePayloads.EncodeResult;
2020
import com.mongodb.internal.connection.DualSplittablePayloads.WritersProviderAndLimitsChecker;
21-
import com.mongodb.internal.connection.DualSplittablePayloads.WritersProviderAndLimitsChecker.OpsBsonWriters;
2221
import com.mongodb.internal.validator.NoOpFieldNameValidator;
2322
import com.mongodb.lang.Nullable;
2423
import org.bson.BsonBinaryWriter;
@@ -109,12 +108,8 @@ static EncodeResult writeDualSplittablePayloads(
109108
final int commandDocumentSizeInBytes,
110109
final BsonOutput firstOutput,
111110
final BsonOutput secondOutput,
112-
final MessageSettings messageSettings,
113-
final boolean validateDocumentSizeLimits) {
114-
BinaryOpsBsonWriters opsWriters = new BinaryOpsBsonWriters(
115-
firstOutput,
116-
dualSplittablePayloads.getFirstFieldNameValidator(),
117-
validateDocumentSizeLimits ? messageSettings : null);
111+
final MessageSettings messageSettings) {
112+
BsonBinaryWriter firstWriter = new BsonBinaryWriter(firstOutput, dualSplittablePayloads.getFirstFieldNameValidator());
118113
BsonBinaryWriter secondWriter = new BsonBinaryWriter(secondOutput, dualSplittablePayloads.getSecondFieldNameValidator());
119114
// the size of operation-agnostic command fields (a.k.a. extra elements) is counted towards `messageOverheadInBytes`
120115
int messageOverheadInBytes = 1000;
@@ -125,7 +120,7 @@ static EncodeResult writeDualSplittablePayloads(
125120
return dualSplittablePayloads.encode(write -> {
126121
int firstBeforeWritePosition = firstOutput.getPosition();
127122
int secondBeforeWritePosition = secondOutput.getPosition();
128-
int batchCountAfterWrite = write.doAndGetBatchCount(opsWriters, secondWriter);
123+
int batchCountAfterWrite = write.doAndGetBatchCount(firstWriter, secondWriter);
129124
assertTrue(batchCountAfterWrite <= maxBatchCount);
130125
int totalSizeInBytes =
131126
firstOutput.getPosition() - firstStart
@@ -237,57 +232,29 @@ private static boolean exceedsLimits(final MessageSettings settings, final int m
237232
private BsonWriterHelper() {
238233
}
239234

240-
private static final class BinaryOpsBsonWriters implements OpsBsonWriters {
241-
private final BsonBinaryWriter writer;
242-
private final BsonWriter storedDocumentWriter;
235+
public static final class DocumentSizeLimitCheckingBsonBinaryWriter extends BsonWriterDecorator {
236+
private final int maxStoredDocumentSize;
237+
private final BsonOutput out;
238+
private int documentStart;
243239

244-
/**
245-
* @param messageSettings Non-{@code null} iff the document size limits must be validated.
246-
*/
247-
BinaryOpsBsonWriters(
248-
final BsonOutput out,
249-
final FieldNameValidator validator,
250-
@Nullable final MessageSettings messageSettings) {
251-
writer = createBsonBinaryWriter(out, validator, messageSettings);
252-
storedDocumentWriter = messageSettings == null
253-
? writer
254-
: new StoredDocumentSizeLimitCheckingBsonBinaryWriter(writer, messageSettings.getMaxDocumentSize());
240+
public DocumentSizeLimitCheckingBsonBinaryWriter(final BsonBinaryWriter writer, final int maxStoredDocumentSize) {
241+
super(writer);
242+
this.maxStoredDocumentSize = maxStoredDocumentSize;
243+
this.out = writer.getBsonOutput();
255244
}
256245

257246
@Override
258-
public BsonWriter getWriter() {
259-
return writer;
247+
public void writeStartDocument() {
248+
documentStart = out.getPosition();
249+
super.writeStartDocument();
260250
}
261251

262252
@Override
263-
public BsonWriter getStoredDocumentWriter() {
264-
return storedDocumentWriter;
265-
}
266-
267-
private static final class StoredDocumentSizeLimitCheckingBsonBinaryWriter extends BsonWriterDecorator {
268-
private final int maxStoredDocumentSize;
269-
private final BsonOutput out;
270-
private int documentStart;
271-
272-
StoredDocumentSizeLimitCheckingBsonBinaryWriter(final BsonBinaryWriter writer, final int maxStoredDocumentSize) {
273-
super(writer);
274-
this.maxStoredDocumentSize = maxStoredDocumentSize;
275-
this.out = writer.getBsonOutput();
276-
}
277-
278-
@Override
279-
public void writeStartDocument() {
280-
documentStart = out.getPosition();
281-
super.writeStartDocument();
282-
}
283-
284-
@Override
285-
public void writeEndDocument() throws BsonMaximumSizeExceededException {
286-
super.writeEndDocument();
287-
int documentSize = out.getPosition() - documentStart;
288-
if (documentSize > maxStoredDocumentSize) {
289-
throw createBsonMaximumSizeExceededException(maxStoredDocumentSize);
290-
}
253+
public void writeEndDocument() throws BsonMaximumSizeExceededException {
254+
super.writeEndDocument();
255+
int documentSize = out.getPosition() - documentStart;
256+
if (documentSize > maxStoredDocumentSize) {
257+
throw createBsonMaximumSizeExceededException(maxStoredDocumentSize);
291258
}
292259
}
293260
}

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.mongodb.internal.TimeoutContext;
2626
import com.mongodb.internal.connection.DualSplittablePayloads.EncodeResult;
2727
import com.mongodb.internal.connection.OpMsgSequences.EmptyOpMsgSequences;
28-
import com.mongodb.internal.operation.ClientBulkWriteOperation.ClientBulkWriteCommand;
2928
import com.mongodb.internal.session.SessionContext;
3029
import com.mongodb.lang.Nullable;
3130
import org.bson.BsonArray;
@@ -234,10 +233,8 @@ protected EncodingMetadata encodeMessageBodyWithMetadata(final ByteBufferBsonOut
234233
bsonOutput.writeByte(0); // payload type
235234
commandStartPosition = bsonOutput.getPosition();
236235
ArrayList<BsonElement> extraElements = getExtraElements(operationContext);
237-
// `DualSplittablePayloads` requires validation only if no response is expected, otherwise we must rely on the server validation
238-
boolean validateDocumentSizeLimits = !(sequences instanceof DualSplittablePayloads) || !responseExpected;
239236

240-
int commandDocumentSizeInBytes = writeDocument(command, bsonOutput, commandFieldNameValidator, validateDocumentSizeLimits);
237+
int commandDocumentSizeInBytes = writeDocument(command, bsonOutput, commandFieldNameValidator, true);
241238
if (sequences instanceof ValidatableSplittablePayload) {
242239
appendElementsToDocument(bsonOutput, commandStartPosition, extraElements);
243240
ValidatableSplittablePayload validatableSplittablePayload = (ValidatableSplittablePayload) sequences;
@@ -258,7 +255,7 @@ bsonOutput, getSettings(), messageStartPosition, payload, getSettings().getMaxDo
258255
writeOpMsgSectionWithPayloadType1(bsonOutputBranch2, dualSplittablePayloads.getSecondSequenceId(), () ->
259256
writeDualSplittablePayloads(
260257
dualSplittablePayloads, commandDocumentSizeInBytes, bsonOutputBranch1,
261-
bsonOutputBranch2, getSettings(), validateDocumentSizeLimits)
258+
bsonOutputBranch2, getSettings())
262259
)
263260
);
264261
dualSplittablePayloadsRequireResponse = encodeResult.isServerResponseRequired();

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

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,9 @@
1616

1717
package com.mongodb.internal.connection;
1818

19+
import org.bson.BsonBinaryWriter;
1920
import org.bson.BsonElement;
20-
import org.bson.BsonWriter;
2121
import org.bson.FieldNameValidator;
22-
import org.bson.io.BsonOutput;
2322

2423
import java.util.List;
2524

@@ -65,30 +64,20 @@ public String getSecondSequenceId() {
6564
public interface WritersProviderAndLimitsChecker {
6665
/**
6766
* Provides writers to the specified {@link WriteAction},
68-
* {@linkplain WriteAction#doAndGetBatchCount(OpsBsonWriters, BsonWriter) executes} it,
67+
* {@linkplain WriteAction#doAndGetBatchCount(BsonBinaryWriter, BsonBinaryWriter) executes} it,
6968
* checks the {@linkplain MessageSettings limits}.
7069
*/
7170
WriteResult tryWrite(WriteAction write);
7271

7372
/**
74-
* @see #doAndGetBatchCount(OpsBsonWriters, BsonWriter)
73+
* @see #doAndGetBatchCount(BsonBinaryWriter, BsonBinaryWriter)
7574
*/
7675
interface WriteAction {
7776
/**.
7877
*
7978
* @return The resulting batch count}.
8079
*/
81-
int doAndGetBatchCount(OpsBsonWriters firstWriter, BsonWriter secondWriter);
82-
}
83-
84-
interface OpsBsonWriters {
85-
BsonWriter getWriter();
86-
87-
/**
88-
* A {@link BsonWriter} to use for writing documents that are intended to be stored in a database.
89-
* Must write to the same {@linkplain BsonOutput output} as {@link #getWriter()} does.
90-
*/
91-
BsonWriter getStoredDocumentWriter();
80+
int doAndGetBatchCount(BsonBinaryWriter firstWriter, BsonBinaryWriter secondWriter);
9281
}
9382

9483
enum WriteResult {

0 commit comments

Comments
 (0)