JAVA-1156: Properly ensure that none of the server-defined limits on the write commands are exceeded.

There is one limit for the number of items that are allowed in each command (maxWriteBatchSize from ismaster).
There is a second limit for the number of bytes in the encoded message (maxBsonObjectSize from ismaster), with an exception for a write command containing just a single item, which is allowed to exceed that limit.

Author: jyemin

src/main/com/mongodb/BaseWriteCommandMessage.java

@@ -87,9 +87,19 @@ private void writeCommandHeader(final OutputBuffer buffer) {
     protected abstract BaseWriteCommandMessage writeTheWrites(final OutputBuffer buffer, final int commandStartPosition,
                                                               final BSONBinaryWriter writer);
 
-    protected boolean maximumCommandDocumentSizeExceeded(final OutputBuffer buffer, final int commandStartPosition) {
-        // Subtract 2 to account for the trailing 0x0 at the end of the enclosing array and command document
-        return buffer.getPosition() - commandStartPosition > getSettings().getMaxDocumentSize() + HEADROOM - 2;
+    protected boolean exceedsLimits(final int batchLength, final int batchItemCount) {
+        return (exceedsBatchLengthLimit(batchLength, batchItemCount) || exceedsBatchItemCountLimit(batchItemCount));
+    }
+
+    // make a special exception for a command with only a single item added to it.  It's allowed to exceed maximum document size so that
+    // it's possible to, say, send a replacement document that is itself 16MB, which would push the size of the containing command
+    // document to be greater than the maximum document size.
+    private boolean exceedsBatchLengthLimit(final int batchLength, final int batchItemCount) {
+        return batchLength > getSettings().getMaxDocumentSize() && batchItemCount > 1;
+    }
+
+    private boolean exceedsBatchItemCountLimit(final int batchItemCount) {
+        return batchItemCount > getSettings().getMaxWriteBatchSize();
     }
 
     public abstract int getItemCount();

src/main/com/mongodb/DBCollectionImpl.java

@@ -372,7 +372,7 @@ private BulkWriteResult insertWithCommandProtocol(final List<DBObject> list, fin
 
         BaseWriteCommandMessage message = new InsertCommandMessage(getNamespace(), writeConcern, list,
                                                                    DefaultDBEncoder.FACTORY.create(), encoder,
-                                                                   getMessageSettings(port.getAddress()));
+                                                                   getMessageSettings(port));
         return writeWithCommandProtocol(port, INSERT, message, writeConcern);
     }
 
@@ -392,7 +392,7 @@ private BulkWriteResult removeWithCommandProtocol(final List<RemoveRequest> remo
                                                       final DBEncoder encoder, final DBPort port) {
         BaseWriteCommandMessage message = new DeleteCommandMessage(getNamespace(), writeConcern, removeList,
                                                                    DefaultDBEncoder.FACTORY.create(), encoder,
-                                                                   getMessageSettings(port.getAddress()));
+                                                                   getMessageSettings(port));
         return writeWithCommandProtocol(port, REMOVE, message, writeConcern);
     }
 
@@ -402,7 +402,7 @@ private BulkWriteResult updateWithCommandProtocol(final List<ModifyRequest> upda
                                                       final DBEncoder encoder, final DBPort port) {
         BaseWriteCommandMessage message = new UpdateCommandMessage(getNamespace(), writeConcern, updates,
                                                                    DefaultDBEncoder.FACTORY.create(), encoder,
-                                                                   getMessageSettings(port.getAddress()));
+                                                                   getMessageSettings(port));
         return writeWithCommandProtocol(port, UPDATE, message, writeConcern);
     }
 
@@ -438,13 +438,19 @@ private boolean useWriteCommands(final WriteConcern concern, final DBPort port)
                db.getConnector().getServerDescription(port.getAddress()).getVersion().compareTo(new ServerVersion(2, 6)) >= 0;
     }
 
