Fix QueryBatchCursor.ResourceManager.close (#1440)

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

@@ -26,6 +26,7 @@
 import com.mongodb.annotations.ThreadSafe;
 import com.mongodb.connection.ConnectionDescription;
 import com.mongodb.connection.ServerType;
+import com.mongodb.internal.VisibleForTesting;
 import com.mongodb.internal.binding.ConnectionSource;
 import com.mongodb.internal.connection.Connection;
 import com.mongodb.internal.connection.QueryResult;
@@ -58,6 +59,7 @@
 import static com.mongodb.assertions.Assertions.isTrueArgument;
 import static com.mongodb.assertions.Assertions.notNull;
 import static com.mongodb.internal.Locks.withLock;
+import static com.mongodb.internal.VisibleForTesting.AccessModifier.PRIVATE;
 import static com.mongodb.internal.operation.CursorHelper.getNumberToReturn;
 import static com.mongodb.internal.operation.DocumentHelper.putIfNotNull;
 import static com.mongodb.internal.operation.SyncOperationHelper.getMoreCursorDocumentToQueryResult;
@@ -359,7 +361,8 @@ private BsonDocument getPostBatchResumeTokenFromResponse(final BsonDocument resu
      * others are not and rely on the total order mentioned above.
      */
     @ThreadSafe
-    private final class ResourceManager {
+    @VisibleForTesting(otherwise = PRIVATE)
+    final class ResourceManager {
         private final Lock lock;
         private volatile State state;
         @Nullable
@@ -416,7 +419,8 @@ <R> R execute(final String exceptionMessageIfClosed, final Supplier<R> operation
          * If {@linkplain #operable() closed}, then returns false, otherwise completes abruptly.
          * @throws IllegalStateException Iff another operation is in progress.
          */
-        private boolean tryStartOperation() throws IllegalStateException {
+        @VisibleForTesting(otherwise = PRIVATE)
+        boolean tryStartOperation() throws IllegalStateException {
             return withLock(lock, () -> {
                 State localState = state;
                 if (!localState.operable()) {
@@ -435,7 +439,8 @@ private boolean tryStartOperation() throws IllegalStateException {
         /**
          * Thread-safe.
          */
-        private void endOperation() {
+        @VisibleForTesting(otherwise = PRIVATE)
+        void endOperation() {
             boolean doClose = withLock(lock, () -> {
                 State localState = state;
                 if (localState == State.OPERATION_IN_PROGRESS) {
@@ -459,7 +464,7 @@ private void endOperation() {
         void close() {
             boolean doClose = withLock(lock, () -> {
                 State localState = state;
-                if (localState == State.OPERATION_IN_PROGRESS) {
+                if (localState.inProgress()) {
                     state = State.CLOSE_PENDING;
                     return false;
                 } else if (localState != State.CLOSED) {

driver-core/src/test/unit/com/mongodb/internal/operation/QueryBatchCursorResourceManagerTest.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.mongodb.internal.operation;
+
+import com.mongodb.ServerAddress;
+import com.mongodb.ServerCursor;
+import com.mongodb.internal.binding.ConnectionSource;
+import com.mongodb.internal.connection.QueryResult;
+import org.bson.BsonDocument;
+import org.bson.codecs.BsonDocumentCodec;
+import org.junit.jupiter.api.Test;
+
+import static java.util.Collections.emptyList;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+final class QueryBatchCursorResourceManagerTest {
+    @Test
+    void doubleCloseExecutedConcurrentlyWithOperationBeingInProgressShouldNotFail() {
+        ConnectionSource connectionSourceMock = mock(ConnectionSource.class);
+        when(connectionSourceMock.retain()).thenReturn(connectionSourceMock);
+        when(connectionSourceMock.release()).thenReturn(1);
+        ServerAddress serverAddress = new ServerAddress();
+        try (QueryBatchCursor<BsonDocument> cursor = new QueryBatchCursor<>(
+                new QueryResult<>(null, emptyList(), 0, serverAddress),
+                1, 1, new BsonDocumentCodec())) {
+            QueryBatchCursor<?>.ResourceManager cursorResourceManager = cursor.new ResourceManager(
+                    connectionSourceMock, null, new ServerCursor(1, serverAddress));
+            cursorResourceManager.tryStartOperation();
+            try {
+                assertDoesNotThrow(() -> {
+                    cursorResourceManager.close();
+                    cursorResourceManager.close();
+                    cursorResourceManager.setServerCursor(null);
+                });
+            } finally {
+                cursorResourceManager.endOperation();
+            }
+        }
+    }
+}