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

Author: jyemin

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

@@ -62,7 +62,6 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
     private final AsyncConnectionSource connectionSource;
     private final AtomicBoolean isClosed = new AtomicBoolean();
     private final AtomicReference<ServerCursor> cursor;
-    private final Object lock = new Object();
     private volatile QueryResult<T> firstBatch;
     private volatile int batchSize;
     private volatile int count;
@@ -80,10 +79,11 @@ class AsyncQueryBatchCursor<T> implements AsyncBatchCursor<T> {
         this.connectionSource = notNull("connectionSource", connectionSource);
         this.count += firstBatch.getResults().size();
 
-        connectionSource.retain();
-        if (firstBatch.getCursor() != null && limitReached()) {
+        if (firstBatch.getCursor() != null) {
             connectionSource.retain();
-            killCursor(connection);
+            if (limitReached()) {
+                killCursor(connection);
+            }
         }
     }
 
@@ -134,8 +134,9 @@ private void next(final SingleResultCallback<List<T>> callback, final boolean tr
             firstBatch = null;
             callback.onResult(results, null);
         } else {
-            ServerCursor localCursor = getCursorForNext();
+            ServerCursor localCursor = getServerCursor();
             if (localCursor == null) {
+                isClosed.set(true);
                 callback.onResult(null, null);
             } else {
                 getMore(localCursor, callback, tryNext);
@@ -152,7 +153,6 @@ private void getMore(final ServerCursor cursor, final SingleResultCallback<List<
             @Override
             public void onResult(final AsyncConnection connection, final Throwable t) {
                 if (t != null) {
-                    connectionSource.release();
                     callback.onResult(null, t);
                 } else {
                     getMore(connection, cursor, callback, tryNext);
@@ -170,13 +170,13 @@ private void getMore(final AsyncConnection connection, final ServerCursor cursor
 
         } else {
             connection.getMoreAsync(namespace, cursor.getId(), getNumberToReturn(limit, batchSize, count),
-                                    decoder, new QueryResultSingleResultCallback(connection, callback, tryNext));
+                    decoder, new QueryResultSingleResultCallback(connection, callback, tryNext));
         }
     }
 
     private BsonDocument asGetMoreCommandDocument(final long cursorId) {
         BsonDocument document = new BsonDocument("getMore", new BsonInt64(cursorId))
-                                .append("collection", new BsonString(namespace.getCollectionName()));
+                .append("collection", new BsonString(namespace.getCollectionName()));
 
         int batchSizeForGetMoreCommand = Math.abs(getNumberToReturn(limit, this.batchSize, count));
         if (batchSizeForGetMoreCommand != 0) {
@@ -201,8 +201,6 @@ public void onResult(final AsyncConnection connection, final Throwable t) {
                     }
                 }
             });
-        } else {
-            connectionSource.release();
         }
     }
 
@@ -239,15 +237,14 @@ public void onResult(final Void result, final Throwable t) {
 
     private BsonDocument asKillCursorsCommandDocument(final ServerCursor localCursor) {
         return new BsonDocument("killCursors", new BsonString(namespace.getCollectionName()))
-                       .append("cursors", new BsonArray(singletonList(new BsonInt64(localCursor.getId()))));
+                .append("cursors", new BsonArray(singletonList(new BsonInt64(localCursor.getId()))));
     }
 
 
     private void handleGetMoreQueryResult(final AsyncConnection connection, final SingleResultCallback<List<T>> callback,
                                           final QueryResult<T> result, final boolean tryNext) {
         if (isClosed()) {
             connection.release();
-            connectionSource.release();
             callback.onResult(null, new MongoException(format("The cursor was closed before %s completed.",
                     tryNext ? "tryNext()" : "next()")));
             return;
@@ -263,7 +260,9 @@ private void handleGetMoreQueryResult(final AsyncConnection connection, final Si
                 connection.release();
             } else {
                 connection.release();
-                connectionSource.release();
+                if (result.getCursor() == null) {
+                    connectionSource.release();
+                }
             }
 
             if (result.getResults().isEmpty()) {
@@ -292,10 +291,9 @@ private class CommandResultSingleResultCallback implements SingleResultCallback<
         public void onResult(final BsonDocument result, final Throwable t) {
             if (t != null) {
                 Throwable translatedException = t instanceof MongoCommandException
-                                                ? translateCommandException((MongoCommandException) t, cursor)
-                                                : t;
+                        ? translateCommandException((MongoCommandException) t, cursor)
+                        : t;
                 connection.release();
-                connectionSource.release();
                 callback.onResult(null, translatedException);
             } else {
                 QueryResult<T> queryResult = getMoreCursorDocumentToQueryResult(result.getDocument("cursor"),
@@ -321,35 +319,14 @@ private class QueryResultSingleResultCallback implements SingleResultCallback<Qu
         public void onResult(final QueryResult<T> result, final Throwable t) {
             if (t != null) {
                 connection.release();
-                connectionSource.release();
                 callback.onResult(null, t);
             } else {
                 handleGetMoreQueryResult(connection, callback, result, tryNext);
             }
         }
     }
 
-    private ServerCursor getCursorForNext() {
-        ServerCursor localCursor;
-        synchronized (lock) {
-            localCursor = cursor.get();
-            if (localCursor == null) {
-                if (!isClosed.getAndSet(true)) {
-                    connectionSource.release();
-                }
-            } else {
-                connectionSource.retain();
-            }
-        }
-        return localCursor;
-    }
-
     ServerCursor getServerCursor() {
-        ServerCursor localCursor;
-        synchronized (lock) {
-            localCursor = cursor.get();
-        }
-        return localCursor;
+        return cursor.get();
     }
-
 }

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

@@ -102,6 +102,40 @@ class AsyncQueryBatchCursorFunctionalSpecification extends OperationFunctionalSp
         !nextBatch()
     }
 
+    def 'should not retain connection and source after cursor is exhausted on first batch'() {
+        given:
+        cursor = new AsyncQueryBatchCursor<Document>(executeQuery(), 0, 0, 0, new DocumentCodec(), connectionSource, connection)
+
+        when:
+        nextBatch()
+
+        then:
+        connection.count == 1
+        connectionSource.count == 1
+    }
+
+    def 'should not retain connection and source after cursor is exhausted on getMore'() {
+        given:
+        cursor = new AsyncQueryBatchCursor<Document>(executeQuery(1, 0), 1, 1, 0, new DocumentCodec(), connectionSource, connection)
+
+        when:
+        nextBatch()
+
+        then:
+        connection.count == 1
+        connectionSource.count == 1
+    }
+
+    def 'should not retain connection and source after cursor is exhausted after first batch'() {
+        when:
+        cursor = new AsyncQueryBatchCursor<Document>(executeQuery(10, 10), 10, 10, 0, new DocumentCodec(), connectionSource,
+                connection)
+
+        then:
+        connection.count == 1
+        connectionSource.count == 1
+    }
+
     def 'should exhaust single batch with limit'() {
         given:
         cursor = new AsyncQueryBatchCursor<Document>(executeQuery(1, 0), 1, 0, 0, new DocumentCodec(), connectionSource, connection)

driver-core/src/test/unit/com/mongodb/operation/AsyncQueryBatchCursorSpecification.groovy

@@ -134,7 +134,7 @@ class AsyncQueryBatchCursorSpecification extends Specification {
         nextBatch(cursor) == FIRST_BATCH
 
         then:
-        connectionSource.getCount() == 1
+        connectionSource.getCount() == 0
 
         then:
         nextBatch(cursor) == null
@@ -203,7 +203,7 @@ class AsyncQueryBatchCursorSpecification extends Specification {
         then:
         batch == thirdBatch
         connectionB.getCount() == 0
-        connectionSource.getCount() == 1
+        connectionSource.getCount() == 0
 
         when:
         batch = tryNextBatch(cursor)
@@ -272,7 +272,7 @@ class AsyncQueryBatchCursorSpecification extends Specification {
         then:
         batch == thirdBatch
         connectionB.getCount() == 0
-        connectionSource.getCount() == 1
+        connectionSource.getCount() == 0
 
         when:
         batch = nextBatch(cursor)
@@ -364,7 +364,7 @@ class AsyncQueryBatchCursorSpecification extends Specification {
 
         then:
         connection.getCount() == 0
-        connectionSource.getCount() == 1
+        connectionSource.getCount() == 0
 
         when:
         cursor.close()
@@ -412,7 +412,7 @@ class AsyncQueryBatchCursorSpecification extends Specification {
 
         then:
         connection.getCount() == 0
-        connectionSource.getCount() == 1
+        connectionSource.getCount() == 0
 
         when:
         cursor.close()