-    private MessageSettings getMessageSettings(final ServerAddress address) {
-        ServerDescription serverDescription = db.getConnector().getServerDescription(address);
-        return MessageSettings.builder().maxDocumentSize(serverDescription.getMaxDocumentSize()).maxMessageSize(serverDescription
-                                                                                                                .getMaxMessageSize())
+    private MessageSettings getMessageSettings(final DBPort port) {
+        ServerDescription serverDescription = db.getConnector().getServerDescription(port.getAddress());
+        return MessageSettings.builder()
+                              .maxDocumentSize(serverDescription.getMaxDocumentSize())
+                              .maxMessageSize(serverDescription.getMaxMessageSize())
+                              .maxWriteBatchSize(serverDescription.getMaxWriteBatchSize())
                               .build();
     }
 
+    private int getMaxWriteBatchSize(final DBPort port) {
+        return db.getConnector().getServerDescription(port.getAddress()).getMaxWriteBatchSize();
+    }
+
     private MongoNamespace getNamespace() {
         return new MongoNamespace(getDB().getName(), getName());
     }
@@ -551,7 +557,7 @@ public OrderedRunGenerator(final List<WriteRequest> writeRequests, final WriteCo
             this.writeRequests = writeRequests;
             this.writeConcern = writeConcern.continueOnError(false);
             this.encoder = encoder;
-            this.maxBatchWriteSize = db.getConnector().getServerDescription(port.getAddress()).getMaxWriteBatchSize();
+            this.maxBatchWriteSize = getMaxWriteBatchSize(port);
         }
 
         @Override
@@ -605,7 +611,7 @@ public UnorderedRunGenerator(final List<WriteRequest> writeRequests, final Write
             this.writeRequests = writeRequests;
             this.writeConcern = writeConcern.continueOnError(true);
             this.encoder = encoder;
-            this.maxBatchWriteSize = db.getConnector().getServerDescription(port.getAddress()).getMaxWriteBatchSize();
+            this.maxBatchWriteSize = getMaxWriteBatchSize(port);
         }
 
         @Override

src/main/com/mongodb/DeleteCommandMessage.java

@@ -51,7 +51,7 @@ protected BaseWriteCommandMessage writeTheWrites(final OutputBuffer buffer, fina
             writer.writeInt32("limit", remove.isMulti() ? 0 : 1);
             writer.popMaxDocumentSize();
             writer.writeEndDocument();
-            if (maximumCommandDocumentSizeExceeded(buffer, commandStartPosition)) {
+            if (exceedsLimits(buffer.getPosition() - commandStartPosition, i + 1)) {
                 writer.reset();
                 nextMessage = new DeleteCommandMessage(getWriteNamespace(), getWriteConcern(), deletes.subList(i, deletes.size()),
                                                        getCommandEncoder(), queryEncoder, getSettings());

src/main/com/mongodb/InsertCommandMessage.java

@@ -43,10 +43,9 @@ protected InsertCommandMessage writeTheWrites(final OutputBuffer buffer, final i
         for (int i = 0; i < documents.size(); i++) {
             writer.mark();
             writer.encodeDocument(encoder, documents.get(i));
-            if (maximumCommandDocumentSizeExceeded(buffer, commandStartPosition)) {
+            if (exceedsLimits(buffer.getPosition() - commandStartPosition, i + 1)) {
                 writer.reset();
-                nextMessage = new InsertCommandMessage(getWriteNamespace(), getWriteConcern(),
-                                                       documents.subList(i, documents.size()),
+                nextMessage = new InsertCommandMessage(getWriteNamespace(), getWriteConcern(), documents.subList(i, documents.size()),
                                                        getCommandEncoder(), encoder, getSettings());
                 break;
             }

src/main/com/mongodb/MessageSettings.java

@@ -22,9 +22,11 @@
 final class MessageSettings {
     private static final int DEFAULT_MAX_DOCUMENT_SIZE = 0x1000000;  // 16MB
     private static final int DEFAULT_MAX_MESSAGE_SIZE = 0x2000000;   // 32MB
+    private static final int DEFAULT_MAX_WRITE_BATCH_SIZE = 1000;
 
     private final int maxDocumentSize;
     private final int maxMessageSize;
+    private final int maxWriteBatchSize;
 
     public static Builder builder() {
         return new Builder();
@@ -33,6 +35,7 @@ public static Builder builder() {
     static final class Builder {
         private int maxDocumentSize = DEFAULT_MAX_DOCUMENT_SIZE;
         private int maxMessageSize = DEFAULT_MAX_MESSAGE_SIZE;
+        private int maxWriteBatchSize = DEFAULT_MAX_WRITE_BATCH_SIZE;
 
         public MessageSettings build() {
             return new MessageSettings(this);
@@ -48,6 +51,12 @@ public Builder maxMessageSize(final int maxMessageSize) {
             this.maxMessageSize = maxMessageSize;
             return this;
         }
+
+        public Builder maxWriteBatchSize(final int maxWriteBatchSize) {
+            this.maxWriteBatchSize = maxWriteBatchSize;
+            return this;
+        }
+
         // CHECKSTYLE:ON
     }
 
@@ -59,8 +68,13 @@ public int getMaxMessageSize() {
         return maxMessageSize;
     }
 
+    public int getMaxWriteBatchSize() {
+        return maxWriteBatchSize;
+    }
+
     MessageSettings(final Builder builder) {
         this.maxDocumentSize = builder.maxDocumentSize;
         this.maxMessageSize = builder.maxMessageSize;
+        this.maxWriteBatchSize = builder.maxWriteBatchSize;
     }
 }

src/main/com/mongodb/UpdateCommandMessage.java

@@ -54,7 +54,7 @@ protected UpdateCommandMessage writeTheWrites(final OutputBuffer buffer, final i
             }
             writer.popMaxDocumentSize();
             writer.writeEndDocument();
-            if (maximumCommandDocumentSizeExceeded(buffer, commandStartPosition)) {
+            if (exceedsLimits(buffer.getPosition() - commandStartPosition, i + 1)) {
                 writer.reset();
                 nextMessage = new UpdateCommandMessage(getWriteNamespace(), getWriteConcern(), updates.subList(i, updates.size()),
                                                        getCommandEncoder(), encoder, getSettings());

src/test/com/mongodb/DBCollectionTest.java

@@ -21,6 +21,7 @@
 import org.junit.Test;
 
 import java.net.UnknownHostException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
@@ -49,6 +50,17 @@ public void testMultiInsert() {
         c.insert(inserted1,inserted2);
     }
 
+    @Test
+    public void testLargeMultiInsert() {
+        List<DBObject> documents = new ArrayList<DBObject>();
+        for (int i = 0; i < 1001; i++) {
+            documents.add(new BasicDBObject());
+        }
+
+        collection.insert(documents);
+        assertEquals(1001, collection.count());
+    }
+
     @Test
     public void testCappedCollection() {
         String collectionName = "testCapped";