JAVA-2778: Release connection source for cursor as soon as the last batch has been retrieved

Author: jyemin

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

@@ -49,11 +49,12 @@
 class QueryBatchCursor<T> implements BatchCursor<T> {
     private static final FieldNameValidator NO_OP_FIELD_NAME_VALIDATOR = new NoOpFieldNameValidator();
     private final MongoNamespace namespace;
+    private final ServerAddress serverAddress;
     private final int limit;
     private final Decoder<T> decoder;
-    private final ConnectionSource connectionSource;
     private final long maxTimeMS;
     private int batchSize;
+    private ConnectionSource connectionSource;
     private ServerCursor serverCursor;
     private List<T> nextBatch;
     private int count;
@@ -73,6 +74,7 @@ class QueryBatchCursor<T> implements BatchCursor<T> {
         isTrueArgument("maxTimeMS >= 0", maxTimeMS >= 0);
         this.maxTimeMS = maxTimeMS;
         this.namespace = firstQueryResult.getNamespace();
+        this.serverAddress = firstQueryResult.getAddress();
         this.limit = limit;
         this.batchSize = batchSize;
         this.decoder = notNull("decoder", decoder);
@@ -89,6 +91,10 @@ class QueryBatchCursor<T> implements BatchCursor<T> {
         if (limitReached()) {
             killCursor(connection);
         }
+        if (serverCursor == null && this.connectionSource != null) {
+            this.connectionSource.release();
+            this.connectionSource = null;
+        }
     }
 
     @Override
@@ -205,7 +211,7 @@ public ServerAddress getServerAddress() {
             throw new IllegalStateException("Iterator has been closed");
         }
 
-        return connectionSource.getServerDescription().getAddress();
+        return serverAddress;
     }
 
     private void getMore() {
@@ -229,6 +235,10 @@ private void getMore() {
             if (limitReached()) {
                 killCursor(connection);
             }
+            if (serverCursor == null) {
+                this.connectionSource.release();
+                this.connectionSource = null;
+            }
         } finally {
             connection.release();
         }

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

@@ -40,6 +40,7 @@ import spock.lang.IgnoreIf
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 
+import static com.mongodb.ClusterFixture.checkReferenceCountReachesTarget
 import static com.mongodb.ClusterFixture.getBinding
 import static com.mongodb.ClusterFixture.isDiscoverableReplicaSet
 import static com.mongodb.ClusterFixture.isSharded
@@ -380,6 +381,35 @@ class QueryBatchCursorFunctionalSpecification extends OperationFunctionalSpecifi
         thrown(MongoCursorNotFoundException)
     }
 
+    def 'should release connection source if limit is reached on initial query'() throws InterruptedException {
+        given:
+        def firstBatch = executeQuery(5)
+        def connection = connectionSource.getConnection()
+
+        when:
+        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 0, 0, new DocumentCodec(), connectionSource, connection)
+
+        then:
+        checkReferenceCountReachesTarget(connectionSource, 1)
+
+        cleanup:
+        connection?.release()
+    }
+
+    def 'should release connection source if limit is reached on get more'() throws InterruptedException {
+        given:
+        def firstBatch = executeQuery(3)
+
+        cursor = new QueryBatchCursor<Document>(firstBatch, 5, 3, new DocumentCodec(), connectionSource)
+
+        when:
+        cursor.next()
+        cursor.next()
+
+        then:
+        checkReferenceCountReachesTarget(connectionSource, 1)
+    }
+
     def 'test limit with get more'() {
         given:
         def firstBatch = executeQuery(2)

driver/src/main/com/mongodb/DBCursor.java

@@ -676,7 +676,11 @@ public int count() {
      */
     public DBObject one() {
         DBCursor findOneCursor = copy().limit(-1);
-        return findOneCursor.hasNext() ? findOneCursor.next() : null;
+        try {
+            return findOneCursor.hasNext() ? findOneCursor.next() : null;
+        } finally {
+            findOneCursor.close();
+        }
     }
 
     /**

driver/src/main/com/mongodb/FindIterableImpl.java

@@ -191,7 +191,11 @@ public FindIterable<TResult> snapshot(final boolean snapshot) {
     public TResult first() {
         FindOperation<TResult> findFirstOperation = createQueryOperation().batchSize(0).limit(-1);
         BatchCursor<TResult> batchCursor = getExecutor().execute(findFirstOperation, getReadPreference(), getClientSession());
-        return batchCursor.hasNext() ? batchCursor.next().iterator().next() : null;
+        try {
+            return batchCursor.hasNext() ? batchCursor.next().iterator().next() : null;
+        } finally {
+            batchCursor.close();
+        }
     }
 
     protected ReadOperation<BatchCursor<TResult>> asReadOperation() {