JAVA-1119: Added maxWriteBatchSize to ServerDescription

Author: jyemin

src/main/com/mongodb/DBCollectionImpl.java

@@ -132,7 +132,7 @@ BulkWriteResult executeBulkWriteOperation(final boolean ordered, final List<Writ
         DBPort port = db.getConnector().getPrimaryPort();
         try {
             BulkWriteBatchCombiner bulkWriteBatchCombiner = new BulkWriteBatchCombiner(port.getAddress(), writeConcern);
-            for (Run run : getRunGenerator(ordered, writeRequests, writeConcern, encoder)) {
+            for (Run run : getRunGenerator(ordered, writeRequests, writeConcern, encoder, port)) {
                 try {
                     BulkWriteResult result = run.execute(port);
                     if (result.isAcknowledged()) {
@@ -506,11 +506,11 @@ private WriteResult insertWithWriteProtocol(final List<DBObject> list, final Wri
     }
 
     private Iterable<Run> getRunGenerator(final boolean ordered, final List<WriteRequest> writeRequests,
-                                          final WriteConcern writeConcern, final DBEncoder encoder) {
+                                          final WriteConcern writeConcern, final DBEncoder encoder, final DBPort port) {
         if (ordered) {
-            return new OrderedRunGenerator(writeRequests, writeConcern, encoder);
+            return new OrderedRunGenerator(writeRequests, writeConcern, encoder, port);
         } else {
-            return new UnorderedRunGenerator(writeRequests, writeConcern, encoder);
+            return new UnorderedRunGenerator(writeRequests, writeConcern, encoder, port);
         }
     }
 
@@ -529,15 +529,16 @@ private Logger getLogger() {
         return TRACE_LOGGER;
     }
 
-    private static final int MAX_COUNT = 1000;
-
     private class OrderedRunGenerator implements Iterable<Run> {
         private final List<WriteRequest> writeRequests;
+        private final DBPort port;
         private final WriteConcern writeConcern;
         private final DBEncoder encoder;
 
-        public OrderedRunGenerator(final List<WriteRequest> writeRequests, final WriteConcern writeConcern, final DBEncoder encoder) {
+        public OrderedRunGenerator(final List<WriteRequest> writeRequests, final WriteConcern writeConcern, final DBEncoder encoder,
+                                   final DBPort port) {
             this.writeRequests = writeRequests;
+            this.port = port;
             this.writeConcern = writeConcern.continueOnError(false);
             this.encoder = encoder;
         }
@@ -566,7 +567,8 @@ public Run next() {
                 private int getStartIndexOfNextRun() {
                     WriteRequest.Type type = writeRequests.get(curIndex).getType();
                     for (int i = curIndex; i < writeRequests.size(); i++) {
-                        if (i == MAX_COUNT || writeRequests.get(i).getType() != type) {
+                        if (i == db.getConnector().getServerDescription(port.getAddress()).getMaxWriteBatchSize()
+                            || writeRequests.get(i).getType() != type) {
                             return i;
                         }
                     }
@@ -584,12 +586,14 @@ public void remove() {
 
     private class UnorderedRunGenerator implements Iterable<Run> {
         private final List<WriteRequest> writeRequests;
+        private final DBPort port;
         private final WriteConcern writeConcern;
         private final DBEncoder encoder;
 
         public UnorderedRunGenerator(final List<WriteRequest> writeRequests, final WriteConcern writeConcern,
-                                     final DBEncoder encoder) {
+                                     final DBEncoder encoder, final DBPort port) {
             this.writeRequests = writeRequests;
+            this.port = port;
             this.writeConcern = writeConcern.continueOnError(true);
             this.encoder = encoder;
         }
@@ -622,7 +626,7 @@ public Run next() {
                         }
                         run.add(writeRequest, curIndex);
                         curIndex++;
-                        if (run.size() > MAX_COUNT) {
+                        if (run.size() > db.getConnector().getServerDescription(port.getAddress()).getMaxWriteBatchSize()) {
                             return runs.remove(run.type);
                         }
                     }

src/main/com/mongodb/ServerDescription.java

@@ -49,6 +49,8 @@ class ServerDescription {
     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 = 512;
+
     private final ServerAddress address;
 
     private final ServerType type;
@@ -57,6 +59,7 @@ class ServerDescription {
     private final Set<String> arbiters;
     private final String primary;
     private final int maxDocumentSize;
+    private final int maxWriteBatchSize;
 
     private final int maxMessageSize;
     private final Tags tags;
@@ -79,6 +82,7 @@ static class Builder {
         private String primary;
         private int maxDocumentSize = DEFAULT_MAX_DOCUMENT_SIZE;
         private int maxMessageSize = DEFAULT_MAX_MESSAGE_SIZE;
+        private int maxWriteBatchSize = DEFAULT_MAX_WRITE_BATCH_SIZE;
         private Tags tags = Tags.freeze(new Tags());
         private String setName;
         private Integer setVersion;
@@ -130,6 +134,11 @@ public Builder maxMessageSize(final int maxMessageSize) {
             return this;
         }
 
+        public Builder maxWriteBatchSize(final int maxWriteBatchSize) {
+            this.maxWriteBatchSize = maxWriteBatchSize;
+            return this;
+        }
+
         public Builder tags(final Tags tags) {
             this.tags = tags == null ? Tags.freeze(new Tags()) : Tags.freeze(tags);
             return this;
@@ -220,6 +229,10 @@ public static int getDefaultMaxWireVersion() {
         return 0;
     }
 
+    public static int getDefaultMaxWriteBatchSize() {
+        return DEFAULT_MAX_WRITE_BATCH_SIZE;
+    }
+
     public static Builder builder() {
         return new Builder();
     }
@@ -272,6 +285,10 @@ public int getMaxMessageSize() {
         return maxMessageSize;
     }
 
+    public int getMaxWriteBatchSize() {
+        return maxWriteBatchSize;
+    }
+
     public Tags getTags() {
         return tags;
     }
@@ -285,8 +302,8 @@ public int getMaxWireVersion() {
     }
 
     /**
-     * Returns true if the server has the given tags.  A server of either type @code{ServerType.StandAlone} or
-     * @code{ServerType.ShardRouter} is considered to have all tags, so this method will always return true for instances of either of
+     * Returns true if the server has the given tags.  A server of either type {@code ServerType.StandAlone} or
+     * {@code ServerType.ShardRouter} is considered to have all tags, so this method will always return true for instances of either of
      * those types.
      *
      * @param desiredTags the tags
@@ -365,6 +382,9 @@ public boolean equals(final Object o) {
         if (maxMessageSize != that.maxMessageSize) {
             return false;
         }
+        if (maxWriteBatchSize != that.maxWriteBatchSize) {
+            return false;
+        }
         if (ok != that.ok) {
             return false;
         }
@@ -422,6 +442,7 @@ public int hashCode() {
         result = 31 * result + (primary != null ? primary.hashCode() : 0);
         result = 31 * result + maxDocumentSize;
         result = 31 * result + maxMessageSize;
+        result = 31 * result + maxWriteBatchSize;
         result = 31 * result + tags.hashCode();
         result = 31 * result + (setName != null ? setName.hashCode() : 0);
         result = 31 * result + (setVersion != null ? setVersion.hashCode() : 0);
@@ -444,6 +465,7 @@ public String toString() {
                + ", primary='" + primary + '\''
                + ", maxDocumentSize=" + maxDocumentSize
                + ", maxMessageSize=" + maxMessageSize
+               + ", maxWriteBatchSize=" + maxWriteBatchSize
                + ", tags=" + tags
                + ", setName='" + setName + '\''
                + ", setVersion='" + setVersion + '\''
@@ -477,6 +499,7 @@ public String getShortDescription() {
         primary = builder.primary;
         maxDocumentSize = builder.maxDocumentSize;
         maxMessageSize = builder.maxMessageSize;
+        maxWriteBatchSize = builder.maxWriteBatchSize;
         tags = builder.tags;
         setName = builder.setName;
         setVersion = builder.setVersion;

src/main/com/mongodb/ServerStateNotifier.java

@@ -140,6 +140,8 @@ private ServerDescription createDescription(final CommandResult commandResult, f
                                 .primary(commandResult.getString("primary"))
                                 .maxDocumentSize(commandResult.getInt("maxBsonObjectSize", ServerDescription.getDefaultMaxDocumentSize()))
                                 .maxMessageSize(commandResult.getInt("maxMessageSizeBytes", ServerDescription.getDefaultMaxMessageSize()))
+                                .maxWriteBatchSize(commandResult.getInt("maxWriteBatchSize",
+                                                                        ServerDescription.getDefaultMaxWriteBatchSize()))
                                 .tags(getTagsFromDocument((DBObject) commandResult.get("tags")))
                                 .setName(commandResult.getString("setName"))
                                 .setVersion((Integer) commandResult.get("setVersion"))

src/test/com/mongodb/Fixture.java

@@ -52,6 +52,16 @@ public static synchronized DB getDefaultDatabase() {
         return defaultDatabase;
     }
 
+    /**
+     *
+     * @param version  must be a major version, e.g. 1.8, 2,0, 2.2
+     * @return true if server is at least specified version
+     */
+    public static boolean serverIsAtLeastVersion(double version) {
+        String serverVersion = (String) getMongoClient().getDB("admin").command("serverStatus").get("version");
+        return Double.parseDouble(serverVersion.substring(0, 3)) >= version;
+    }
+
     static class ShutdownHook extends Thread {
         @Override
         public void run() {

src/test/com/mongodb/ServerDescriptionTest.java

@@ -71,6 +71,7 @@ public void testDefaults() throws UnknownHostException {
 
         assertEquals(0x1000000, serverDescription.getMaxDocumentSize());
         assertEquals(0x2000000, serverDescription.getMaxMessageSize());
+        assertEquals(512, serverDescription.getMaxWriteBatchSize());
 
         assertNull(serverDescription.getPrimary());
         assertEquals(Collections.<String>emptySet(), serverDescription.getHosts());
@@ -94,6 +95,7 @@ public void testObjectOverrides() throws UnknownHostException {
                                                              .setVersion(11)
                                                              .maxDocumentSize(100)
                                                              .maxMessageSize(200)
+                                                             .maxWriteBatchSize(1024)
                                                              .averagePingTime(50000, java.util.concurrent.TimeUnit.NANOSECONDS)
                                                              .primary("localhost:27017")
                                                              .hosts(new HashSet<String>(Arrays.asList("localhost:27017",

src/test/com/mongodb/ServerStateNotifierSpecification.groovy

@@ -0,0 +1,62 @@
+package com.mongodb
+
+import static com.mongodb.Fixture.getMongoClient
+import static com.mongodb.Fixture.serverIsAtLeastVersion
+import static org.junit.Assume.assumeFalse
+import static org.junit.Assume.assumeTrue
+
+class ServerStateNotifierSpecification extends FunctionalSpecification {
+    ServerDescription newDescription
+    ServerStateNotifier serverStateNotifier
+
+    def setup() {
+        serverStateNotifier = new ServerStateNotifier(new ServerAddress(),
+                                                      new ChangeListener<ServerDescription>() {
+                                                          @Override
+                                                          void stateChanged(final ChangeEvent<ServerDescription> event) {
+                                                              newDescription = event.newValue
+                                                          }
+                                                      },
+                                                      SocketSettings.builder().build(), getMongoClient())
+    }
+
+    def cleanup() {
+        serverStateNotifier.close();
+    }
+
+    def 'should set server version'() {
+        given:
+        CommandResult commandResult = database.command(new BasicDBObject('buildinfo', 1))
+        def expectedVersion = new ServerVersion((commandResult.get('versionArray') as List<Integer>).subList(0, 3))
+
+        when:
+        serverStateNotifier.run()
+
+        then:
+        newDescription.version == expectedVersion
+    }
+
+    def 'should set max wire batch size when provided by server'() {
+        assumeTrue(serverIsAtLeastVersion(2.5))
+
+        given:
+        CommandResult commandResult = database.command(new BasicDBObject('ismaster', 1))
+        def expected = commandResult.get('maxWriteBatchSize')
+
+        when:
+        serverStateNotifier.run()
+
+        then:
+        newDescription.maxWriteBatchSize == expected
+    }
+
+    def 'should set default max wire batch size when not provided by server'() {
+        assumeFalse(serverIsAtLeastVersion(2.5))
+
+        when:
+        serverStateNotifier.run()
+
+        then:
+        newDescription.maxWriteBatchSize == ServerDescription.getDefaultMaxWriteBatchSize()
+    }
+}