When the limit is reached after the initial batch of results of a FindOperation, use the same connection that was used for the initial OP_QUERY to execute the OP_KILL_CURSOR. This will ensure that a thread never attempts to checkout more than one connection at a time.

 JAVA-1882

Author: jyemin

driver-core/src/main/com/mongodb/operation/FindOperation.java

@@ -406,7 +406,7 @@ public BatchCursor<T> call(final ConnectionSource source, final Connection conne
                                                               isPartial(),
                                                               isOplogReplay(),
                                                               decoder);
-                return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source);
+                return new QueryBatchCursor<T>(queryResult, limit, batchSize, decoder, source, connection);
             }
         });
     }

driver-core/src/main/com/mongodb/operation/QueryBatchCursor.java

@@ -50,6 +50,11 @@ class QueryBatchCursor<T> implements BatchCursor<T> {
 
     QueryBatchCursor(final QueryResult<T> firstQueryResult, final int limit, final int batchSize,
                      final Decoder<T> decoder, final ConnectionSource connectionSource) {
+        this(firstQueryResult, limit, batchSize, decoder, connectionSource, null);
+    }
+
+    QueryBatchCursor(final QueryResult<T> firstQueryResult, final int limit, final int batchSize,
+                     final Decoder<T> decoder, final ConnectionSource connectionSource, final Connection connection) {
         this.namespace = firstQueryResult.getNamespace();
         this.limit = limit;
         this.batchSize = batchSize;
@@ -65,7 +70,8 @@ class QueryBatchCursor<T> implements BatchCursor<T> {
 
         initFromQueryResult(firstQueryResult);
         if (limitReached()) {
-            killCursor();
+            notNull("connection", connection);
+            killCursor(connection);
         }
     }
 

driver-core/src/test/functional/com/mongodb/operation/QueryBatchCursorSpecification.groovy

@@ -116,12 +116,16 @@ class QueryBatchCursorSpecification extends OperationFunctionalSpecification {
     def 'test limit exhaustion'() {
         given:
         def firstBatch = executeQuery(2)
+        def connection = connectionSource.getConnection()
 
         when:
-        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 2, new DocumentCodec(), connectionSource)
+        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 2, new DocumentCodec(), connectionSource, connection)
 
         then:
         cursor.iterator().sum { it.size } == 5
+
+        cleanup:
+        connection?.release()
     }
 
     def 'test remove'() {
@@ -241,8 +245,9 @@ class QueryBatchCursorSpecification extends OperationFunctionalSpecification {
     def 'should kill cursor if limit is reached on initial query'() throws InterruptedException {
         given:
         def firstBatch = executeQuery(5)
+        def connection = connectionSource.getConnection()
 
-        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 0, new DocumentCodec(), connectionSource)
+        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 0, new DocumentCodec(), connectionSource, connection)
 
         ServerCursor serverCursor = cursor.getServerCursor()
         Thread.sleep(1000) //Note: waiting for some time for killCursor operation to be performed on a server.
@@ -252,6 +257,9 @@ class QueryBatchCursorSpecification extends OperationFunctionalSpecification {
 
         then:
         thrown(MongoCursorNotFoundException)
+
+        cleanup:
+        connection?.release()
     }
 
     @IgnoreIf({ isSharded() && !serverVersionAtLeast([2, 4, 0]) })