increase batchSize when limit is equal to it

Author: NathanQingyangXu

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

@@ -188,11 +188,18 @@ <TResult> FindOperation<TResult> find(final MongoNamespace findNamespace, @Nulla
 
     private <TResult> FindOperation<TResult> createFindOperation(final MongoNamespace findNamespace, @Nullable final Bson filter,
                                                                  final Class<TResult> resultClass, final FindOptions options) {
+        final int effectiveBatchSize;
+        if (options.getBatchSize() > 0 && options.getBatchSize() == options.getLimit()) {
+            // Eliminate unnecessary killCursors command when batchSize == limit
+            effectiveBatchSize = options.getBatchSize() + 1;
+        } else {
+            effectiveBatchSize = options.getBatchSize();
+        }
         FindOperation<TResult> operation = new FindOperation<>(
                 findNamespace, codecRegistry.get(resultClass))
                 .retryReads(retryReads)
                 .filter(filter == null ? new BsonDocument() : filter.toBsonDocument(documentClass, codecRegistry))
-                .batchSize(options.getBatchSize())
+                .batchSize(effectiveBatchSize)
                 .skip(options.getSkip())
                 .limit(options.getLimit())
                 .projection(toBsonDocument(options.getProjection()))

driver-core/src/test/resources/unified-test-format/crud/find.json

@@ -237,6 +237,68 @@
           ]
         }
       ]
+    },
+    {
+      "description": "Find with batchSize equal to limit",
+      "operations": [
+        {
+          "object": "collection0",
+          "name": "find",
+          "arguments": {
+            "filter": {
+              "_id": {
+                "$gt": 1
+              }
+            },
+            "sort": {
+              "_id": 1
+            },
+            "limit": 4,
+            "batchSize": 4
+          },
+          "expectResult": [
+            {
+              "_id": 2,
+              "x": 22
+            },
+            {
+              "_id": 3,
+              "x": 33
+            },
+            {
+              "_id": 4,
+              "x": 44
+            },
+            {
+              "_id": 5,
+              "x": 55
+            }
+          ]
+        }
+      ],
+      "expectEvents": [
+        {
+          "client": "client0",
+          "events": [
+            {
+              "commandStartedEvent": {
+                "command": {
+                  "find": "coll0",
+                  "filter": {
+                    "_id": {
+                      "$gt": 1
+                    }
+                  },
+                  "limit": 4,
+                  "batchSize": 5
+                },
+                "commandName": "find",
+                "databaseName": "find-tests"
+              }
+            }
+          ]
+        }
+      ]
     }
   ]
 }

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/FindPublisherImpl.java

@@ -233,6 +233,13 @@ Function<AsyncOperations<?>, TimeoutSettings> getTimeoutSettings() {
 
     @Override
     AsyncReadOperation<AsyncBatchCursor<T>> asAsyncFirstReadOperation() {
-        return getOperations().findFirst(filter, getDocumentClass(), findOptions);
+        final FindOptions effectiveFindOptions;
+        if (findOptions.getBatchSize() > 0 && findOptions.getBatchSize() == findOptions.getLimit()) {
+            // Eliminate unnecessary killCursors command when batchSize == limit
+            effectiveFindOptions = findOptions.withBatchSize(findOptions.getBatchSize() + 1);
+        } else {
+            effectiveFindOptions = findOptions;
+        }
+        return getOperations().findFirst(filter, getDocumentClass(), effectiveFindOptions);
     }
 }

driver-reactive-streams/src/test/unit/com/mongodb/reactivestreams/client/internal/FindPublisherImplTest.java

@@ -88,7 +88,7 @@ void shouldBuildTheExpectedOperation() {
                 .retryReads(true)
                 .filter(new BsonDocument())
                 .allowDiskUse(false)
-                .batchSize(100)
+                .batchSize(101)
                 .collation(COLLATION)
                 .comment(new BsonString("my comment"))
                 .cursorType(CursorType.NonTailable)

driver-sync/src/main/com/mongodb/client/internal/FindIterableImpl.java

@@ -245,7 +245,14 @@ private <E> E executeExplain(final Class<E> explainResultClass, @Nullable final
     }
 
     public ExplainableReadOperation<BatchCursor<TResult>> asReadOperation() {
-        return operations.find(filter, resultClass, findOptions);
+        final FindOptions effectiveFindOptions;
+        if (findOptions.getBatchSize() > 0 && findOptions.getBatchSize() == findOptions.getLimit()) {
+            // Eliminate unnecessary killCursors command when batchSize == limit
+            effectiveFindOptions = findOptions.withBatchSize(findOptions.getBatchSize() + 1);
+        } else {
+            effectiveFindOptions = findOptions;
+        }
+        return operations.find(filter, resultClass, effectiveFindOptions);
     }
 
 }

driver-sync/src/test/unit/com/mongodb/client/gridfs/GridFSFindIterableSpecification.groovy

@@ -88,7 +88,7 @@ class GridFSFindIterableSpecification extends Specification {
         expect operation, isTheSameAs(new FindOperation<GridFSFile>(namespace, gridFSFileCodec)
                 .filter(new BsonDocument('filter', new BsonInt32(2)))
                 .sort(new BsonDocument('sort', new BsonInt32(2)))
-                .batchSize(99)
+                .batchSize(100)
                 .limit(99)
                 .skip(9)
                 .noCursorTimeout(true)

driver-sync/src/test/unit/com/mongodb/client/internal/FindIterableSpecification.groovy

@@ -87,7 +87,7 @@ class FindIterableSpecification extends Specification {
                 .filter(new BsonDocument('filter', new BsonInt32(1)))
                 .sort(new BsonDocument('sort', new BsonInt32(1)))
                 .projection(new BsonDocument('projection', new BsonInt32(1)))
-                .batchSize(100)
+                .batchSize(101)
                 .limit(100)
                 .skip(10)
                 .cursorType(CursorType.NonTailable)
@@ -132,7 +132,7 @@ class FindIterableSpecification extends Specification {
                         .filter(new BsonDocument('filter', new BsonInt32(2)))
                         .sort(new BsonDocument('sort', new BsonInt32(2)))
                         .projection(new BsonDocument('projection', new BsonInt32(2)))
-                        .batchSize(99)
+                        .batchSize(100)
                         .limit(99)
                         .skip(9)
                         .cursorType(CursorType.Tailable)