JAVA-1948: Use nToReturn calculation to determine batchSize for the getMore command, and in addition take the absolute value and then check for zero, as getMore requires that any specified batchSize be a positive integer.

jyemin · jyemin · commit 0d36abb39cd4 · 2015-10-21T22:59:18.000-04:00
diff --git a/driver-core/src/main/com/mongodb/operation/AsyncQueryBatchCursor.java b/driver-core/src/main/com/mongodb/operation/AsyncQueryBatchCursor.java
@@ -130,7 +130,7 @@ ServerCursor getServerCursor() {
     }
 
     private boolean limitReached() {
-        return limit != 0 && count >= limit;
+        return Math.abs(limit) != 0 && count >= Math.abs(limit);
     }
 
     private void getMore(final SingleResultCallback<List<T>> callback) {
@@ -162,8 +162,9 @@ private BsonDocument asGetMoreCommandDocument() {
         BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursor.getId()))
                                 .append("collection", new BsonString(namespace.getCollectionName()));
 
-        if (batchSize != 0) {
-            document.append("batchSize", new BsonInt32(Math.abs(batchSize)));
+        int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count));
+        if (batchSizeForGetMoreCommand != 0) {
+            document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));
         }
 
         return document;
diff --git a/driver-core/src/main/com/mongodb/operation/FindOperation.java b/driver-core/src/main/com/mongodb/operation/FindOperation.java
@@ -728,7 +728,7 @@ private Function<BsonDocument, BatchCursor<T>> transformer(final ConnectionSourc
             public BatchCursor<T> apply(final BsonDocument result) {
                 QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),
                                                                          connection.getDescription().getServerAddress());
-                return new QueryBatchCursor<T>(queryResult, 0, batchSize, decoder, source);
+                return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source);
             }
         };
     }
@@ -740,7 +740,7 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async
             public AsyncBatchCursor<T> apply(final BsonDocument result) {
                 QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),
                                                                          connection.getDescription().getServerAddress());
-                return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize, decoder, source, connection);
+                return new AsyncQueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source, connection);
             }
         };
     }
diff --git a/driver-core/src/main/com/mongodb/operation/QueryBatchCursor.java b/driver-core/src/main/com/mongodb/operation/QueryBatchCursor.java
@@ -77,7 +77,6 @@ class QueryBatchCursor<T> implements BatchCursor<T> {
 
         initFromQueryResult(firstQueryResult);
         if (limitReached()) {
-            notNull("connection", connection);
             killCursor(connection);
         }
     }
@@ -228,8 +227,9 @@ private BsonDocument asGetMoreCommandDocument() {
         BsonDocument document = new BsonDocument("getMore", new BsonInt64(serverCursor.getId()))
                                 .append("collection", new BsonString(namespace.getCollectionName()));
 
-        if (batchSize != 0) {
-            document.append("batchSize", new BsonInt32(Math.abs(batchSize)));
+        int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count));
+        if (batchSizeForGetMoreCommand != 0) {
+            document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));
         }
 
         return document;
