Make implementations of AsyncAggregateResponseBatchCursor more responsive to close signals (#654)

1) Change implementations of AsyncAggregateResponseBatchCursor to make them more responsive to close signals

AsyncAggregateResponseBatchCursor wraps AsyncQueryBatchCursor.
AsyncQueryBatchCursor implements pending cancel functionality.
AsyncQueryBatchCursor may enter a pending close state as a result of calling close
concurrently with the object being in the pending operation state.
AsyncAggregateResponseBatchCursor does not need to duplicate the pending functionality,
in fact, it prevented AsyncQueryBatchCursor from receiving close signals while
being in the pending operation state.

2) Release resources in AsyncQueryBatchCursor

The changes made is this commit appear to be correct from the standpoint of tests,
but it is very hard for me to prove that by inspecting the code. There are two reasons:
- the reference counting approach used in the code
- the presence of multiple places in code where the instance field connectionSource
is released, combined with the fact that it is not always retained in the constructor
and must be released not more than once per the lifetime of an AsyncQueryBatchCursor object.

The following text expresses my thoughts about reference counting and is not directly
related to the changes made in this commit.
The current approach for reference counting in the driver code is what we may call
a release-early approach, similar to the one described in the Netty documentation
(https://github.com/netty/netty/wiki/Reference-counted-objects#who-destroys-it).
In this approach "the party that accesses a reference-counted object last
is also responsible for the destruction of that reference-counted object".
The problem with this approach is that it is not structured.

"Structured" reference counting.
The idea is to try and get our reference counting code closer to
following the structured programming approach
(https://en.wikipedia.org/wiki/Structured_programming)
while fixing a bug: AsyncQueryBatchCursor does not always release a connection
when it is not needed anymore.

In simple words, the following
is desirable: given a place in code where retain/release is called,
it should be simple to tell where the corresponding release/retain is called.

As a side note, structured programming applied to concurrent code
is called structured concurrency. This idea is realized in Kotlin:
https://kotlinlang.org/docs/reference/coroutines/basics.html#structured-concurrency,
https://elizarov.medium.com/structured-concurrency-722d765aa952
(the article is written by the author of Kotlin coroutines functionality).

As much as I would like us to use the structured approach, it is generally speaking
not possible to reconcile it with the release-early in the same codebase.
In other words, a complete refactoring may be needed to switch to the structured approach.
Following is an example demonstrating a situation when synchronous code that was written
with the release-early approach behaves differently when combined with the structured approach.

Consider the following legacy code in which we want to refactor releaseEarlyLegacyMethod
(the only code that will stay unchanged is the releaseEarlyUnchangedCode method):

    void releaseEarlyLegacyMethod() {
        ReferenceCounted resource = pool.borrow();//count is 1
        /* Since ReleaseEarlyLegacyConsumer is the last party using the resource, it must release the resource.
         * A reader cannot find the release call corresponding to the borrow call inside the releaseEarlyLegacyMethod method.*/
        releaseEarlyUnchangedCode(new ReleaseEarlyLegacyOneTimeConsumer(resource));
    }

    class ReleaseEarlyLegacyOneTimeConsumer implements Runnable {
        private final ReferenceCounted resource;

        ReleaseEarlyLegacyOneTimeConsumer(ReferenceCounted resource) {
            this.resource = resource;//count is 1
        }

        public void run() {
            try {
                System.out.println(resource.toString());//count is 1; use the resource
            } finally {//release the resource because ReleaseEarlyLegacyConsumer.accept is the last party accessing it
                /* Note that by looking at this release call, a reader cannot easily find the corresponding retain/borrow call.
                 * The corresponding retain/borrow call is not even in this class.*/
                resource.release();//count is 0; the resource returns to the pool
            }
        }
    }

    void releaseEarlyUnchangedCode(Runnable action) {
        action.run();
        ReferenceCounted resource = pool.borrow();//count is 1; succeeds only if the single resource has been returned to the pool
        try {//use the resource
            System.out.println(resource.toString());//count is 1
        } finally {//release the resource because releaseEarlyUnchangedCode is the last party accessing it
            resource.release();//count is 0; the resource is returned to the pool
        }
    }

The pool.borrow call in releaseEarlyUnchangedCode works because the resource is returned to the pool before action.run() returns.
Here is the refactored code that uses structured reference counting:

    void structuredNewMethod() {
        ReferenceCounted resource = pool.borrow();//count is 1
        //a reader can easily find the release call corresponding to the borrow call because it is in the finally block in this method
        try {
            releaseEarlyUnchangedCode(new StructuredNewOneTimeConsumer(resource));
        } finally {
            resource.release();//count is 0; the resource returns to the pool
        }
    }

    class StructuredNewOneTimeConsumer implements Runnable {
        private final ReferenceCounted resource;

        StructuredNewOneTimeConsumer(ReferenceCounted resource) {
            this.resource = resource.retain();//count is 2; we store a new reference to the resource
        }

        public void run() {
            try {
                System.out.println(resource.toString());//count is 2; use the resource
            } finally {//release the resource because StructuredNewConsumer is not going to use it anymore
                /* The corresponding retain call is in the constructor of the class because resource is the instance field.
                 * A reader may easily locate it.*/
                resource.release();//count is 1
            }
        }
    }

With this new code, the pool.borrow call in releaseEarlyUnchangedCode fails because the resource is not returned to the pool before action.run() returns.
It is only returned to the pool before structuredNewMethod returns.
Thereby, we have shown that the release-early and the structured approaches are not reconcilable.

This would not be a problem if the code was written following the structured approach from the beginning
because in that case, releaseEarlyUnchangedCode would have been written differently
(using either return-early or structured reference counting definitely affects the structure of the rest of the code).

JAVA-3907

Author: stIncMale

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

@@ -30,6 +30,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static com.mongodb.internal.async.ErrorHandlingResultCallback.errorHandlingCallback;
 import static com.mongodb.internal.operation.ChangeStreamBatchCursorHelper.isRetryableError;
@@ -44,12 +45,7 @@ final class AsyncChangeStreamBatchCursor<T> implements AsyncAggregateResponseBat
 
     private volatile BsonDocument resumeToken;
     private volatile AsyncAggregateResponseBatchCursor<RawBsonDocument> wrapped;
-
-    /* protected by `this` */
-    private boolean isClosed = false;
-    private boolean isOperationInProgress = false;
-    private boolean isClosePending = false;
-    /* protected by `this` */
+    private final AtomicBoolean isClosed;
 
     AsyncChangeStreamBatchCursor(final ChangeStreamOperation<T> changeStreamOperation,
                                  final AsyncAggregateResponseBatchCursor<RawBsonDocument> wrapped,
@@ -62,6 +58,7 @@ final class AsyncChangeStreamBatchCursor<T> implements AsyncAggregateResponseBat
         binding.retain();
         this.resumeToken = resumeToken;
         this.maxWireVersion = maxWireVersion;
+        isClosed = new AtomicBoolean();
     }
 
     AsyncAggregateResponseBatchCursor<RawBsonDocument> getWrapped() {
@@ -94,22 +91,13 @@ public void apply(final AsyncAggregateResponseBatchCursor<RawBsonDocument> curso
 
     @Override
     public void close() {
-        boolean closeCursor = false;
-
-        synchronized (this) {
-            if (isOperationInProgress) {
-                isClosePending = true;
-            } else {
-                closeCursor = !isClosed;
-                isClosed = true;
-                isClosePending = false;
+        if (isClosed.compareAndSet(false, true)) {
+            try {
+                wrapped.close();
+            } finally {
+                binding.release();
             }
         }
-
-        if (closeCursor) {
-            wrapped.close();
-            binding.release();
-        }
     }
 
     @Override
@@ -124,9 +112,7 @@ public int getBatchSize() {
 
     @Override
     public boolean isClosed() {
-        synchronized (this) {
-            return isClosed;
-        }
+        return isClosed.get();
     }
 
     @Override
@@ -155,17 +141,6 @@ private void cachePostBatchResumeToken(final AsyncAggregateResponseBatchCursor<R
         }
     }
 
-    private void endOperationInProgress() {
-        boolean closePending = false;
-        synchronized (this) {
-            isOperationInProgress = false;
-            closePending = this.isClosePending;
-        }
-        if (closePending) {
-            close();
-        }
-    }
-
     private SingleResultCallback<List<RawBsonDocument>> convertResultsCallback(final SingleResultCallback<List<T>> callback) {
         return errorHandlingCallback(new SingleResultCallback<List<RawBsonDocument>>() {
             @Override
@@ -203,25 +178,20 @@ private interface AsyncBlock {
 
     private void resumeableOperation(final AsyncBlock asyncBlock, final SingleResultCallback<List<RawBsonDocument>> callback,
                                      final boolean tryNext) {
-        synchronized (this) {
-            if (isClosed) {
-                callback.onResult(null, new MongoException(format("%s called after the cursor was closed.",
-                        tryNext ? "tryNext()" : "next()")));
-                return;
-            }
-            isOperationInProgress = true;
+        if (isClosed.get()) {
+            callback.onResult(null, new MongoException(format("%s called after the cursor was closed.",
+                    tryNext ? "tryNext()" : "next()")));
+            return;
         }
         asyncBlock.apply(wrapped, new SingleResultCallback<List<RawBsonDocument>>() {
             @Override
             public void onResult(final List<RawBsonDocument> result, final Throwable t) {
                 if (t == null) {
-                    endOperationInProgress();
                     callback.onResult(result, null);
                 } else if (isRetryableError(t, maxWireVersion)) {
                     wrapped.close();
                     retryOperation(asyncBlock, callback, tryNext);
                 } else {
-                    endOperationInProgress();
                     callback.onResult(null, t);
                 }
             }

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

@@ -69,15 +69,15 @@ class AsyncQueryBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T> {
     private volatile int batchSize;
     private final AtomicInteger count = new AtomicInteger();
     private volatile BsonDocument postBatchResumeToken;
-    private volatile BsonTimestamp operationTime;
-    private volatile boolean firstBatchEmpty;
-    private volatile int maxWireVersion = 0;
+    private final BsonTimestamp operationTime;
+    private final boolean firstBatchEmpty;
+    private final int maxWireVersion;
 
     /* protected by `this` */
     private boolean isOperationInProgress = false;
     private boolean isClosed = false;
-    private boolean isClosePending = false;
     /* protected by `this` */
+    private volatile boolean isClosePending = false;
 
     AsyncQueryBatchCursor(final QueryResult<T> firstBatch, final int limit, final int batchSize, final long maxTimeMS,
                           final Decoder<T> decoder, final AsyncConnectionSource connectionSource, final AsyncConnection connection) {
@@ -100,6 +100,8 @@ class AsyncQueryBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T> {
         if (result != null) {
             this.operationTime = result.getTimestamp(OPERATION_TIME, null);
             this.postBatchResumeToken = getPostBatchResumeTokenFromResponse(result);
+        } else {
+            this.operationTime = null;
         }
 
         firstBatchEmpty = firstBatch.getResults().isEmpty();
@@ -109,26 +111,24 @@ class AsyncQueryBatchCursor<T> implements AsyncAggregateResponseBatchCursor<T> {
                 killCursor(connection);
             }
         }
-        if (connection != null) {
-            this.maxWireVersion = connection.getDescription().getMaxWireVersion();
-        }
+        this.maxWireVersion = connection == null ? 0 : connection.getDescription().getMaxWireVersion();
     }
 
     @Override
     public void close() {
-        boolean killCursor = false;
+        boolean doClose = false;
 
         synchronized (this) {
             if (isOperationInProgress) {
                 isClosePending = true;
-            } else {
-                killCursor = !isClosed;
+            } else if (!isClosed) {
                 isClosed = true;
                 isClosePending = false;
+                doClose = true;
             }
         }
 
-        if (killCursor) {
+        if (doClose) {
             killCursorOnClose();
         }
     }
@@ -336,7 +336,17 @@ private void endOperationInProgress() {
     private void handleGetMoreQueryResult(final AsyncConnection connection, final SingleResultCallback<List<T>> callback,
                                           final QueryResult<T> result, final boolean tryNext) {
         cursor.set(result.getCursor());
-        if (!tryNext && result.getResults().isEmpty() && result.getCursor() != null) {
+        if (isClosePending) {
+            try {
+                connection.release();
+                if (result.getCursor() == null) {
+                    connectionSource.release();
+                }
+                endOperationInProgress();
+            } finally {
+                callback.onResult(null, null);
+            }
+        } else if (!tryNext && result.getResults().isEmpty() && result.getCursor() != null) {
             getMore(connection, result.getCursor(), callback, false);
         } else {
             count.addAndGet(result.getResults().size());

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

@@ -35,6 +35,7 @@ import com.mongodb.internal.connection.QueryResult
 import org.bson.BsonBoolean
 import org.bson.BsonDocument
 import org.bson.BsonInt32
+import org.bson.BsonNull
 import org.bson.BsonString
 import org.bson.BsonTimestamp
 import org.bson.Document
@@ -303,6 +304,29 @@ class AsyncQueryBatchCursorFunctionalSpecification extends OperationFunctionalSp
         false     | 0
     }
 
+    @Slow
+    def 'should unblock if closed while waiting for more data from tailable cursor'() {
+        given:
+        collectionHelper.create(collectionName, new CreateCollectionOptions().capped(true).sizeInBytes(1000))
+        collectionHelper.insertDocuments(new DocumentCodec(), Document.parse('{}'))
+        def firstBatch = executeQuery(new BsonDocument('_id', BsonNull.VALUE), 0, 1, true, true);
+
+        when:
+        cursor = new AsyncQueryBatchCursor<Document>(firstBatch, 0, 1, 500, new DocumentCodec(), connectionSource, connection)
+        Thread.start {
+            Thread.sleep(SECONDS.toMillis(2))
+            cursor.close()
+        }
+        def batch = nextBatch()
+
+        then:
+        cursor.isClosed()
+        batch == null
+        //both connection and connectionSource have reference count 1 when we pass them to the AsyncQueryBatchCursor constructor
+        connection.getCount() == 1
+        waitForRelease(connectionSource, 1)
+    }
+
     def 'should respect limit'() {
         given:
         cursor = new AsyncQueryBatchCursor<Document>(executeQuery(6, 3), 6, 2, 0, new DocumentCodec(), connectionSource, connection)

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

@@ -474,7 +474,9 @@ class AsyncQueryBatchCursorSpecification extends Specification {
 
         then:
         connectionA.getCount() == 0
-        connectionSource.getCount() == 0
+        if (response2 == null) { //otherwise connectionSource is released asynchronously, which is not easy to verify
+            connectionSource.getCount() == 0
+        }
         cursor.isClosed()
 
         where:
@@ -709,6 +711,9 @@ class AsyncQueryBatchCursorSpecification extends Specification {
         int counter = 0
         def mock = Mock(AsyncConnectionSource)
         mock.getConnection(_) >> {
+            if (counter == 0) {
+                throw new IllegalStateException('Tried to use released AsyncConnectionSource')
+            }
             def (result, error) = connectionCallbackResults()
             it[0].onResult(result, error)
         }