Skip to content

Commit 056d411

Browse files
stIncMalejyemin
andcommitted
Get rid of OrdinaryAndStoredBsonWriters
JAVA-5529 Co-authored-by: Jeff Yemin <jeff.yemin@mongodb.com>
1 parent 9fe35e4 commit 056d411

File tree

3 files changed

+70
-89
lines changed

3 files changed

+70
-89
lines changed

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

Lines changed: 31 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -110,14 +110,10 @@ static EncodeDocumentsResult writeDocumentsOfDualMessageSequences(
110110
final BsonOutput secondOutput,
111111
final MessageSettings messageSettings,
112112
final boolean validateDocumentSizeLimits) {
113-
BinaryOrdinaryAndStoredBsonWriters firstWriters = new BinaryOrdinaryAndStoredBsonWriters(
114-
firstOutput,
115-
dualMessageSequences.getFirstFieldNameValidator(),
116-
validateDocumentSizeLimits ? messageSettings : null);
117-
BinaryOrdinaryAndStoredBsonWriters secondWriters = new BinaryOrdinaryAndStoredBsonWriters(
118-
secondOutput,
119-
dualMessageSequences.getSecondFieldNameValidator(),
120-
validateDocumentSizeLimits ? messageSettings : null);
113+
BsonBinaryWriter firstWriter = createBsonBinaryWriter(
114+
firstOutput, dualMessageSequences.getFirstFieldNameValidator(), validateDocumentSizeLimits ? messageSettings : null);
115+
BsonBinaryWriter secondWriter = createBsonBinaryWriter(
116+
secondOutput, dualMessageSequences.getSecondFieldNameValidator(), validateDocumentSizeLimits ? messageSettings : null);
121117
// the size of operation-agnostic command fields (a.k.a. extra elements) is counted towards `messageOverheadInBytes`
122118
int messageOverheadInBytes = 1000;
123119
int maxSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
@@ -127,7 +123,7 @@ static EncodeDocumentsResult writeDocumentsOfDualMessageSequences(
127123
return dualMessageSequences.encodeDocuments(write -> {
128124
int firstBeforeWritePosition = firstOutput.getPosition();
129125
int secondBeforeWritePosition = secondOutput.getPosition();
130-
int batchCountAfterWrite = write.doAndGetBatchCount(firstWriters, secondWriters);
126+
int batchCountAfterWrite = write.doAndGetBatchCount(firstWriter, secondWriter);
131127
assertTrue(batchCountAfterWrite <= maxBatchCount);
132128
int writtenSizeInBytes =
133129
firstOutput.getPosition() - firstStart
@@ -235,64 +231,43 @@ private static boolean exceedsLimits(final MessageSettings settings, final int m
235231
return false;
236232
}
237233

234+
/**
235+
* @return {@code writer} if {@code maxDocumentSize} is {@code null}, otherwise decorates it.
236+
*/
237+
public static BsonWriter decorateWithDocumentSizeChecking(final BsonBinaryWriter writer, @Nullable final Integer maxDocumentSize) {
238+
return maxDocumentSize == null ? writer : new DocumentSizeLimitCheckingBsonBinaryWriter(writer, maxDocumentSize);
239+
}
238240

239241
private BsonWriterHelper() {
240242
}
241243

242-
private static final class BinaryOrdinaryAndStoredBsonWriters implements WritersProviderAndLimitsChecker.OrdinaryAndStoredBsonWriters {
243-
private final BsonBinaryWriter writer;
244-
private final BsonWriter storedDocumentWriter;
244+
private static final class DocumentSizeLimitCheckingBsonBinaryWriter extends LevelCountingBsonWriter {
245+
private final int maxStoredDocumentSize;
246+
private final BsonOutput out;
247+
private int documentStart;
245248

246-
/**
247-
* @param messageSettings Non-{@code null} iff the document size limits must be validated.
248-
*/
249-
BinaryOrdinaryAndStoredBsonWriters(
250-
final BsonOutput out,
251-
final FieldNameValidator validator,
252-
@Nullable final MessageSettings messageSettings) {
253-
writer = createBsonBinaryWriter(out, validator, messageSettings);
254-
storedDocumentWriter = messageSettings == null
255-
? writer
256-
: new StoredDocumentSizeLimitCheckingBsonBinaryWriter(writer, messageSettings.getMaxDocumentSize());
249+
DocumentSizeLimitCheckingBsonBinaryWriter(final BsonBinaryWriter writer, final int maxStoredDocumentSize) {
250+
super(writer);
251+
assertTrue(maxStoredDocumentSize > 0);
252+
this.maxStoredDocumentSize = maxStoredDocumentSize;
253+
this.out = writer.getBsonOutput();
257254
}
258255

259256
@Override
260-
public BsonWriter getWriter() {
261-
return writer;
257+
public void writeStartDocument() {
258+
if (getCurrentLevel() == INITIAL_LEVEL) {
259+
documentStart = out.getPosition();
260+
}
261+
super.writeStartDocument();
262262
}
263263

264264
@Override
265-
public BsonWriter getStoredDocumentWriter() {
266-
return storedDocumentWriter;
267-
}
268-
269-
private static final class StoredDocumentSizeLimitCheckingBsonBinaryWriter extends LevelCountingBsonWriter {
270-
private final int maxStoredDocumentSize;
271-
private final BsonOutput out;
272-
private int documentStart;
273-
274-
StoredDocumentSizeLimitCheckingBsonBinaryWriter(final BsonBinaryWriter writer, final int maxStoredDocumentSize) {
275-
super(writer);
276-
this.maxStoredDocumentSize = maxStoredDocumentSize;
277-
this.out = writer.getBsonOutput();
278-
}
279-
280-
@Override
281-
public void writeStartDocument() {
282-
if (getCurrentLevel() == INITIAL_LEVEL) {
283-
documentStart = out.getPosition();
284-
}
285-
super.writeStartDocument();
286-
}
287-
288-
@Override
289-
public void writeEndDocument() throws BsonMaximumSizeExceededException {
290-
super.writeEndDocument();
291-
if (getCurrentLevel() == INITIAL_LEVEL) {
292-
int documentSize = out.getPosition() - documentStart;
293-
if (documentSize > maxStoredDocumentSize) {
294-
throw createBsonMaximumSizeExceededException(maxStoredDocumentSize);
295-
}
265+
public void writeEndDocument() throws BsonMaximumSizeExceededException {
266+
super.writeEndDocument();
267+
if (getCurrentLevel() == INITIAL_LEVEL) {
268+
int documentSize = out.getPosition() - documentStart;
269+
if (documentSize > maxStoredDocumentSize) {
270+
throw createBsonMaximumSizeExceededException(maxStoredDocumentSize);
296271
}
297272
}
298273
}

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

Lines changed: 4 additions & 16 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

@@ -70,15 +69,15 @@ String getSecondSequenceId() {
7069
public interface WritersProviderAndLimitsChecker {
7170
/**
7271
* Provides writers to the specified {@link WriteAction},
73-
* {@linkplain WriteAction#doAndGetBatchCount(OrdinaryAndStoredBsonWriters, OrdinaryAndStoredBsonWriters) executes} it,
72+
* {@linkplain WriteAction#doAndGetBatchCount(BsonBinaryWriter, BsonBinaryWriter) executes} it,
7473
* checks the {@linkplain MessageSettings limits}.
7574
* <p>
7675
* May be called multiple times per {@link #encodeDocuments(WritersProviderAndLimitsChecker)}.</p>
7776
*/
7877
WriteResult tryWrite(WriteAction write);
7978

8079
/**
81-
* @see #doAndGetBatchCount(OrdinaryAndStoredBsonWriters, OrdinaryAndStoredBsonWriters)
80+
* @see #doAndGetBatchCount(BsonBinaryWriter, BsonBinaryWriter)
8281
*/
8382
interface WriteAction {
8483
/**
@@ -87,18 +86,7 @@ interface WriteAction {
8786
* @return The resulting batch count since the beginning of {@link #encodeDocuments(WritersProviderAndLimitsChecker)}.
8887
* It is generally allowed to be greater than {@link MessageSettings#getMaxBatchCount()}.
8988
*/
90-
int doAndGetBatchCount(OrdinaryAndStoredBsonWriters firstWriter, OrdinaryAndStoredBsonWriters secondWriter);
91-
}
92-
93-
interface OrdinaryAndStoredBsonWriters {
94-
BsonWriter getWriter();
95-
96-
/**
97-
* A {@link BsonWriter} to use for writing documents that are intended to be stored in a database.
98-
* Must write to the same {@linkplain BsonOutput output} as {@link #getWriter()} does,
99-
* and may be the same as {@link #getWriter()}.
100-
*/
101-
BsonWriter getStoredDocumentWriter();
89+
int doAndGetBatchCount(BsonBinaryWriter firstWriter, BsonBinaryWriter secondWriter);
10290
}
10391

10492
enum WriteResult {

driver-core/src/main/com/mongodb/internal/operation/ClientBulkWriteOperation.java

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,6 @@
7070
import com.mongodb.internal.client.model.bulk.UnacknowledgedClientBulkWriteResult;
7171
import com.mongodb.internal.connection.Connection;
7272
import com.mongodb.internal.connection.DualMessageSequences;
73-
import com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.OrdinaryAndStoredBsonWriters;
7473
import com.mongodb.internal.connection.IdHoldingBsonWriter;
7574
import com.mongodb.internal.connection.MongoWriteConcernWithResponseException;
7675
import com.mongodb.internal.connection.OperationContext;
@@ -81,6 +80,7 @@
8180
import com.mongodb.internal.validator.UpdateFieldNameValidator;
8281
import com.mongodb.lang.Nullable;
8382
import org.bson.BsonArray;
83+
import org.bson.BsonBinaryWriter;
8484
import org.bson.BsonBoolean;
8585
import org.bson.BsonDocument;
8686
import org.bson.BsonElement;
@@ -109,6 +109,7 @@
109109
import static com.mongodb.assertions.Assertions.assertNotNull;
110110
import static com.mongodb.assertions.Assertions.assertTrue;
111111
import static com.mongodb.assertions.Assertions.fail;
112+
import static com.mongodb.internal.connection.BsonWriterHelper.decorateWithDocumentSizeChecking;
112113
import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
113114
import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_NOT_REACHED;
114115
import static com.mongodb.internal.operation.BulkWriteBatch.logWriteModelDoesNotSupportRetries;
@@ -230,7 +231,8 @@ private Integer executeBatch(
230231
retryState.attach(AttachmentKeys.maxWireVersion(), connectionDescription.getMaxWireVersion(), true)
231232
.attach(AttachmentKeys.commandDescriptionSupplier(), () -> BULK_WRITE_COMMAND_NAME, false);
232233
ClientBulkWriteCommand bulkWriteCommand = createBulkWriteCommand(
233-
retryState, effectiveRetryWrites, effectiveWriteConcern, sessionContext, unexecutedModels, batchEncoder,
234+
retryState, effectiveRetryWrites, effectiveWriteConcern, sessionContext, unexecutedModels,
235+
connectionDescription.getMaxDocumentSize(), batchEncoder,
234236
() -> retryState.attach(AttachmentKeys.retryableCommandFlag(), true, true));
235237
return executeBulkWriteCommandAndExhaustOkResponse(
236238
retryState, connectionSource, connection, bulkWriteCommand, effectiveWriteConcern, operationContext);
@@ -337,6 +339,7 @@ private ClientBulkWriteCommand createBulkWriteCommand(
337339
final WriteConcern effectiveWriteConcern,
338340
final SessionContext sessionContext,
339341
final List<? extends ClientNamespacedWriteModel> unexecutedModels,
342+
final int maxStoredDocumentSize,
340343
final BatchEncoder batchEncoder,
341344
final Runnable retriesEnabler) {
342345
BsonDocument commandDocument = new BsonDocument(BULK_WRITE_COMMAND_NAME, new BsonInt32(1))
@@ -353,7 +356,11 @@ private ClientBulkWriteCommand createBulkWriteCommand(
353356
return new ClientBulkWriteCommand(
354357
commandDocument,
355358
new ClientBulkWriteCommand.OpsAndNsInfo(
356-
effectiveRetryWrites, unexecutedModels, batchEncoder, options,
359+
effectiveRetryWrites, unexecutedModels,
360+
// we must validate the size only if no response is expected, otherwise we must rely on the server validation
361+
effectiveWriteConcern.isAcknowledged() ? null : maxStoredDocumentSize,
362+
batchEncoder,
363+
options,
357364
() -> {
358365
retriesEnabler.run();
359366
return retryState.isFirstAttempt()
@@ -669,19 +676,23 @@ OpsAndNsInfo getOpsAndNsInfo() {
669676
public static final class OpsAndNsInfo extends DualMessageSequences {
670677
private final boolean effectiveRetryWrites;
671678
private final List<? extends ClientNamespacedWriteModel> models;
679+
@Nullable
680+
private final Integer maxStoredDocumentSize;
672681
private final BatchEncoder batchEncoder;
673682
private final ConcreteClientBulkWriteOptions options;
674683
private final Supplier<Long> doIfCommandIsRetryableAndAdvanceGetTxnNumber;
675684

676685
OpsAndNsInfo(
677686
final boolean effectiveRetryWrites,
678687
final List<? extends ClientNamespacedWriteModel> models,
688+
@Nullable final Integer maxStoredDocumentSize,
679689
final BatchEncoder batchEncoder,
680690
final ConcreteClientBulkWriteOptions options,
681691
final Supplier<Long> doIfCommandIsRetryableAndAdvanceGetTxnNumber) {
682692
super("ops", new OpsFieldNameValidator(models), "nsInfo", NoOpFieldNameValidator.INSTANCE);
683693
this.effectiveRetryWrites = effectiveRetryWrites;
684694
this.models = models;
695+
this.maxStoredDocumentSize = maxStoredDocumentSize;
685696
this.batchEncoder = batchEncoder;
686697
this.options = options;
687698
this.doIfCommandIsRetryableAndAdvanceGetTxnNumber = doIfCommandIsRetryableAndAdvanceGetTxnNumber;
@@ -704,10 +715,10 @@ public EncodeDocumentsResult encodeDocuments(final WritersProviderAndLimitsCheck
704715
int namespaceIndexInBatch = indexedNamespaces.computeIfAbsent(namespace, k -> indexedNamespacesSizeBeforeCompute);
705716
boolean writeNewNamespace = indexedNamespaces.size() != indexedNamespacesSizeBeforeCompute;
706717
int finalModelIndexInBatch = modelIndexInBatch;
707-
writeResult = writersProviderAndLimitsChecker.tryWrite((opsWriters, nsInfoWriters) -> {
708-
batchEncoder.encodeWriteModel(opsWriters, namespacedModel.getModel(), finalModelIndexInBatch, namespaceIndexInBatch);
718+
writeResult = writersProviderAndLimitsChecker.tryWrite((opsWriter, nsInfoWriter) -> {
719+
batchEncoder.encodeWriteModel(opsWriter, maxStoredDocumentSize,
720+
namespacedModel.getModel(), finalModelIndexInBatch, namespaceIndexInBatch);
709721
if (writeNewNamespace) {
710-
BsonWriter nsInfoWriter = nsInfoWriters.getWriter();
711722
nsInfoWriter.writeStartDocument();
712723
nsInfoWriter.writeString("ns", namespace.getFullName());
713724
nsInfoWriter.writeEndDocument();
@@ -940,7 +951,7 @@ private final class BatchEncoder {
940951
/**
941952
* Must be called at most once.
942953
* Must not be called before calling
943-
* {@link #encodeWriteModel(OrdinaryAndStoredBsonWriters, ClientWriteModel, int, int)} at least once.
954+
* {@link #encodeWriteModel(BsonBinaryWriter, Integer, ClientWriteModel, int, int)} at least once.
944955
* Renders {@code this} unusable.
945956
*/
946957
EncodedBatchInfo intoEncodedBatchInfo() {
@@ -961,16 +972,16 @@ void reset(final int modelIndexInBatch) {
961972
}
962973

963974
void encodeWriteModel(
964-
final OrdinaryAndStoredBsonWriters writers,
975+
final BsonBinaryWriter writer,
976+
@Nullable final Integer maxStoredDocumentSize,
965977
final ClientWriteModel model,
966978
final int modelIndexInBatch,
967979
final int namespaceIndexInBatch) {
968980
assertNotNull(encodedBatchInfo).modelsCount++;
969-
BsonWriter writer = writers.getWriter();
970981
writer.writeStartDocument();
971982
if (model instanceof ConcreteClientInsertOneModel) {
972983
writer.writeInt32("insert", namespaceIndexInBatch);
973-
encodeWriteModelInternals(writers, (ConcreteClientInsertOneModel) model, modelIndexInBatch);
984+
encodeWriteModelInternals(writer, maxStoredDocumentSize, (ConcreteClientInsertOneModel) model, modelIndexInBatch);
974985
} else if (model instanceof ConcreteClientUpdateOneModel) {
975986
writer.writeInt32("update", namespaceIndexInBatch);
976987
writer.writeBoolean("multi", false);
@@ -981,7 +992,7 @@ void encodeWriteModel(
981992
encodeWriteModelInternals(writer, (ConcreteClientUpdateManyModel) model);
982993
} else if (model instanceof ConcreteClientReplaceOneModel) {
983994
writer.writeInt32("update", namespaceIndexInBatch);
984-
encodeWriteModelInternals(writers, (ConcreteClientReplaceOneModel) model);
995+
encodeWriteModelInternals(writer, maxStoredDocumentSize, (ConcreteClientReplaceOneModel) model);
985996
} else if (model instanceof ConcreteClientDeleteOneModel) {
986997
writer.writeInt32("delete", namespaceIndexInBatch);
987998
writer.writeBoolean("multi", false);
@@ -996,12 +1007,16 @@ void encodeWriteModel(
9961007
writer.writeEndDocument();
9971008
}
9981009

999-
private void encodeWriteModelInternals(final OrdinaryAndStoredBsonWriters writers, final ConcreteClientInsertOneModel model, final int modelIndexInBatch) {
1000-
writers.getWriter().writeName("document");
1010+
private void encodeWriteModelInternals(
1011+
final BsonBinaryWriter writer,
1012+
@Nullable final Integer maxStoredDocumentSize,
1013+
final ConcreteClientInsertOneModel model,
1014+
final int modelIndexInBatch) {
1015+
writer.writeName("document");
10011016
Object document = model.getDocument();
10021017
assertNotNull(encodedBatchInfo).insertModelDocumentIds.compute(modelIndexInBatch, (k, knownModelDocumentId) -> {
10031018
IdHoldingBsonWriter documentIdHoldingBsonWriter = new IdHoldingBsonWriter(
1004-
writers.getStoredDocumentWriter(),
1019+
decorateWithDocumentSizeChecking(writer, maxStoredDocumentSize),
10051020
// Reuse `knownModelDocumentId` if it may have been generated by `IdHoldingBsonWriter` in a previous attempt.
10061021
// If its type is not `BsonObjectId`, we know it could not have been generated.
10071022
knownModelDocumentId instanceof BsonObjectId ? knownModelDocumentId.asObjectId() : null);
@@ -1040,13 +1055,16 @@ private void encodeWriteModelInternals(final BsonWriter writer, final AbstractCl
10401055
options.isUpsert().ifPresent(value -> writer.writeBoolean("upsert", value));
10411056
}
10421057

1043-
private void encodeWriteModelInternals(final OrdinaryAndStoredBsonWriters writers, final ConcreteClientReplaceOneModel model) {
1044-
BsonWriter writer = writers.getWriter();
1058+
private void encodeWriteModelInternals(
1059+
final BsonBinaryWriter writer,
1060+
@Nullable final Integer maxStoredDocumentSize,
1061+
final ConcreteClientReplaceOneModel model) {
10451062
writer.writeBoolean("multi", false);
10461063
writer.writeName("filter");
10471064
encodeUsingRegistry(writer, model.getFilter());
10481065
writer.writeName("updateMods");
1049-
encodeUsingRegistry(writers.getStoredDocumentWriter(), model.getReplacement(), COLLECTIBLE_DOCUMENT_ENCODER_CONTEXT);
1066+
encodeUsingRegistry(decorateWithDocumentSizeChecking(writer, maxStoredDocumentSize), model.getReplacement(),
1067+
COLLECTIBLE_DOCUMENT_ENCODER_CONTEXT);
10501068
ConcreteClientReplaceOptions options = model.getOptions();
10511069
options.getCollation().ifPresent(value -> {
10521070
writer.writeName("collation");

0 commit comments

Comments
 (0)