@@ -248,7 +248,7 @@ private void initFromCommandResult(final BsonDocument getMoreCommandResultDocume
     }
 
     private boolean limitReached() {
-        return limit != 0 && count >= limit;
+        return Math.abs(limit) != 0 && count >= Math.abs(limit);
     }
 
     private void killCursor() {
@@ -264,6 +264,7 @@ private void killCursor() {
 
     private void killCursor(final Connection connection) {
         if (serverCursor != null) {
+            notNull("connection", connection);
             connection.killCursor(namespace, singletonList(serverCursor.getId()));
             serverCursor = null;
         }
diff --git a/driver-core/src/test/functional/com/mongodb/operation/AsyncQueryBatchCursorSpecification.groovy b/driver-core/src/test/functional/com/mongodb/operation/AsyncQueryBatchCursorSpecification.groovy
@@ -109,6 +109,25 @@ class AsyncQueryBatchCursorSpecification extends OperationFunctionalSpecificatio
         !nextBatch()
     }
 
+    def 'should exhaust multiple batches with limit'() {
+        given:
+        cursor = new AsyncQueryBatchCursor<Document>(executeQuery(limit, 2), limit, 2, new DocumentCodec(), connectionSource, connection)
+
+        when:
+        def next = nextBatch()
+        def total = 0
+        while (next) {
+            total +=next.size()
+            next = nextBatch()
+        }
+
+        then:
+        total == Math.abs(limit)
+
+        where:
+        limit << [5, -5]
+    }
+
     def 'should exhaust multiple batches'() {
         given:
         cursor = new AsyncQueryBatchCursor<Document>(executeQuery(3), 0, 2, new DocumentCodec(), connectionSource, connection)
@@ -311,11 +330,9 @@ class AsyncQueryBatchCursorSpecification extends OperationFunctionalSpecificatio
         if (serverIsAtLeastVersionThreeDotTwo(connection.getDescription())) {
             def findCommand = new BsonDocument('find', new BsonString(getCollectionName()))
                     .append('filter', filter)
+                    .append('batchSize', new BsonInt32(batchSize))
                     .append('tailable', BsonBoolean.valueOf(tailable))
                     .append('awaitData', BsonBoolean.valueOf(awaitData))
-            if (batchSize > 0) {
-                findCommand.append('batchSize', new BsonInt32(batchSize))
-            }
             if (limit > 0) {
                 findCommand.append('limit', new BsonInt32(limit))
             }
diff --git a/driver-core/src/test/functional/com/mongodb/operation/QueryBatchCursorSpecification.groovy b/driver-core/src/test/functional/com/mongodb/operation/QueryBatchCursorSpecification.groovy
@@ -147,17 +147,20 @@ class QueryBatchCursorSpecification extends OperationFunctionalSpecification {
 
     def 'test limit exhaustion'() {
         given:
-        def firstBatch = executeQuery(5, 2)
+        def firstBatch = executeQuery(limit, 2)
         def connection = connectionSource.getConnection()
 
         when:
-        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 2, new DocumentCodec(), connectionSource, connection)
+        cursor = new QueryBatchCursor<Document>(firstBatch, limit, 2, new DocumentCodec(), connectionSource, connection)
 
         then:
-        cursor.iterator().sum { it.size } == 5
+        cursor.iterator().sum { it.size } == Math.abs(limit)
 
         cleanup:
         connection?.release()
+
+        where:
+        limit << [5, -5]
     }
 
     def 'test remove'() {
@@ -488,12 +491,10 @@ class QueryBatchCursorSpecification extends OperationFunctionalSpecification {
             if (serverIsAtLeastVersionThreeDotTwo(connection.getDescription())) {
                 def findCommand = new BsonDocument('find', new BsonString(getCollectionName()))
                         .append('filter', filter)
+                        .append('limit', new BsonInt32(Math.abs(limit)))
                         .append('batchSize', new BsonInt32(batchSize))
                         .append('tailable', BsonBoolean.valueOf(tailable))
                         .append('awaitData', BsonBoolean.valueOf(awaitData))
-                if (limit > 0) {
-                    findCommand.append('limit', new BsonInt32(limit))
-                }
                 def response = connection.command(getDatabaseName(), findCommand,
                                                   slaveOk, new NoOpFieldNameValidator(),
                                                   CommandResultDocumentCodec.create(new DocumentCodec(), 'firstBatch'))
diff --git a/driver-core/src/test/resources/command-monitoring/find.json b/driver-core/src/test/resources/command-monitoring/find.json
@@ -469,7 +469,7 @@
                 "$numberLong": "42"
               },
               "collection": "test",
-              "batchSize": 3
+              "batchSize": 1
             },
             "command_name": "getMore",
             "database_name": "command-monitoring-tests"

Original file line number	Diff line number	Diff line change
`@@ -130,7 +130,7 @@ ServerCursor getServerCursor() {`
`130`	`130`	`}`
`131`	`131`
`132`	`132`	`private boolean limitReached() {`
`133`		`- return limit != 0 && count >= limit;`
	`133`	`+ return Math.abs(limit) != 0 && count >= Math.abs(limit);`
`134`	`134`	`}`
`135`	`135`
`136`	`136`	`private void getMore(final SingleResultCallback<List<T>> callback) {`
`@@ -162,8 +162,9 @@ private BsonDocument asGetMoreCommandDocument() {`
`162`	`162`	`BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursor.getId()))`
`163`	`163`	`.append("collection", new BsonString(namespace.getCollectionName()));`
`164`	`164`
`165`		`- if (batchSize != 0) {`
`166`		`- document.append("batchSize", new BsonInt32(Math.abs(batchSize)));`
	`165`	`+ int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count));`
	`166`	`+ if (batchSizeForGetMoreCommand != 0) {`
	`167`	`+ document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));`
`167`	`168`	`}`
`168`	`169`
`169`	`170`	`return document;`
Original file line number	Diff line number	Diff line change
`@@ -728,7 +728,7 @@ private Function<BsonDocument, BatchCursor<T>> transformer(final ConnectionSourc`
`728`	`728`	`public BatchCursor<T> apply(final BsonDocument result) {`
`729`	`729`	`QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),`
`730`	`730`	`connection.getDescription().getServerAddress());`
`731`		`- return new QueryBatchCursor<T>(queryResult, 0, batchSize, decoder, source);`
	`731`	`+ return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source);`
`732`	`732`	`}`
`733`	`733`	`};`
`734`	`734`	`}`
`@@ -740,7 +740,7 @@ private Function<BsonDocument, AsyncBatchCursor<T>> asyncTransformer(final Async`
`740`	`740`	`public AsyncBatchCursor<T> apply(final BsonDocument result) {`
`741`	`741`	`QueryResult<T> queryResult = cursorDocumentToQueryResult(result.getDocument("cursor"),`
`742`	`742`	`connection.getDescription().getServerAddress());`
`743`		`- return new AsyncQueryBatchCursor<T>(queryResult, 0, batchSize, decoder, source, connection);`
	`743`	`+ return new AsyncQueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source, connection);`
`744`	`744`	`}`
`745`	`745`	`};`
`746`	`746`	`}`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,6 @@ class QueryBatchCursor<T> implements BatchCursor<T> {`
`77`	`77`
`78`	`78`	`initFromQueryResult(firstQueryResult);`
`79`	`79`	`if (limitReached()) {`
`80`		`- notNull("connection", connection);`
`81`	`80`	`killCursor(connection);`
`82`	`81`	`}`
`83`	`82`	`}`
`@@ -228,8 +227,9 @@ private BsonDocument asGetMoreCommandDocument() {`
`228`	`227`	`BsonDocument document = new BsonDocument("getMore", new BsonInt64(serverCursor.getId()))`
`229`	`228`	`.append("collection", new BsonString(namespace.getCollectionName()));`
`230`	`229`
`231`		`- if (batchSize != 0) {`
`232`		`- document.append("batchSize", new BsonInt32(Math.abs(batchSize)));`
	`230`	`+ int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count));`
	`231`	`+ if (batchSizeForGetMoreCommand != 0) {`
	`232`	`+ document.append("batchSize", new BsonInt32(batchSizeForGetMoreCommand));`
`233`	`233`	`}`
`234`	`234`
`235`	`235`	`return document;`
`@@ -248,7 +248,7 @@ private void initFromCommandResult(final BsonDocument getMoreCommandResultDocume`
`248`	`248`	`}`
`249`	`249`
`250`	`250`	`private boolean limitReached() {`
`251`		`- return limit != 0 && count >= limit;`
	`251`	`+ return Math.abs(limit) != 0 && count >= Math.abs(limit);`
`252`	`252`	`}`
`253`	`253`
`254`	`254`	`private void killCursor() {`
`@@ -264,6 +264,7 @@ private void killCursor() {`
`264`	`264`
`265`	`265`	`private void killCursor(final Connection connection) {`
`266`	`266`	`if (serverCursor != null) {`
	`267`	`+ notNull("connection", connection);`
`267`	`268`	`connection.killCursor(namespace, singletonList(serverCursor.getId()));`
`268`	`269`	`serverCursor = null;`
`269`	`270`	`}`