Implement splitting and encoding ops, nsInfo as separate OP_MSG sections, implement prose tests

bson/src/main/org/bson/AbstractBsonWriter.java

@@ -748,6 +748,15 @@ protected void throwInvalidState(final String methodName, final State... validSt
                                                        methodName, validStatesString, state));
     }
 
+    /**
+     * {@inheritDoc}
+     * <p>
+     * The {@link #flush()} method of {@link AbstractBsonWriter} does nothing.</p>
+     */
+    @Override
+    public void flush() {
+    }
+
     @Override
     public void close() {
         closed = true;

bson/src/main/org/bson/BSONCallbackAdapter.java

@@ -39,11 +39,6 @@ protected BSONCallbackAdapter(final BsonWriterSettings settings, final BSONCallb
         this.bsonCallback = bsonCallback;
     }
 
-    @Override
-    public void flush() {
-        //Looks like should be no-op?
-    }
-
     @Override
     public void doWriteStartDocument() {
         BsonContextType contextType = getState() == State.SCOPE_DOCUMENT

bson/src/main/org/bson/BsonBinaryWriter.java

@@ -108,10 +108,6 @@ public BsonBinaryWriterSettings getBinaryWriterSettings() {
         return binaryWriterSettings;
     }
 
-    @Override
-    public void flush() {
-    }
-
     @Override
     protected Context getContext() {
         return (Context) super.getContext();

bson/src/main/org/bson/BsonDocumentWriter.java

@@ -194,10 +194,6 @@ public void doWriteUndefined() {
         write(new BsonUndefined());
     }
 
-    @Override
-    public void flush() {
-    }
-
     @Override
     protected Context getContext() {
         return (Context) super.getContext();

bson/src/main/org/bson/io/OutputBuffer.java

@@ -41,6 +41,16 @@ public void write(final byte[] b) {
     public void close() {
     }
 
+    /**
+     * {@inheritDoc}
+     * <p>
+     * The {@link #flush()} method of {@link OutputBuffer} does nothing.</p>
+     */
+    @Override
+    public void flush() throws IOException {
+        super.flush();
+    }
+
     @Override
     public void write(final byte[] bytes, final int offset, final int length) {
         writeBytes(bytes, offset, length);

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

@@ -50,8 +50,7 @@ <T> void commandAsync(String database, BsonDocument command, FieldNameValidator
 
     <T> void commandAsync(String database, BsonDocument command, FieldNameValidator commandFieldNameValidator,
             @Nullable ReadPreference readPreference, Decoder<T> commandResultDecoder,
-            OperationContext operationContext, boolean responseExpected, @Nullable SplittablePayload payload,
-            @Nullable FieldNameValidator payloadFieldNameValidator, SingleResultCallback<T> callback);
+            OperationContext operationContext, boolean responseExpected, MessageSequences sequences, SingleResultCallback<T> callback);
 
     void markAsPinned(Connection.PinningMode pinningMode);
 }

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

@@ -16,32 +16,63 @@
 
 package com.mongodb.internal.connection;
 
+import com.mongodb.internal.connection.DualMessageSequences.EncodeDocumentsResult;
+import com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker;
+import com.mongodb.internal.validator.NoOpFieldNameValidator;
+import com.mongodb.lang.Nullable;
+import org.bson.BsonBinaryWriter;
+import org.bson.BsonBinaryWriterSettings;
 import org.bson.BsonDocument;
 import org.bson.BsonElement;
 import org.bson.BsonMaximumSizeExceededException;
 import org.bson.BsonValue;
 import org.bson.BsonWriter;
+import org.bson.BsonWriterSettings;
+import org.bson.FieldNameValidator;
 import org.bson.codecs.BsonValueCodecProvider;
-import org.bson.codecs.Codec;
+import org.bson.codecs.Encoder;
 import org.bson.codecs.EncoderContext;
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.io.BsonOutput;
 
 import java.util.List;
 
+import static com.mongodb.assertions.Assertions.assertTrue;
+import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.FAIL_LIMIT_EXCEEDED;
+import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_NOT_REACHED;
+import static com.mongodb.internal.connection.DualMessageSequences.WritersProviderAndLimitsChecker.WriteResult.OK_LIMIT_REACHED;
+import static com.mongodb.internal.connection.MessageSettings.DOCUMENT_HEADROOM_SIZE;
 import static java.lang.String.format;
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
 
-final class BsonWriterHelper {
-    private static final int DOCUMENT_HEADROOM = 1024 * 16;
+/**
+ * This class is not part of the public API and may be removed or changed at any time.
+ */
+public final class BsonWriterHelper {
     private static final CodecRegistry REGISTRY = fromProviders(new BsonValueCodecProvider());
     private static final EncoderContext ENCODER_CONTEXT = EncoderContext.builder().build();
 
-    static void writeElements(final BsonWriter writer, final List<BsonElement> bsonElements) {
-        for (BsonElement bsonElement : bsonElements) {
-            writer.writeName(bsonElement.getName());
-            getCodec(bsonElement.getValue()).encode(writer, bsonElement.getValue(), ENCODER_CONTEXT);
+    static void appendElementsToDocument(
+            final BsonOutput bsonOutputWithDocument,
+            final int documentStartPosition,
+            @Nullable final List<BsonElement> bsonElements) {
+        int bsonDocumentEndingSize = 1;
+        int appendFrom = bsonOutputWithDocument.getPosition() - bsonDocumentEndingSize;
+        BsonBinaryWriter writer = createBsonBinaryWriter(bsonOutputWithDocument, NoOpFieldNameValidator.INSTANCE, null);
+        // change `writer`s state so that we can append elements
+        writer.writeStartDocument();
+        bsonOutputWithDocument.truncateToPosition(appendFrom);
+        if (bsonElements != null) {
+            for (BsonElement element : bsonElements) {
+                String name = element.getName();
+                BsonValue value = element.getValue();
+                writer.writeName(name);
+                encodeUsingRegistry(writer, value);
+            }
         }
+        // write the BSON document ending
+        bsonOutputWithDocument.writeByte(0);
+        backpatchLength(documentStartPosition, bsonOutputWithDocument);
     }
 
     static void writePayloadArray(final BsonWriter writer, final BsonOutput bsonOutput, final MessageSettings settings,
@@ -65,16 +96,89 @@ static void writePayload(final BsonWriter writer, final BsonOutput bsonOutput, f
         }
 
         if (payload.getPosition() == 0) {
-            throw new BsonMaximumSizeExceededException(format("Payload document size is larger than maximum of %d.",
-                    payloadSettings.getMaxDocumentSize()));
+            throw createBsonMaximumSizeExceededException(payloadSettings.getMaxDocumentSize());
         }
     }
 
+    /**
+     * @return See {@link DualMessageSequences#encodeDocuments(WritersProviderAndLimitsChecker)}.
+     */
+    static EncodeDocumentsResult writeDocumentsOfDualMessageSequences(
+            final DualMessageSequences dualMessageSequences,
+            final int commandDocumentSizeInBytes,
+            final BsonOutput firstOutput,
+            final BsonOutput secondOutput,
+            final MessageSettings messageSettings,
+            final boolean validateDocumentSizeLimits) {
+        BsonBinaryWriter firstWriter = createBsonBinaryWriter(
+                firstOutput, dualMessageSequences.getFirstFieldNameValidator(), validateDocumentSizeLimits ? messageSettings : null);
+        BsonBinaryWriter secondWriter = createBsonBinaryWriter(
+                secondOutput, dualMessageSequences.getSecondFieldNameValidator(), validateDocumentSizeLimits ? messageSettings : null);
+        // the size of operation-agnostic command fields (a.k.a. extra elements) is counted towards `messageOverheadInBytes`
+        int messageOverheadInBytes = 1000;
+        int maxSizeInBytes = messageSettings.getMaxMessageSize() - (messageOverheadInBytes + commandDocumentSizeInBytes);
+        int firstStart = firstOutput.getPosition();
+        int secondStart = secondOutput.getPosition();
+        int maxBatchCount = messageSettings.getMaxBatchCount();
+        return dualMessageSequences.encodeDocuments(write -> {
+            int firstBeforeWritePosition = firstOutput.getPosition();
+            int secondBeforeWritePosition = secondOutput.getPosition();
+            int batchCountAfterWrite = write.doAndGetBatchCount(firstWriter, secondWriter);
+            assertTrue(batchCountAfterWrite <= maxBatchCount);
+            int writtenSizeInBytes =
+                    firstOutput.getPosition() - firstStart
+                    + secondOutput.getPosition() - secondStart;
+            if (writtenSizeInBytes < maxSizeInBytes && batchCountAfterWrite < maxBatchCount) {
+                return OK_LIMIT_NOT_REACHED;
+            } else if (writtenSizeInBytes > maxSizeInBytes) {
+                firstOutput.truncateToPosition(firstBeforeWritePosition);
+                secondOutput.truncateToPosition(secondBeforeWritePosition);
+                if (batchCountAfterWrite == 1) {
+                    // we have failed to write a single document
+                    throw createBsonMaximumSizeExceededException(messageSettings.getMaxDocumentSize());
+                }
+                return FAIL_LIMIT_EXCEEDED;
+            } else {
+                return OK_LIMIT_REACHED;
+            }
+        });
+    }
+
+    /**
+     * @param messageSettings Non-{@code null} iff the document size limit must be validated.
+     */
+    static BsonBinaryWriter createBsonBinaryWriter(
+            final BsonOutput out,
+            final FieldNameValidator validator,
+            @Nullable final MessageSettings messageSettings) {
+        return new BsonBinaryWriter(
+                new BsonWriterSettings(),
+                messageSettings == null
+                        ? new BsonBinaryWriterSettings()
+                        : new BsonBinaryWriterSettings(messageSettings.getMaxDocumentSize() + DOCUMENT_HEADROOM_SIZE),
+                out,
+                validator);
+    }
+
+    /**
+     * Backpatches the document/message/sequence length into the beginning of the document/message/sequence.
+     *
+     * @param startPosition The start position of the document/message/sequence in {@code bsonOutput}.
+     */
+    static void backpatchLength(final int startPosition, final BsonOutput bsonOutput) {
+        int messageLength = bsonOutput.getPosition() - startPosition;
+        bsonOutput.writeInt32(startPosition, messageLength);
+    }
+
+    private static BsonMaximumSizeExceededException createBsonMaximumSizeExceededException(final int maxSize) {
+        return new BsonMaximumSizeExceededException(format("Payload document size is larger than maximum of %d.", maxSize));
+    }
+
     private static boolean writeDocument(final BsonWriter writer, final BsonOutput bsonOutput, final MessageSettings settings,
                                          final BsonDocument document, final int messageStartPosition, final int batchItemCount,
                                          final int maxSplittableDocumentSize) {
         int currentPosition = bsonOutput.getPosition();
-        getCodec(document).encode(writer, document, ENCODER_CONTEXT);
+        encodeUsingRegistry(writer, document);
         int messageSize = bsonOutput.getPosition() - messageStartPosition;
         int documentSize = bsonOutput.getPosition() - currentPosition;
         if (exceedsLimits(settings, messageSize, documentSize, batchItemCount)
@@ -85,24 +189,25 @@ private static boolean writeDocument(final BsonWriter writer, final BsonOutput b
         return true;
     }
 
-    @SuppressWarnings({"unchecked"})
-    private static Codec<BsonValue> getCodec(final BsonValue bsonValue) {
-        return (Codec<BsonValue>) REGISTRY.get(bsonValue.getClass());
+    static void encodeUsingRegistry(final BsonWriter writer, final BsonValue value) {
+        @SuppressWarnings("unchecked")
+        Encoder<BsonValue> encoder = (Encoder<BsonValue>) REGISTRY.get(value.getClass());
+        encoder.encode(writer, value, ENCODER_CONTEXT);
     }
 
     private static MessageSettings getPayloadMessageSettings(final SplittablePayload.Type type, final MessageSettings settings) {
         MessageSettings payloadMessageSettings = settings;
         if (type != SplittablePayload.Type.INSERT) {
             payloadMessageSettings = createMessageSettingsBuilder(settings)
-                    .maxDocumentSize(settings.getMaxDocumentSize() + DOCUMENT_HEADROOM)
+                    .maxDocumentSize(settings.getMaxDocumentSize() + DOCUMENT_HEADROOM_SIZE)
                     .build();
         }
         return payloadMessageSettings;
     }
 
     private static MessageSettings getDocumentMessageSettings(final MessageSettings settings) {
         return createMessageSettingsBuilder(settings)
-                    .maxMessageSize(settings.getMaxDocumentSize() + DOCUMENT_HEADROOM)
+                    .maxMessageSize(settings.getMaxDocumentSize() + DOCUMENT_HEADROOM_SIZE)
                     .build();
     }
 
@@ -126,8 +231,45 @@ private static boolean exceedsLimits(final MessageSettings settings, final int m
         return false;
     }
 
+    /**
+     * @return {@code writer} if {@code maxDocumentSize} is {@code null}, otherwise decorates it.
+     */
+    public static BsonWriter decorateWithDocumentSizeChecking(final BsonBinaryWriter writer, @Nullable final Integer maxDocumentSize) {
+        return maxDocumentSize == null ? writer : new DocumentSizeLimitCheckingBsonBinaryWriter(writer, maxDocumentSize);
+    }
 
     private BsonWriterHelper() {
     }
 
+    private static final class DocumentSizeLimitCheckingBsonBinaryWriter extends LevelCountingBsonWriter {
+        private final int maxStoredDocumentSize;
+        private final BsonOutput out;
+        private int documentStart;
+
+        DocumentSizeLimitCheckingBsonBinaryWriter(final BsonBinaryWriter writer, final int maxStoredDocumentSize) {
+            super(writer);
+            assertTrue(maxStoredDocumentSize > 0);
+            this.maxStoredDocumentSize = maxStoredDocumentSize;
+            this.out = writer.getBsonOutput();
+        }
+
+        @Override
+        public void writeStartDocument() {
+            if (getCurrentLevel() == INITIAL_LEVEL) {
+                documentStart = out.getPosition();
+            }
+            super.writeStartDocument();
+        }
+
+        @Override
+        public void writeEndDocument() throws BsonMaximumSizeExceededException {
+            super.writeEndDocument();
+            if (getCurrentLevel() == INITIAL_LEVEL) {
+                int documentSize = out.getPosition() - documentStart;
+                if (documentSize > maxStoredDocumentSize) {
+                    throw createBsonMaximumSizeExceededException(maxStoredDocumentSize);
+                }
+            }
+        }
+    }
 }